Add restifier#use

lib/index.js

@@ -11,14 +11,18 @@ var Model = require('./model');
  */
 function RestAPI() {
   this.models = {};
+  this.router = express.Router();
+  this.router.use(require('res-error')({
+    log: false
+  }));
 }
 
 /**
  * Gateway function that can add model(s) or return initialization
  * middleware depending on the arguments passed to it.
  *
  * @returns {Model|Function} A model or middleware depending on the arguments.
- * * `setup()` - Returns `initialize()`
+ * * `setup()` - Returns `this.router`
  * * `setup(mong)` - Returns `model(mong)`
  * * `setup(mong1, mong2)` - Returns `modelsMiddleware(mong1, mong2)`
  */
@@ -29,20 +33,19 @@ RestAPI.prototype.setup = function() {
   }
 
   if (args.length === 0 || !args[0].schema) {
-    return this.initialize();
+    return this.router;
   }
   return this.model(args[0]);
 };
 
 /**
- * Initialization middleware.
+ * Applies middleware to this RestAPI. Uses express.Router#use.
  *
- * @returns {Function} Express middleware
+ * @returns {RestAPI} This RestAPI.
  */
-RestAPI.prototype.initialize = function() {
-  return require('res-error')({
-    log: false
-  });
+RestAPI.prototype.use = function() {
+  this.router.use.apply(this.router, arguments);
+  return this;
 };
 
 /**
@@ -101,7 +104,7 @@ RestAPI.prototype.modelsMiddleware = function(mongs) {
  */
 RestAPI.prototype.middleware = function() {
   var router = express.Router();
-  router.use(this.initialize());
+  router.use(this.router);
   _.forEach(this.models, function(model) {
     router.use(model.middleware());
     _.forEach(model.submodels, function(submodel) {

test/index.js

@@ -147,4 +147,30 @@ describe('restifier', function() {
       });
     });
   });
+
+  describe('middleware', function(done) {
+    it('should apply middleware', function(done) {
+      restifier(User, Post);
+      restifier.use(function(req, res, next) {
+        res.set('yolo', 'swag');
+        return next();
+      });
+      restifier.use('/asdf', function(req, res, next) {
+        res.set('young', 'mani');
+        return res.status(200).send({
+          message: 'hello mundo'
+        });
+      });
+      app.use(restifier.middleware());
+
+      request(app).get('/asdf').end(function(err, res) {
+        expect(err).to.be.null;
+        expect(res.status).to.equal(200);
+        expect(res.header.yolo).to.equal('swag');
+        expect(res.header.young).to.equal('mani');
+        expect(res.body.message).to.equal('hello mundo');
+        done();
+      });
+    });
+  });
 });