Proper support for after with a passed async function in wrapped mode

boot.js

@@ -362,11 +362,27 @@ function encapsulateTwoParam (func, that) {
 function encapsulateThreeParam (func, that) {
   return _encapsulateThreeParam.bind(that)
   function _encapsulateThreeParam (err, cb) {
+    var res
     if (!func) {
       process.nextTick(cb)
-    } else if (func.length === 0 || func.length === 1) {
-      func(err)
-      process.nextTick(cb)
+    } else if (func.length === 0) {
+      res = func()
+      if (res && res.then) {
+        res.then(function () {
+          process.nextTick(cb, err)
+        }, cb)
+      } else {
+        process.nextTick(cb, err)
+      }
+    } else if (func.length === 1) {
+      res = func(err)
+      if (res && res.then) {
+        res.then(function () {
+          process.nextTick(cb)
+        }, cb)
+      } else {
+        process.nextTick(cb)
+      }
     } else if (func.length === 2) {
       func(err, cb)
     } else {

test/after-pass-through.test.js

@@ -0,0 +1,28 @@
+'use strict'
+
+const t = require('tap')
+const boot = require('..')
+const app = {}
+
+boot(app)
+
+t.plan(5)
+
+app.use(function (f, opts) {
+  return Promise.reject(new Error('kaboom'))
+}).after(function (err, cb) {
+  t.pass('this is just called')
+  cb(err)
+}).after(function () {
+  t.pass('this is just called')
+}).after(function (err, cb) {
+  t.pass('this is just called')
+  cb(err)
+})
+
+app.ready().then(() => {
+  t.fail('this should not be called')
+}).catch(err => {
+  t.ok(err)
+  t.strictEqual(err.message, 'kaboom')
+})

test/async-await.js

@@ -180,3 +180,49 @@ test('after', async (t) => {
   t.ok(thirdLoaded, 'third is loaded')
   t.pass('booted')
 })
+
+test('after wrapped', async (t) => {
+  t.plan(15)
+
+  const app = {}
+  boot(app)
+  let firstLoaded = false
+  let secondLoaded = false
+  let thirdLoaded = false
+
+  app.use(first)
+
+  async function first (s, opts) {
+    t.notOk(firstLoaded, 'first is not loaded')
+    t.notOk(secondLoaded, 'second is not loaded')
+    t.notOk(thirdLoaded, 'third is not loaded')
+    firstLoaded = true
+    s.after(second)
+    s.after(third)
+  }
+
+  async function second (err) {
+    t.error(err)
+    t.ok(firstLoaded, 'first is loaded')
+    t.notOk(secondLoaded, 'second is not loaded')
+    t.notOk(thirdLoaded, 'third is not loaded')
+    await sleep(10)
+    secondLoaded = true
+  }
+
+  async function third () {
+    t.ok(firstLoaded, 'first is loaded')
+    t.ok(secondLoaded, 'second is loaded')
+    t.notOk(thirdLoaded, 'third is not loaded')
+    await sleep(10)
+    thirdLoaded = true
+  }
+
+  const readyContext = await app.ready()
+
+  t.equal(app, readyContext)
+  t.ok(firstLoaded, 'first is loaded')
+  t.ok(secondLoaded, 'second is loaded')
+  t.ok(thirdLoaded, 'third is loaded')
+  t.pass('booted')
+})