Add resolves() and rejects() method

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,21 @@ 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 `resolves()` and `rejects()` 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
+```
+
+```typescript
+calculator.heavyOperation(1, 2).rejects(new Error());
+//same as calculator.heavyOperation(1, 2).returns(Promise.reject(new Error()));
+console.log(await calculator.heavyOperation(1, 2)); //throws Error
+```
+
 ## 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();
 

spec/throws.spec.ts

@@ -5,6 +5,7 @@ import { Substitute, Arg } from '../src/index'
 interface Calculator {
   add(a: number, b: number): number
   divide(a: number, b: number): number
+  heavyOperation(): Promise<number>
   mode: boolean
   fakeSetting: boolean
 }
@@ -16,6 +17,13 @@ test('throws on a method with arguments', t => {
   t.throws(() => calculator.divide(1, 0), { instanceOf: Error, message: 'Cannot divide by 0' })
 })
 
+test('rejects on a method', async t => {
+  const calculator = Substitute.for<Calculator>()
+  calculator.heavyOperation().rejects(new Error('Error'))
+
+  await t.throwsAsync(calculator.heavyOperation, { instanceOf: Error, message: 'Error' })
+})
+
 test('throws on a property being called', t => {
   const calculator = Substitute.for<Calculator>()
   calculator.mode.throws(new Error('Property not set'))

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,16 @@ type MockObjectMixin<TArguments extends any[], TReturnType> = BaseMockObjectMixi
     mimicks: (func: (...args: TArguments) => TReturnType) => void;
 }
 
+type NoArgumentMockObjectPromise<TReturnType> = NoArgumentMockObjectMixin<TReturnType> & {
+    resolves: (...args: Unpacked<TReturnType>[]) => void;
+    rejects: (exception: any) => void;
+}
+
+type MockObjectPromise<TArguments extends any[], TReturnType> = MockObjectMixin<TArguments, TReturnType> & {
+    resolves: (...args: Unpacked<TReturnType>[]) => void;
+    rejects: (exception: any) => 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,9 @@ export enum SubstituteMethods {
     didNotReceive = 'didNotReceive',
     mimicks = 'mimicks',
     throws = 'throws',
-    returns = 'returns'
+    returns = 'returns',
+    resolves = 'resolves',
+    rejects = 'rejects'
 }
 
 export const Nothing = Symbol();

src/states/FunctionState.ts

@@ -130,16 +130,27 @@ export class FunctionState implements ContextState {
             }
         }
 
-        if (property === SubstituteMethods.returns) {
+        if (property === SubstituteMethods.returns
+            || property === SubstituteMethods.resolves
+            || property === SubstituteMethods.rejects
+        ) {
             return (...returns: any[]) => {
                 if (!this._lastArgs) {
                     throw SubstituteException.generic('Eh, there\'s a bug, no args recorded for this return :/')
                 }
-                this.returns.push({
-                    returnValues: returns,
-                    returnIndex: 0,
-                    args: this._lastArgs
-                })
+                const returnMock: Partial<ReturnMock> = { returnIndex: 0, args: this._lastArgs };
+                switch (property) {
+                    case SubstituteMethods.returns:
+                        returnMock.returnValues = returns;
+                        break;
+                    case SubstituteMethods.resolves:
+                        returnMock.returnValues = returns.map(value => Promise.resolve(value));
+                        break;
+                    case SubstituteMethods.rejects:
+                        returnMock.returnValues = returns.map(value => Promise.reject(value));
+                        break;
+                }
+                this.returns.push(<ReturnMock>returnMock);
                 this._calls.pop()
 
                 if (this.callCount === 0) {