Adds support for write calls with defined outputs.

term-finance · Oct 25, 2024 · 36fa79f · 36fa79f
1 parent bf6b94b
commit 36fa79f
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 14 deletions.
diff --git a/src/compat/waffle.ts b/src/compat/waffle.ts
@@ -146,18 +146,21 @@ class Stub<T extends AbiFunction> implements StubInterface {
     if (!this.func.outputs)
       this.err("Cannot mock return values from a void function");
 
-    const kind =
-      this.func.stateMutability === "view" ||
-      this.func.stateMutability === "pure"
-        ? "read"
-        : "write";
-
-    this.calls.push({
-      kind,
-      abi: this.func,
-      inputs: this.inputs,
-      outputs: args,
-    });
+    if (this.func.stateMutability === "pure" || this.func.stateMutability === "view") {
+      this.calls.push({
+        kind: "read",
+        abi: this.func,
+        inputs: this.inputs,
+        outputs: args,
+      });
+    } else {
+      this.calls.push({
+        kind: "write",
+        abi: this.func,
+        inputs: this.inputs,
+        outputs: args.length === 0 ? undefined : args,
+      });
+    }
 
     return this;
   }

diff --git a/src/mock-contract.ts b/src/mock-contract.ts
@@ -22,6 +22,7 @@ export type MockWriteCallExpectation<T extends AbiFunction> = {
   kind: "write";
   abi: T;
   inputs?: AbiParametersToPrimitiveTypes<T["inputs"]>;
+  outputs?: AbiParametersToPrimitiveTypes<T["outputs"]>;
 };
 export type MockRevertExpectation<T extends AbiFunction> = {
   kind: "revert";
@@ -118,6 +119,11 @@ export const deployMock = async (
           }
           case "write": {
             const fnSigHash = calculateFnSigHash(call);
+            const encodedOutputs = call.outputs ? encodeFunctionResult({
+              abi: [call.abi as AbiFunction],
+              functionName: call.abi.name,
+              result: call.outputs,
+            }) : "0x";
             // Use a mock function to return the expected return value
             if (firstCall) {
               await signer.writeContract({
@@ -126,7 +132,7 @@ export const deployMock = async (
                 account: signer.account,
                 abi: abi,
                 functionName: "__doppelganger__mockReturns",
-                args: [fnSigHash, "0x"],
+                args: [fnSigHash, encodedOutputs],
               });
               firstCall = false;
             } else {
@@ -136,7 +142,7 @@ export const deployMock = async (
                 account: signer.account,
                 abi: abi,
                 functionName: "__doppelganger__queueReturn",
-                args: [fnSigHash, "0x"],
+                args: [fnSigHash, encodedOutputs],
               });
             }
             break;