diff --git a/lib/application.js b/lib/application.js index 551bfab9c..efc9cedce 100644 --- a/lib/application.js +++ b/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; } diff --git a/package.json b/package.json index 94be43923..5b8f80e70 100644 --- a/package.json +++ b/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", diff --git a/test/application/wrap.js b/test/application/wrap.js new file mode 100644 index 000000000..d64449959 --- /dev/null +++ b/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(); + }); + }); +});