add support for chained calls in typings (#31695)

types/jest-when/index.d.ts

@@ -6,22 +6,22 @@
 
 /// <reference types="jest" />
 
-export interface PartialMockInstance<T> {
-  mockReturnValue: jest.MockInstance<T>['mockReturnValue'];
-  mockReturnValueOnce: jest.MockInstance<T>['mockReturnValueOnce'];
-  mockResolvedValue: jest.MockInstance<T>['mockResolvedValue'];
-  mockResolvedValueOnce: jest.MockInstance<T>['mockResolvedValueOnce'];
-  mockRejectedValue: jest.MockInstance<T>['mockRejectedValue'];
-  mockRejectedValueOnce: jest.MockInstance<T>['mockRejectedValueOnce'];
-}
-
 export interface When {
   <T>(fn: jest.Mocked<T> | jest.Mock<T>): When;
+
   // due to no-unnecessary-generics lint rule, the generics have been replaced with 'any'
   // calledWith<T>(...matchers: any[]): PartialMockInstance<T>;
   // expectCalledWith<T>(...matchers: any[]): PartialMockInstance<T>;
-  calledWith(...matchers: any[]): PartialMockInstance<any>;
-  expectCalledWith(...matchers: any[]): PartialMockInstance<any>;
+  calledWith(...matchers: any[]): When;
+
+  expectCalledWith(...matchers: any[]): When;
+
+  mockReturnValue: (value: any) => jest.MockInstance<any>['mockReturnValue'] & When;
+  mockReturnValueOnce: (value: any) => jest.MockInstance<any>['mockReturnValue'] & When;
+  mockResolvedValue: (value: any) => jest.MockInstance<any>['mockReturnValue'] & When;
+  mockResolvedValueOnce: (value: any) => jest.MockInstance<any>['mockReturnValue'] & When;
+  mockRejectedValue: (value: any) => jest.MockInstance<any>['mockReturnValue'] & When;
+  mockRejectedValueOnce: (value: any) => jest.MockInstance<any>['mockReturnValue'] & When;
 }
 
 export const when: When;

types/jest-when/jest-when-tests.ts

@@ -43,6 +43,37 @@ describe('mock-when test', () => {
     expect(fn(5)).toEqual(undefined);
   });
 
+  it('Supports chained calls:', () => {
+    const fn = jest.fn();
+    when(fn)
+      .calledWith(1).mockReturnValue('no')
+      .calledWith(2).mockReturnValue('way?')
+      .calledWith(3).mockReturnValue('yes')
+      .calledWith(4).mockReturnValue('way!');
+
+    expect(fn(1)).toEqual('no');
+    expect(fn(2)).toEqual('way?');
+    expect(fn(3)).toEqual('yes');
+    expect(fn(4)).toEqual('way!');
+    expect(fn(5)).toEqual(undefined);
+  });
+
+  it('Supports chained calls with defaults:', () => {
+    const fn = jest.fn();
+    when(fn)
+      .mockReturnValue('nice')
+      .calledWith(1).mockReturnValue('no')
+      .calledWith(2).mockReturnValue('way?')
+      .calledWith(3).mockReturnValue('yes')
+      .calledWith(4).mockReturnValue('way!');
+
+    expect(fn(1)).toEqual('no');
+    expect(fn(2)).toEqual('way?');
+    expect(fn(3)).toEqual('yes');
+    expect(fn(4)).toEqual('way!');
+    expect(fn(5)).toEqual('nice');
+  });
+
   it('Assert the args:', () => {
     const fn = jest.fn();
     when(fn).expectCalledWith(1).mockReturnValue('x');