From c311f83c1f1f4636659011bd02ff7bd7a2a1f218 Mon Sep 17 00:00:00 2001
From: Sergii Stotskyi <sergiy.stotskiy@gmail.com>
Date: Thu, 15 Jun 2017 21:14:34 +0300
Subject: [PATCH] feat(sanbox): adds support for replacing implementation of
 object's method in `on` method

---
 lib/spy.js    |  9 ++++++---
 test/spies.js | 17 +++++++++++++++++
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/lib/spy.js b/lib/spy.js
index 4b45992..568af73 100644
--- a/lib/spy.js
+++ b/lib/spy.js
@@ -42,17 +42,20 @@ module.exports = function (chai, _) {
    *      const array = []
    *      const spy = chai.spy.sandbox();
    *      const [push, pop] = spy.on(array, ['push', 'pop']);
+   *      
+   *      spy.on(array, 'push', returns => 1)
    *
    * @param {Object} object
    * @param {String|String[]} method name or methods names to spy on
+   * @param {Function} [fn] mock implementation
    * @returns created spy or created spies
    * @api public
    */
 
-  Sandbox.prototype.on = function (object, methodName) {
+  Sandbox.prototype.on = function (object, methodName, fn) {
     if (Array.isArray(methodName)) {
       return methodName.map(function (name) {
-        return this.on(object, name);
+        return this.on(object, name, fn);
       }, this);
     }
 
@@ -69,7 +72,7 @@ module.exports = function (chai, _) {
       throw new Error('"' + methodName + '" is already a spy')
     }
 
-    var method = chai.spy('object.' + methodName, object[methodName]);
+    var method = chai.spy('object.' + methodName, fn || object[methodName]);
     var trackingId = ++spyAmount
 
     this[STATE_KEY][trackingId] = method;
diff --git a/test/spies.js b/test/spies.js
index c02e51d..415f56e 100644
--- a/test/spies.js
+++ b/test/spies.js
@@ -442,6 +442,23 @@ describe('Chai Spies', function () {
         chai.spy.on(object, 'push');
       }).should.throw(Error)
     });
+
+    it('should allow to overwrite method implementation', function () {
+      chai.spy.on(object, 'push', function() {
+        return 5;
+      });
+
+      object.push().should.equal(5);
+    });
+
+    it('should overwrite all methods with the same implementation', function () {
+      chai.spy.on(object, ['push', 'pop'], function() {
+        return 5;
+      });
+
+      object.push().should.equal(5);
+      object.pop().should.equal(5);
+    })
   });
 
   describe('spy interface', function () {