Return new assertion instead of this for overwriteMethod

chaijs · Sep 15, 2016 · 1d05b40 · 1d05b40
1 parent 4101fdf
commit 1d05b40
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 1 deletion.
diff --git a/lib/chai/utils/overwriteMethod.js b/lib/chai/utils/overwriteMethod.js
@@ -4,7 +4,9 @@
  * MIT Licensed
  */
 
+var chai = require('../../chai');
 var flag = require('./flag');
+var transferFlags = require('./transferFlags');
 
 /**
  * ### overwriteMethod (ctx, name, fn)
@@ -59,6 +61,12 @@ module.exports = function (ctx, name, method) {
     var result = method(_super).apply(this, arguments);
     flag(this, 'keep_ssfi', false);
 
-    return result === undefined ? this : result;
+    if (result !== undefined) {
+      return result;
+    }
+
+    var newAssertion = new chai.Assertion();
+    transferFlags(this, newAssertion);
+    return newAssertion;
   }
 };
diff --git a/test/utilities.js b/test/utilities.js
@@ -392,6 +392,54 @@ describe('utilities', function () {
         }
       }
     });
+
+    it.only('should return a new assertion with flags copied over', function () {
+      var assertionConstructor = chai.Assertion;
+
+      chai.use(function(_chai, utils) {
+        assertionConstructor = _chai.Assertion;
+        _chai.Assertion.addMethod('foo', function(str) {
+          var obj = utils.flag(this, 'object');
+          new _chai.Assertion(obj).to.be.equal(str);
+        });
+
+        _chai.Assertion.overwriteMethod('foo', function() {
+          return function(str) {
+            utils.flag(this, 'mySpecificFlag', 'value1');
+            utils.flag(this, 'ultraSpecificFlag', 'value2');
+            var obj = utils.flag(this, 'object');
+            new _chai.Assertion(obj.toLowerCase()).to.be.equal(str.toLowerCase());
+          };
+        });
+
+        _chai.Assertion.addMethod('checkFlags', function() {
+          this.assert(
+              utils.flag(this, 'mySpecificFlag') === 'value1' &&
+              utils.flag(this, 'ultraSpecificFlag') === 'value2'
+            , 'expected assertion to have specific flags'
+            , "this doesn't matter"
+          );
+        });
+      });
+
+      var assertion1 = expect('bar');
+      var assertion2 = assertion1.foo('bar');
+
+      // Checking if a new assertion was returned
+      expect(assertion1).to.not.be.equal(assertion2);
+
+      // Check if flags were copied
+      assertion2.checkFlags();
+
+      // Checking if it's really an instance of an Assertion
+      expect(assertion2).to.be.instanceOf(assertionConstructor);
+
+      // Test chaining `.length` after a method to guarantee it is not a function's `length`
+      expect('bar').to.be.a.foo('bar').length.above(2);
+
+      // Ensure that foo returns an Assertion (not a function)
+      expect(expect('bar').foo('bar')).to.be.an.instanceOf(assertionConstructor);
+    });
   });
 
   it('addProperty', function () {