Allow stubbing all getBlock and encodeFunctionData calls by omitt…

…ing params.
delvtech · Nov 4, 2024 · 4bb4281 · 4bb4281
1 parent 25a1e75
commit 4bb4281
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 19 deletions.
diff --git a/packages/drift/src/adapter/MockAdapter.test.ts b/packages/drift/src/adapter/MockAdapter.test.ts
@@ -47,7 +47,6 @@ describe("MockAdapter", () => {
       } catch (e) {
         error = e;
       }
-      console.log("ERROR", error);
       expect(error).toBeInstanceOf(Error);
     });
 
@@ -59,6 +58,7 @@ describe("MockAdapter", () => {
       };
       adapter.onGetBlock().resolves(block);
       expect(await adapter.getBlock()).toBe(block);
+      expect(await adapter.getBlock({ blockNumber: 123n })).toBe(block);
     });
 
     it("Can be stubbed with specific args", async () => {
@@ -219,13 +219,7 @@ describe("MockAdapter", () => {
 
     it("Can be stubbed", async () => {
       const adapter = new MockAdapter();
-      const foo = adapter
-        .onEncodeFunctionData({
-          abi: erc20.abi,
-          fn: "balanceOf",
-          args: { account: "0x" },
-        })
-        .returns("0x123");
+      adapter.onEncodeFunctionData().returns("0x123");
       expect(
         adapter.encodeFunctionData({
           abi: erc20.abi,

diff --git a/packages/drift/src/adapter/MockAdapter.ts b/packages/drift/src/adapter/MockAdapter.ts
@@ -51,11 +51,16 @@ export class MockAdapter implements ReadWriteAdapter {
   // getBlock //
 
   onGetBlock(params?: Partial<NetworkGetBlockParams>) {
-    return this.stubs
-      .get<[NetworkGetBlockParams?], Promise<Block | undefined>>({
-        method: "getBlock",
-      })
-      .withArgs(params);
+    let stub = this.stubs.get<
+      [NetworkGetBlockParams?],
+      Promise<Block | undefined>
+    >({
+      method: "getBlock",
+    });
+    if (params) {
+      stub = stub.withArgs(params);
+    }
+    return stub;
   }
 
   async getBlock(params?: NetworkGetBlockParams) {
@@ -142,18 +147,24 @@ export class MockAdapter implements ReadWriteAdapter {
     TAbi extends Abi,
     TFunctionName extends FunctionName<TAbi>,
   >(
-    params: OptionalKeys<
+    params?: OptionalKeys<
       AdapterEncodeFunctionDataParams<TAbi, TFunctionName>,
       "args"
     >,
   ) {
-    return this.stubs
-      .get<[AdapterEncodeFunctionDataParams<TAbi, TFunctionName>], Bytes>({
-        method: "encodeFunctionData",
-      })
-      .withArgs(
+    let stub = this.stubs.get<
+      [AdapterEncodeFunctionDataParams<TAbi, TFunctionName>],
+      Bytes
+    >({
+      method: "encodeFunctionData",
+    });
+
+    if (params) {
+      stub = stub.withArgs(
         params as Partial<AdapterEncodeFunctionDataParams<TAbi, TFunctionName>>,
       );
+    }
+    return stub;
   }
 
   encodeFunctionData<