Adding async versions of mockReturnValue & mockreturnValueOnce

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

@@ -391,6 +391,53 @@ 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 expect(promise).resolves.toBe('abcd');
+    });
+
+    it('supports mocking resolvable async functions only once', () => {
+      const fn = moduleMocker.fn();
+      fn.mockResolvedValue('abcd');
+      fn.mockResolvedValueOnce('abcde');
+
+      return Promise.all([
+        expect(fn()).resolves.toBe('abcde'),
+        expect(fn()).resolves.toBe('abcd'),
+      ]);
+    });
+
+    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 expect(promise).rejects.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);
+
+      return Promise.all([
+        expect(fn()).rejects.toBe(err),
+        expect(fn()).rejects.toBe(defaultErr),
+      ]);
+    });
+
     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