From 83f2b621b46afbb1633b224a10a3f2264ce0356f Mon Sep 17 00:00:00 2001
From: nicholaspai <npai.nyc@gmail.com>
Date: Sat, 30 Nov 2024 21:10:44 -0500
Subject: [PATCH] throw error if can't fund the DonationBox or
 loadEthForL2Calls call

---
 src/dataworker/Dataworker.ts             | 15 +++++++++
 test/Dataworker.executePoolRebalances.ts | 43 ++++++++++++++++++++++--
 2 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/src/dataworker/Dataworker.ts b/src/dataworker/Dataworker.ts
index d1cdae38f..87b9a3319 100644
--- a/src/dataworker/Dataworker.ts
+++ b/src/dataworker/Dataworker.ts
@@ -1603,6 +1603,21 @@ export class Dataworker {
           requiredAmount,
         });
         if (submitExecution) {
+          const canFund = await balanceAllocator.requestBalanceAllocations([
+            {
+              tokens: [feeToken],
+              amount: requiredAmount,
+              holder: await signer.getAddress(),
+              chainId: hubPoolChainId,
+            },
+          ]);
+          if (!canFund) {
+            throw new Error(
+              `Failed to fund ${requiredAmount.toString()} of orbit gas token ${feeToken} for message to ${getNetworkName(
+                leaf.chainId
+              )}!`
+            );
+          }
           if (feeToken === ZERO_ADDRESS) {
             this.clients.multiCallerClient.enqueueTransaction({
               contract: this.clients.hubPoolClient.hubPool,
diff --git a/test/Dataworker.executePoolRebalances.ts b/test/Dataworker.executePoolRebalances.ts
index 635e122d3..15cebc260 100644
--- a/test/Dataworker.executePoolRebalances.ts
+++ b/test/Dataworker.executePoolRebalances.ts
@@ -554,7 +554,7 @@ describe("Dataworker: Execute pool rebalances", async function () {
         )
       ).to.equal(toBNWei("-1"));
     });
-    it("contains arbitrum leaf", async function () {
+    it("funds arbitrum leaf", async function () {
       // Adds one fee per net send amount + one extra if groupIndex = 0
       const leaves: PoolRebalanceLeaf[] = [
         {
@@ -598,7 +598,7 @@ describe("Dataworker: Execute pool rebalances", async function () {
       expect(queuedTransactions[2].method).to.equal("executeRootBundle");
       expect(queuedTransactions[3].method).to.equal("executeRootBundle");
     });
-    it("contains custom gas token orbit leaf", async function () {
+    it("funds custom gas token orbit leaf", async function () {
       // Replicate custom gas token setups:
       const azero = await smock.fake(ERC20.abi, {
         address: TOKEN_SYMBOLS_MAP.AZERO.addresses[CHAIN_IDs.MAINNET],
@@ -634,6 +634,9 @@ describe("Dataworker: Execute pool rebalances", async function () {
       const expectedFee = toBNWei("0.49");
       const expectedFeeLeaf1 = expectedFee.mul(2).add(expectedFee);
       const expectedFeeLeaf2 = expectedFee.mul(2);
+      azero.balanceOf
+        .whenCalledWith(await hubPoolClient.hubPool.signer.getAddress())
+        .returns(expectedFeeLeaf1.add(expectedFeeLeaf2));
       await dataworkerInstance._executePoolRebalanceLeaves(
         spokePoolClients,
         leaves,
@@ -652,5 +655,41 @@ describe("Dataworker: Execute pool rebalances", async function () {
       expect(queuedTransactions[2].method).to.equal("executeRootBundle");
       expect(queuedTransactions[3].method).to.equal("executeRootBundle");
     });
+    it("fails to fund custom gas token orbit leaf", async function () {
+      // Replicate custom gas token setups, but this time do not set a balance for the custom gas token funder.
+      const azero = await smock.fake(ERC20.abi, {
+        address: TOKEN_SYMBOLS_MAP.AZERO.addresses[CHAIN_IDs.MAINNET],
+        provider: hubPoolClient.hubPool.signer.provider,
+      });
+      // Custom gas token funder for AZERO
+      const customGasTokenFunder = "0x0d57392895Db5aF3280e9223323e20F3951E81B1";
+      azero.balanceOf.whenCalledWith(customGasTokenFunder).returns(0);
+      expect(await balanceAllocator.getBalance(hubPoolClient.chainId, azero.address, customGasTokenFunder)).to.equal(0);
+
+      // Adds one fee per net send amount + one extra if groupIndex = 0
+      const leaves: PoolRebalanceLeaf[] = [
+        {
+          chainId: 41455,
+          groupIndex: 0,
+          bundleLpFees: [toBNWei("1"), toBNWei("1")],
+          netSendAmounts: [toBNWei("1"), toBNWei("1")],
+          runningBalances: [toBNWei("1"), toBNWei("1")],
+          leafId: 0,
+          l1Tokens: [randomAddress(), randomAddress()],
+        },
+      ];
+      // Should throw an error if caller doesn't have enough custom gas token to fund
+      // DonationBox.
+      await assertPromiseError(
+        dataworkerInstance._executePoolRebalanceLeaves(
+          spokePoolClients,
+          leaves,
+          balanceAllocator,
+          buildPoolRebalanceLeafTree(leaves),
+          true
+        ),
+        "Failed to fund"
+      );
+    });
   });
 });