From 969b6a360cfe6b55f643ca0b30d84360b0134190 Mon Sep 17 00:00:00 2001
From: Robert Chu <allegrormc@gmail.com>
Date: Thu, 1 Aug 2024 14:47:51 -0700
Subject: [PATCH] Adds MockWriteCallExpectation.

---
 src/index.ts         |  1 +
 src/mock-contract.ts | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/src/index.ts b/src/index.ts
index 16f7356..1c9e484 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -1,5 +1,6 @@
 export {
   MockReadCallExpectation,
+  MockWriteCallExpectation,
   MockRevertExpectation,
   MockCallExpectation,
   MockContractController,
diff --git a/src/mock-contract.ts b/src/mock-contract.ts
index 34d2670..098e6cc 100644
--- a/src/mock-contract.ts
+++ b/src/mock-contract.ts
@@ -17,6 +17,11 @@ export type MockReadCallExpectation<T extends AbiFunction> = {
   inputs: AbiParametersToPrimitiveTypes<T["inputs"]>;
   outputs: AbiParametersToPrimitiveTypes<T["outputs"]>;
 };
+export type MockWriteCallExpectation<T extends AbiFunction> = {
+  kind: "write";
+  abi: T;
+  inputs: AbiParametersToPrimitiveTypes<T["inputs"]>;
+};
 export type MockRevertExpectation<T extends AbiFunction> = {
   kind: "revert";
   abi: T;
@@ -25,6 +30,7 @@ export type MockRevertExpectation<T extends AbiFunction> = {
 };
 export type MockCallExpectation<T extends AbiFunction> =
   | MockReadCallExpectation<T>
+  | MockWriteCallExpectation<T>
   | MockRevertExpectation<T>;
 
 export type MockContractController = {
@@ -97,6 +103,35 @@ export const deployMock = async (
             }
             break;
           }
+          case "write": {
+            const fnSigHash = encodeFunctionData({
+              abi: [call.abi as AbiFunction],
+              args: call.inputs,
+              functionName: call.abi.name,
+            });
+            // Use a mock function to return the expected return value
+            if (firstCall) {
+              await signer.writeContract({
+                address,
+                chain: signer.chain,
+                account: signer.account,
+                abi: abi,
+                functionName: "__doppelganger__mockReturns",
+                args: [fnSigHash, "0x"],
+              });
+              firstCall = false;
+            } else {
+              await signer.writeContract({
+                address,
+                chain: signer.chain,
+                account: signer.account,
+                abi: abi,
+                functionName: "__doppelganger__queueReturn",
+                args: [fnSigHash, "0x"],
+              });
+            }
+            break;
+          }
           case "revert": {
             const fnSigHash = encodeFunctionData({
               abi: [call.abi as AbiFunction],