Add resolves() method

Author: abdala

README.md

@@ -18,6 +18,7 @@ interface Calculator {
   add(a: number, b: number): number;
   subtract(a: number, b: number): number;
   divide(a: number, b: number): number;
+  async heavyOperation(): Promise<number>;
 
   isEnabled: boolean;
 }
@@ -53,6 +54,15 @@ console.log(calculator.add(1, 2)); //prints 9
 console.log(calculator.add(1, 2)); //prints undefined
 ```
 
+## Working with promises
+When working with promises you can also use the method `resolves()` to return a promise.
+
+```typescript
+calculator.heavyOperation(1, 2).resolves(4); 
+//same as calculator.heavyOperation(1, 2).returns(Promise.resolve(4));
+console.log(await calculator.heavyOperation(1, 2)); //prints 4
+```
+
 ## Verifying calls
 ```typescript
 calculator.enabled = true;

spec/index.test.ts

@@ -174,6 +174,22 @@ test('returning resolved promises works', async t => {
 	t.is(1338, await substitute.returnPromise());
 });
 
+test('resolving other fake works', async t => {
+	initialize();
+
+	const otherSubstitute = Substitute.for<Dummy>();
+	substitute.returnPromise().resolves(otherSubstitute);
+	t.is(otherSubstitute, await substitute.returnPromise());
+});
+
+test('resolving promises works', async t => {
+	initialize();
+
+	substitute.returnPromise().resolves(1338);
+
+	t.is(1338, await substitute.returnPromise());
+});
+
 test('class void returns', t => {
 	initialize();
 

src/Transformations.ts

@@ -1,14 +1,23 @@
 import { AllArguments } from "./Arguments";
 
 export type NoArgumentFunctionSubstitute<TReturnType> =
-    (() => (TReturnType & NoArgumentMockObjectMixin<TReturnType>))
+    TReturnType extends Promise<infer U> ?
+        (() => (TReturnType & NoArgumentMockObjectPromise<TReturnType>)) :
+        (() => (TReturnType & NoArgumentMockObjectMixin<TReturnType>))
 
 export type FunctionSubstitute<TArguments extends any[], TReturnType> =
-    ((...args: TArguments) => (TReturnType & MockObjectMixin<TArguments, TReturnType>)) &
-    ((allArguments: AllArguments) => (TReturnType & MockObjectMixin<TArguments, TReturnType>))
+    TReturnType extends Promise<infer U> ?
+        ((...args: TArguments) => (TReturnType & MockObjectPromise<TArguments, TReturnType>)) &
+        ((allArguments: AllArguments) => (TReturnType & MockObjectPromise<TArguments, TReturnType>)) :
+        ((...args: TArguments) => (TReturnType & MockObjectMixin<TArguments, TReturnType>)) &
+        ((allArguments: AllArguments) => (TReturnType & MockObjectMixin<TArguments, TReturnType>))
 
 export type PropertySubstitute<TReturnType> = (TReturnType & Partial<NoArgumentMockObjectMixin<TReturnType>>);
 
+type Unpacked<T> =
+    T extends Promise<infer U> ? U :
+    T;
+
 type BaseMockObjectMixin<TReturnType> = {
     returns: (...args: TReturnType[]) => void;
     throws: (exception: any) => void;
@@ -22,6 +31,14 @@ type MockObjectMixin<TArguments extends any[], TReturnType> = BaseMockObjectMixi
     mimicks: (func: (...args: TArguments) => TReturnType) => void;
 }
 
+type NoArgumentMockObjectPromise<TReturnType> = NoArgumentMockObjectMixin<TReturnType> & {
+    resolves: (...args: Unpacked<TReturnType>[]) => void;
+}
+
+type MockObjectPromise<TArguments extends any[], TReturnType> = MockObjectMixin<TArguments, TReturnType> & {
+    resolves: (...args: Unpacked<TReturnType>[]) => void;
+}
+
 export type ObjectSubstitute<T extends Object, K extends Object = T> = ObjectSubstituteTransformation<T> & {
     received(amount?: number): TerminatingObject<K>;
     didNotReceive(): TerminatingObject<K>;

src/Utilities.ts

@@ -16,7 +16,8 @@ export enum SubstituteMethods {
     didNotReceive = 'didNotReceive',
     mimicks = 'mimicks',
     throws = 'throws',
-    returns = 'returns'
+    returns = 'returns',
+    resolves = 'resolves'
 }
 
 export const Nothing = Symbol();

src/states/FunctionState.ts

@@ -156,6 +156,23 @@ export class FunctionState implements ContextState {
             };
         }
 
+        if (property === SubstituteMethods.resolves) {
+            return (...returns: any[]) => {
+                if (!this._lastArgs) {
+                    throw SubstituteException.generic('Eh, there\'s a bug, no args recorded for this return :/')
+                }
+                returns = returns.map(value => Promise.resolve(value))
+                this.returns.push({
+                    returnValues: returns,
+                    returnIndex: 0,
+                    args: this._lastArgs
+                })
+                this._calls.pop()
+
+                context.state = context.initialState;
+            };
+        }
+
         return context.proxy;
     }
 }