@0x:contracts-integrations Addressed more review feedback

contracts/integrations/test/utils/function_assertions.ts

@@ -1,28 +1,31 @@
 import { PromiseWithTransactionHash } from '@0x/base-contract';
 import { TransactionReceiptWithDecodedLogs } from 'ethereum-types';
 
-export interface ContractGetterFunction {
-    callAsync: (...args: any[]) => Promise<any>;
-    getABIEncodedTransactionData: (...args: any[]) => string;
+export interface ContractGetterFunction<TCallArgs extends any[]> {
+    callAsync: (...args: TCallArgs) => Promise<any>;
+    getABIEncodedTransactionData: (...args: TCallArgs) => string;
 }
 
-export interface ContractWrapperFunction extends ContractGetterFunction {
-    awaitTransactionSuccessAsync?: (...args: any[]) => PromiseWithTransactionHash<TransactionReceiptWithDecodedLogs>;
+export interface ContractWrapperFunction<TAwaitArgs extends any[], TCallArgs extends any[]>
+    extends ContractGetterFunction<TCallArgs> {
+    awaitTransactionSuccessAsync?: (
+        ...args: TAwaitArgs
+    ) => PromiseWithTransactionHash<TransactionReceiptWithDecodedLogs>;
 }
 
 export interface Condition {
-    before: (...args: any[]) => Promise<void>;
-    after: (result: any, receipt: TransactionReceiptWithDecodedLogs | undefined, ...args: any[]) => Promise<void>;
+    before: (...args: any[]) => Promise<any>;
+    after: (beforeInfo: any, result: Result, ...args: any[]) => Promise<any>;
 }
 
-export class FunctionAssertion {
+export class FunctionAssertion<TAwaitArgs extends any[], TCallArgs extends any[]> {
     // A before and an after assertion that will be called around the wrapper function.
     public condition: Condition;
 
     // The wrapper function that will be wrapped in assertions.
-    public wrapperFunction: ContractWrapperFunction;
+    public wrapperFunction: ContractWrapperFunction<TAwaitArgs, TCallArgs>;
 
-    constructor(wrapperFunction: ContractWrapperFunction, condition: Condition) {
+    constructor(wrapperFunction: ContractWrapperFunction<TAwaitArgs, TCallArgs>, condition: Condition) {
         this.condition = condition;
         this.wrapperFunction = wrapperFunction;
     }
@@ -31,13 +34,33 @@ export class FunctionAssertion {
      * Runs the wrapped function and fails if the before or after assertions fail.
      * @param ...args The args to the contract wrapper function.
      */
-    public async runAsync(...args: any[]): Promise<void> {
-        await this.condition.before(...args);
-        const result = await this.wrapperFunction.callAsync(...args);
-        const receipt =
-            this.wrapperFunction.awaitTransactionSuccessAsync !== undefined
-                ? await this.wrapperFunction.awaitTransactionSuccessAsync(...args)
-                : undefined;
-        await this.condition.after(result, receipt, ...args);
+    public async runAsync(args: TAwaitArgs & TCallArgs): Promise<{ beforeInfo: any; afterInfo: any }> {
+        // Call the before condition.
+        const beforeInfo = await this.condition.before(...args);
+
+        // Initialize the callResult so that the default success value is true.
+        let callResult: Result = { success: true };
+
+        // Try to make the call to the function. If it is successful, pass the
+        // result and receipt to the after condition.
+        try {
+            callResult.data = await this.wrapperFunction.callAsync(...args);
+            callResult.receipt =
+                this.wrapperFunction.awaitTransactionSuccessAsync !== undefined
+                    ? await this.wrapperFunction.awaitTransactionSuccessAsync(...args)
+                    : undefined;
+        } catch (error) {
+            callResult.data = error;
+            callResult.success = false;
+            callResult.receipt = undefined;
+        }
+
+        // Call the after condition.
+        const afterInfo = await this.condition.after(beforeInfo, callResult, ...args);
+
+        return {
+            beforeInfo,
+            afterInfo,
+        };
     }
 }