From 36fa79f75b623275d62bd875ee5d735e007834ec Mon Sep 17 00:00:00 2001
From: Robert Chu <allegrormc@gmail.com>
Date: Fri, 25 Oct 2024 08:59:47 -0700
Subject: [PATCH] Adds support for write calls with defined outputs.

---
 src/compat/waffle.ts | 27 +++++++++++++++------------
 src/mock-contract.ts | 10 ++++++++--
 2 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/src/compat/waffle.ts b/src/compat/waffle.ts
index 893a41c..6493dbe 100644
--- 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
index b81043f..c6f26e4 100644
--- 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;