Adding async utility versions of mockReturnValue & mockreturnValueOnce

Author: m4n3z40

packages/jest-mock/src/__tests__/jest_mock.test.js

@@ -391,6 +391,55 @@ describe('moduleMocker', () => {
       expect(fake(2)).toEqual(4);
     });
 
+    it('supports mocking resolvable async functions', () => {
+      const fn = moduleMocker.fn();
+      fn.mockResolvedValue('abcd');
+
+      const promise = fn();
+
+      expect(promise).toBeInstanceOf(Promise);
+
+      return promise.then(value => expect(value).toBe('abcd'));
+    });
+
+    it('supports mocking resolvable async functions only once', () => {
+      const fn = moduleMocker.fn();
+      fn.mockResolvedValue('abcd');
+      fn.mockResolvedValueOnce('abcde');
+
+      const promise1 = fn().then(value => expect(value).toBe('abcde'));
+      const promise2 = fn().then(value => expect(value).toBe('abcd'));
+
+      return Promise.all([promise1, promise2]);
+    });
+
+    it('supports mocking rejectable async functions', () => {
+      const err = new Error('rejected');
+      const fn = moduleMocker.fn();
+      fn.mockRejectedValue(err);
+
+      const promise = fn();
+
+      expect(promise).toBeInstanceOf(Promise);
+
+      return promise.catch(rejection => expect(rejection).toBe(err));
+    });
+
+    it('supports mocking rejectable async functions only once', () => {
+      const defaultErr = new Error('default rejected');
+      const err = new Error('rejected');
+      const fn = moduleMocker.fn();
+      fn.mockRejectedValue(defaultErr);
+      fn.mockRejectedValueOnce(err);
+
+      const promise1 = fn().catch(rejection => expect(rejection).toBe(err));
+      const promise2 = fn().catch(rejection =>
+        expect(rejection).toBe(defaultErr),
+      );
+
+      return Promise.all([promise1, promise2]);
+    });
+
     describe('timestamps', () => {
       const RealDate = Date;
 

packages/jest-mock/src/index.js

@@ -227,6 +227,17 @@ function getSlots(object?: Object): Array<string> {
   return Object.keys(slots);
 }
 
+function wrapAsyncParam(
+  fn: any => any,
+  asyncAction: 'resolve' | 'reject',
+): any => any {
+  if (asyncAction === 'reject') {
+    return value => fn(Promise.reject(value));
+  }
+
+  return value => fn(Promise.resolve(value));
+}
+
 class ModuleMockerClass {
   _environmentGlobal: Global;
   _mockState: WeakMap<Function, MockFunctionState>;
@@ -407,6 +418,13 @@ class ModuleMockerClass {
         return f;
       };
 
+      f.mockResolvedValueOnce = wrapAsyncParam(
+        f.mockReturnValueOnce,
+        'resolve',
+      );
+
+      f.mockRejectedValueOnce = wrapAsyncParam(f.mockReturnValueOnce, 'reject');
+
       f.mockReturnValue = value => {
         // next function call will return specified return value or this one
         const mockConfig = this._ensureMockConfig(f);
@@ -415,6 +433,10 @@ class ModuleMockerClass {
         return f;
       };
 
+      f.mockResolvedValue = wrapAsyncParam(f.mockReturnValue, 'resolve');
+
+      f.mockRejectedValue = wrapAsyncParam(f.mockReturnValue, 'reject');
+
       f.mockImplementationOnce = fn => {
         // next function call will use this mock implementation return value
         // or default mock implementation return value