Add overwriteChainableMethod utility

component.json

@@ -35,6 +35,7 @@
       , "lib/chai/utils/objDisplay.js"
       , "lib/chai/utils/overwriteMethod.js"
       , "lib/chai/utils/overwriteProperty.js"
+      , "lib/chai/utils/overwriteChainableMethod.js"
       , "lib/chai/utils/test.js"
       , "lib/chai/utils/transferFlags.js"
       , "lib/chai/utils/type.js"

lib/chai/assertion.js

@@ -81,6 +81,10 @@ module.exports = function (_chai, util) {
     util.overwriteMethod(this.prototype, name, fn);
   };
 
+  Assertion.overwriteChainableMethod = function (name, fn, chainingBehavior) {
+    util.overwriteChainableMethod(this.prototype, name, fn, chainingBehavior);
+  };
+
   /*!
    * ### .assert(expression, message, negateMessage, expected, actual)
    *

lib/chai/utils/addChainableMethod.js

@@ -55,15 +55,27 @@ var call  = Function.prototype.call,
  */
 
 module.exports = function (ctx, name, method, chainingBehavior) {
-  if (typeof chainingBehavior !== 'function')
+  if (typeof chainingBehavior !== 'function') {
     chainingBehavior = function () { };
+  }
+
+  var chainableBehavior = {
+      method: method
+    , chainingBehavior: chainingBehavior
+  };
+
+  // save the methods so we can overwrite them later, if we need to.
+  if (!ctx.__methods) {
+    ctx.__methods = {};
+  }
+  ctx.__methods[name] = chainableBehavior;
 
   Object.defineProperty(ctx, name,
     { get: function () {
-        chainingBehavior.call(this);
+        chainableBehavior.chainingBehavior.call(this);
 
         var assert = function () {
-          var result = method.apply(this, arguments);
+          var result = chainableBehavior.method.apply(this, arguments);
           return result === undefined ? this : result;
         };
 

lib/chai/utils/index.js

@@ -106,3 +106,9 @@ exports.overwriteMethod = require('./overwriteMethod');
 
 exports.addChainableMethod = require('./addChainableMethod');
 
+/*!
+ * Overwrite chainable method
+ */
+
+exports.overwriteChainableMethod = require('./overwriteChainableMethod');
+

lib/chai/utils/overwriteChainableMethod.js

@@ -0,0 +1,53 @@
+/*!
+ * Chai - overwriteChainableMethod utility
+ * Copyright(c) 2012-2013 Jake Luer <jake@alogicalparadox.com>
+ * MIT Licensed
+ */
+
+/**
+ * ### overwriteChainableMethod (ctx, name, fn)
+ *
+ * Overwites an already existing chainable method
+ * and provides access to the previous function or
+ * property.  Must return functions to be used for
+ * name.
+ *
+ *     utils.overwriteChainableMethod(chai.Assertion.prototype, 'length',
+ *       function (_super) {
+ *       }
+ *     , function (_super) {
+ *       }
+ *     );
+ *
+ * Can also be accessed directly from `chai.Assertion`.
+ *
+ *     chai.Assertion.overwriteChainableMethod('foo', fn, fn);
+ *
+ * Then can be used as any other assertion.
+ *
+ *     expect(myFoo).to.have.length(3);
+ *     expect(myFoo).to.have.length.above(3);
+ *
+ * @param {Object} ctx object whose method / property is to be overwritten
+ * @param {String} name of method / property to overwrite
+ * @param {Function} method function that returns a function to be used for name
+ * @param {Function} chainingBehavior function that returns a function to be used for property
+ * @name overwriteChainableMethod
+ * @api public
+ */
+
+module.exports = function (ctx, name, method, chainingBehavior) {
+  var chainableBehavior = ctx.__methods[name];
+
+  var _chainingBehavior = chainableBehavior.chainingBehavior;
+  chainableBehavior.chainingBehavior = function () {
+    var result = chainingBehavior(_chainingBehavior).call(this);
+    return result === undefined ? this : result;
+  };
+
+  var _method = chainableBehavior.method;
+  chainableBehavior.method = function () {
+    var result = method(_method).apply(this, arguments);
+    return result === undefined ? this : result;
+  };
+};

test/utilities.js

@@ -263,5 +263,47 @@ describe('utilities', function () {
       expect(obj).x.to.be.ok;
       expect(obj).to.have.property('__x', 'X!');
     })
-  })
+  });
+
+  it('overwriteChainableMethod', function () {
+    chai.use(function (_chai, _) {
+      _chai.Assertion.overwriteChainableMethod('x',
+        function(_super) {
+          return function() {
+            if (_.flag(this, 'marked')) {
+              new chai.Assertion(this._obj).to.be.equal('spot');
+            } else {
+              _super.apply(this, arguments);
+            }
+          };
+        }
+      , function(_super) {
+          return function() {
+            _.flag(this, 'message', 'x marks the spot');
+            _super.apply(this, arguments);
+          };
+        }
+      );
+
+      // Make sure the original behavior of 'x' remains the same
+      expect('foo').x.to.equal("foo");
+      expect("x").x();
+      expect(function () {
+        expect("foo").x();
+      }).to.throw(_chai.AssertionError);
+      var obj = {};
+      expect(obj).x.to.be.ok;
+      expect(obj).to.have.property('__x', 'X!');
+
+      // Test the new behavior of 'x'
+      var assertion = expect('foo').x.to.be.ok;
+      expect(_.flag(assertion, 'message')).to.equal('x marks the spot');
+      expect(function () {
+        var assertion = expect('x');
+        _.flag(assertion, 'marked', true);
+        assertion.x()
+      }).to.throw(_chai.AssertionError);
+    });
+  });
+
 });