From ce5558e196ff81fb7dc5fcde3d4a797bc6ea9ac2 Mon Sep 17 00:00:00 2001
From: Lee Bousfield <ljbousfield@gmail.com>
Date: Thu, 24 Mar 2016 10:55:40 -0600
Subject: [PATCH] Add support for a wrapper array in context

---
 index.js     | 15 ++++++++--
 package.json |  3 +-
 test/test.js | 80 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 94 insertions(+), 4 deletions(-)

diff --git a/index.js b/index.js
index 431c0ff..7958157 100644
--- a/index.js
+++ b/index.js
@@ -1,6 +1,7 @@
 'use strict'
 
 const Promise = require('any-promise')
+const callMiddleware = require('koa-call-middleware')
 
 /**
  * Expose compositor.
@@ -41,9 +42,17 @@ function compose (middleware) {
       if (i === middleware.length) fn = next
       if (!fn) return Promise.resolve()
       try {
-        return Promise.resolve(fn(context, function next () {
-          return dispatch(i + 1)
-        }))
+        if (fn) {
+          return Promise.resolve(callMiddleware(fn, context, function next () {
+            return dispatch(i + 1)
+          }))
+        } else {
+          if (next) {
+            return Promise.resolve(next())
+          } else {
+            return Promise.resolve()
+          }
+        }
       } catch (err) {
         return Promise.reject(err)
       }
diff --git a/package.json b/package.json
index efc8d0d..3fa4694 100644
--- a/package.json
+++ b/package.json
@@ -12,7 +12,8 @@
     "index.js"
   ],
   "dependencies": {
-    "any-promise": "^1.1.0"
+    "any-promise": "^1.1.0",
+    "koa-call-middleware": "^1.0.0"
   },
   "devDependencies": {
     "co": "^4.6.0",
diff --git a/test/test.js b/test/test.js
index 5600909..7f17b48 100644
--- a/test/test.js
+++ b/test/test.js
@@ -85,6 +85,66 @@ describe('Koa Compose', function () {
     })
   })
 
+  it('should wrap functions with wrappers', function () {
+    var arr = []
+    var stack = []
+
+    stack.push(function * (context, next) {
+      arr.push(1)
+      yield wait(1)
+      yield next()
+      yield wait(1)
+      arr.push(6)
+    })
+
+    stack.push(function * (context, next) {
+      arr.push(2)
+      yield wait(1)
+      yield next()
+      yield wait(1)
+      arr.push(5)
+    })
+
+    stack.push(function * (context, next) {
+      arr.push(3)
+      yield wait(1)
+      yield next()
+      yield wait(1)
+      arr.push(4)
+    })
+
+    return compose(stack)({
+      wrappers: [co.wrap]
+    }).then(function () {
+      arr.should.eql([1, 2, 3, 4, 5, 6])
+    })
+  })
+
+  it('should use wrappers in the correct order', function () {
+    let completed = []
+    let stack = []
+    let arr = []
+
+    stack.push(function * (context, next) {
+      arr.push(1)
+      yield next()
+    })
+
+    stack.push(function * (context, next) {
+      arr.push(2)
+      yield next()
+    })
+
+    return compose(stack)({
+      wrappers: [co.wrap, function (fn) {
+        return (ctx, next) => fn(ctx, next).then(completed.push(fn))
+      }]
+    }).then(function () {
+      arr.should.eql([1, 2])
+      completed.length.should.eql(stack.length)
+    })
+  })
+
   it('should only accept an array', function () {
     var err
     try {
@@ -95,6 +155,26 @@ describe('Koa Compose', function () {
     return (err).should.be.instanceof(TypeError)
   })
 
+  it('should throw an error if calling generators or non-functions', function () {
+    var stack = [
+      function * () {}
+    ]
+    return compose(stack)({}).then(function () {
+      throw new Error('No error thrown')
+    }).catch(function (err) {
+      (err).should.be.instanceof(TypeError)
+    }).then(function () {
+      stack = [
+        { key: 'value' }
+      ]
+      return compose(stack)({})
+    }).then(function () {
+      throw new Error('No error thrown')
+    }).catch(function (err) {
+      (err).should.be.instanceof(TypeError)
+    })
+  })
+
   it('should work with 0 middleware', function () {
     return compose([])({})
   })