reference app.wrappers[] for each mw

lib/application.js

@@ -42,6 +42,7 @@ module.exports = class Application extends Emitter {
 
     this.proxy = false;
     this.middleware = [];
+    this.wrappers = [];
     this.subdomainOffset = 2;
     this.env = process.env.NODE_ENV || 'development';
     this.context = Object.create(context);
@@ -110,8 +111,22 @@ module.exports = class Application extends Emitter {
                 'https://github.com/koajs/koa/tree/v2.x#old-signature-middleware-v1x');
       fn = convert(fn);
     }
+
     debug('use %s', fn._name || fn.name || '-');
     this.middleware.push(fn);
+    if (this.wrappers && this.wrappers.length) this.wrap();
+    return this;
+  }
+
+  wrap(wrappers) {
+    wrappers = wrappers || this.wrappers;
+    if (!wrappers || !this.wrappers.length) return this;
+    this.middleware[this.middleware.length - 1]._wrappers = wrappers;
+    return this;
+  }
+
+  dewrap() {
+    delete this.middleware[this.middleware.length - 1]._wrappers;
     return this;
   }
 

package.json

@@ -32,7 +32,7 @@
     "http-assert": "^1.1.0",
     "http-errors": "^1.2.8",
     "is-generator-function": "^1.0.3",
-    "koa-compose": "^3.0.0",
+    "koa-compose": "fl0w/compose#next",
     "koa-convert": "^1.2.0",
     "koa-is-json": "^1.0.0",
     "mime-types": "^2.0.7",

test/application/wrap.js

@@ -0,0 +1,132 @@
+'use strict';
+
+const request = require('supertest');
+const Koa = require('../..');
+
+describe('app.wrap', () => {
+  it('should wrap mw correctly', done => {
+    const app = new Koa();
+    const wrapped = [];
+
+    app.wrappers = [(ctx, next) => {
+      wrapped.push('pre');
+      return next().then(() => {
+        wrapped.push('post');
+      });
+    }];
+
+    app.use((ctx, next) => {
+      wrapped.push('middleware 1');
+      return next();
+    });
+
+    app.use((ctx, next) => {
+      wrapped.push('middleware 2');
+      return next();
+    });
+
+    request(app.listen())
+      .get('/')
+      .end(() => {
+        wrapped.should.eql([
+          'pre',
+          'middleware 1',
+          'pre',
+          'middleware 2',
+          'post',
+          'post'
+        ]);
+        done();
+      });
+  });
+
+  it('should not wrap when mw is dewrapped', done => {
+    const app = new Koa();
+    const wrapped = [];
+
+    app.wrappers = [(ctx, next) => {
+      wrapped.push('pre');
+      return next().then(() => {
+        wrapped.push('post');
+      });
+    }];
+
+    app.use((ctx, next) => {
+      wrapped.push('middleware 1');
+      return next();
+    }).dewrap();
+
+    app.use((ctx, next) => {
+      wrapped.push('middleware 2');
+      return next();
+    }).dewrap();
+
+    request(app.listen())
+      .get('/')
+      .end(() => {
+        wrapped.should.eql(['middleware 1', 'middleware 2']);
+        done();
+      });
+  });
+
+  it('should not wrap when app.wrappers is false', done => {
+    const app = new Koa();
+    const wrapped = [];
+    app.wrappers = false;
+
+    app.use((ctx, next) => {
+      wrapped.push('middleware 1');
+      return next();
+    });
+
+    request(app.listen())
+      .get('/')
+      .end(() => {
+        wrapped.should.eql(['middleware 1']);
+        done();
+      });
+  });
+
+  it('app.wrap() should wrap with provided wrappers', done => {
+    const app = new Koa();
+    const wrapped = [];
+
+    const profiler = (ctx, next) => {
+      wrapped.push('profiler pre');
+      return next().then(() => {
+        wrapped.push('profiler post');
+      });
+    };
+
+    app.wrappers = [(ctx, next) => {
+      wrapped.push('should not wrap mw1');
+      return next().then(() => {
+        wrapped.push('should not wrap mw1');
+      });
+    }];
+
+    app.use((ctx, next) => {
+      wrapped.push('middleware 1');
+      return next();
+    }).wrap([profiler]);
+
+    app.use((ctx, next) => {
+      wrapped.push('middleware 2');
+      return next();
+    });
+
+    request(app.listen())
+      .get('/')
+      .end(() => {
+        wrapped.should.eql([
+          'profiler pre',
+          'middleware 1',
+          'should not wrap mw1',
+          'middleware 2',
+          'should not wrap mw1',
+          'profiler post'
+        ]);
+        done();
+      });
+  });
+});