diff --git a/packages/utils/src/task.js b/packages/utils/src/task.js index a6a82995c8..f9873f5c25 100644 --- a/packages/utils/src/task.js +++ b/packages/utils/src/task.js @@ -13,24 +13,28 @@ export const chainFn = function chainFn (base, fn) { if (typeof fn !== 'function') { return base } - return function (...args) { - if (typeof base !== 'function') { - return fn.apply(this, args) + + if (typeof base !== 'function') { + return fn + } + + return function (arg0, ...args) { + const next = (previous = arg0) => { + const fnResult = fn.call(this, previous, ...args) + + if (fnResult && typeof fnResult.then === 'function') { + return fnResult.then(res => res || previous) + } + + return fnResult || previous } - let baseResult = base.apply(this, args) - // Allow function to mutate the first argument instead of returning the result - if (baseResult === undefined) { - [baseResult] = args + + const baseResult = base.call(this, arg0, ...args) + + if (baseResult && typeof baseResult.then === 'function') { + return baseResult.then(res => next(res)) } - const fnResult = fn.call( - this, - baseResult, - ...Array.prototype.slice.call(args, 1) - ) - // Return mutated argument if no result was returned - if (fnResult === undefined) { - return baseResult - } - return fnResult + + return next(baseResult) } } diff --git a/packages/utils/test/task.test.js b/packages/utils/test/task.test.js index 69cbbcdecc..e2bc390701 100644 --- a/packages/utils/test/task.test.js +++ b/packages/utils/test/task.test.js @@ -103,4 +103,15 @@ describe('util: task', () => { const chainedFn = chainFn(firstFn, secondFn) expect(chainedFn({}, 10)).toEqual({ bar: 12 }) }) + + test('chainFn (promise)', async () => { + const firstFn = () => Promise.resolve({ foo: 1 }) + const secondFn = function (obj) { + obj.foo++ + return Promise.resolve() + } + + const chainedFn = chainFn(firstFn, secondFn) + expect(await chainedFn()).toEqual({ foo: 2 }) + }) })