mocha beforeEach/afterEach can still take callbacks

isaacs · isaacs · commit c363aebf69d2 · 2021-02-16T18:13:08.000-08:00
diff --git a/lib/cb-promise.js b/lib/cb-promise.js
@@ -0,0 +1,9 @@
+// Used by mocha.js, for when we want a callback to pass to
+// a child function, but still want to return a Promise.
+module.exports = () => {
+  let cb
+  const p = new Promise((res, rej) => {
+    cb = er => er ? rej(er) : res()
+  })
+  return [cb, p]
+}
diff --git a/lib/mocha.js b/lib/mocha.js
@@ -110,19 +110,23 @@ function after (name, fn, options) {
   suite.after.push([name, fn, options])
 }
 
+const cbPromise = require('./cb-promise.js')
+
 function moment (when, fn) {
   const t = tapStack[ tapStack.length - 1 ]
   // need function because 'this' tells us which tap object
   // has the tapMochaTest thing in its options object
-  t[when](function (cb) {
+  t[when](function () {
     if (!this.options.tapMochaTest)
-      return cb()
+      return
     const suite = suiteStack[ suiteStack.length - 1 ]
+
+    const [cb, p] = cbPromise()
     const ret = fn.call(this, cb)
     if (ret && ret.then)
       return ret
-    else if (fn.length === 0)
-      return cb()
+    else if (fn.length !== 0)
+      return p
   })
 }
 
diff --git a/test/cb-promise.js b/test/cb-promise.js
@@ -0,0 +1,13 @@
+const t = require('../')
+const cbPromise = require('../lib/cb-promise.js')
+t.test('promise resolved when cb called', async t => {
+  const [cb, p] = cbPromise()
+  cb()
+  return p
+})
+t.test('promise rejects when cb failed', t => {
+  const [cb, p] = cbPromise()
+  const poop = new Error('poop')
+  cb(poop)
+  return t.rejects(p, poop)
+})
