From 442fcd1ffd769cc448a3e964d42b16325ea9ae49 Mon Sep 17 00:00:00 2001
From: dominikeinkemmer <dominik_einkemmer@gmx.at>
Date: Tue, 1 Jan 2019 02:02:59 +0900
Subject: [PATCH] add support for chained calls in typings (#31695)

---
 types/jest-when/index.d.ts         | 22 ++++++++++-----------
 types/jest-when/jest-when-tests.ts | 31 ++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 11 deletions(-)

diff --git a/types/jest-when/index.d.ts b/types/jest-when/index.d.ts
index f484502f4d..0579419425 100644
--- a/types/jest-when/index.d.ts
+++ b/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;
diff --git a/types/jest-when/jest-when-tests.ts b/types/jest-when/jest-when-tests.ts
index 13da65928d..179ea2fae1 100644
--- a/types/jest-when/jest-when-tests.ts
+++ b/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');