2512 add estimate/transact to withdrawal manager completeQueuedWithdr…

…awal
PufferFinance · Sep 19, 2024 · 33560ba · 33560ba
1 parent cbae81a
commit 33560ba
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 6 deletions.
diff --git a/lib/contracts/handlers/puffer-withdrawal-manager.test.ts b/lib/contracts/handlers/puffer-withdrawal-manager.test.ts
@@ -180,7 +180,12 @@ describe('PufferWithdrawalHandler', () => {
     });
 
     const withdrawalIdx = BigInt(5);
-    await handler.completeQueueWithdrawal(mockAddress, withdrawalIdx);
+    const { transact } = await handler.completeQueueWithdrawal(
+      mockAddress,
+      withdrawalIdx,
+    );
+
+    await transact();
 
     expect(
       handler['getContract']().write.completeQueuedWithdrawal,
@@ -190,6 +195,29 @@ describe('PufferWithdrawalHandler', () => {
     );
   });
 
+  it('should estimate gas for a queued withdrawal', async () => {
+    jest.spyOn(handler as any, 'getContract').mockReturnValue({
+      estimateGas: {
+        completeQueuedWithdrawal: jest.fn().mockResolvedValue(BigInt(100000)),
+      },
+    });
+
+    const withdrawalIdx = BigInt(5);
+    const { estimate } = await handler.completeQueueWithdrawal(
+      mockAddress,
+      withdrawalIdx,
+    );
+
+    await estimate();
+
+    expect(
+      handler['getContract']().estimateGas.completeQueuedWithdrawal,
+    ).toHaveBeenCalledWith(
+      expect.arrayContaining([withdrawalIdx]),
+      expect.any(Object),
+    );
+  });
+
   it('should get a withdrawal by its index', async () => {
     const withdrawalIdx = BigInt(1);
     const mockWithdrawal = { amount: '10', recipient: mockAddress };

diff --git a/lib/contracts/handlers/puffer-withdrawal-manager.ts b/lib/contracts/handlers/puffer-withdrawal-manager.ts
@@ -143,13 +143,19 @@ export class PufferWithdrawalHandler {
     walletAddress: Address,
     withdrawalIdx: bigint,
   ) {
-    return await this.getContract().write.completeQueuedWithdrawal(
-      [withdrawalIdx],
-      {
+    const transact = async () =>
+      await this.getContract().write.completeQueuedWithdrawal([withdrawalIdx], {
         account: walletAddress,
         chain: this.viemChain,
-      },
-    );
+      });
+
+    const estimate = async () =>
+      await this.getContract().estimateGas.completeQueuedWithdrawal(
+        [withdrawalIdx],
+        { account: walletAddress },
+      );
+
+    return { transact, estimate };
   }
 
   /**