fix test

smartcontractkit · Mar 26, 2024 · b546312 · b546312
1 parent 8f7ea0a
commit b546312
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 26 deletions.
diff --git a/contracts/src/v0.8/automation/dev/test/AutomationRegistry2_3.t.sol b/contracts/src/v0.8/automation/dev/test/AutomationRegistry2_3.t.sol
@@ -146,32 +146,34 @@ contract AddFunds is SetUp {
   }
 
   function test_anyoneCanAddFunds() public {
-    assertEq(registry.getBalance(linkUpkeepID), 0);
+    uint256 startAmount = registry.getBalance(linkUpkeepID);
     vm.prank(UPKEEP_ADMIN);
     registry.addFunds(linkUpkeepID, 1);
-    assertEq(registry.getBalance(linkUpkeepID), 1);
+    assertEq(registry.getBalance(linkUpkeepID), startAmount + 1);
     vm.prank(STRANGER);
     registry.addFunds(linkUpkeepID, 1);
-    assertEq(registry.getBalance(linkUpkeepID), 2);
+    assertEq(registry.getBalance(linkUpkeepID), startAmount + 2);
   }
 
   function test_movesFundFromCorrectToken() public {
     vm.startPrank(UPKEEP_ADMIN);
 
     uint256 startBalanceLINK = linkToken.balanceOf(address(registry));
     uint256 startBalanceUSDToken = usdToken.balanceOf(address(registry));
+    uint256 startLinkUpkeepBalance = registry.getBalance(linkUpkeepID);
+    uint256 startUSDUpkeepBalance = registry.getBalance(usdUpkeepID);
 
     registry.addFunds(linkUpkeepID, 1);
-    assertEq(registry.getBalance(linkUpkeepID), 1);
-    assertEq(registry.getBalance(usdUpkeepID), 0);
-    assertEq(linkToken.balanceOf(address(registry)), startBalanceLINK + 1);
-    assertEq(usdToken.balanceOf(address(registry)), startBalanceUSDToken);
+    assertEq(registry.getBalance(linkUpkeepID), startBalanceLINK + 1);
+    assertEq(registry.getBalance(usdUpkeepID), startBalanceUSDToken);
+    assertEq(linkToken.balanceOf(address(registry)), startLinkUpkeepBalance + 1);
+    assertEq(usdToken.balanceOf(address(registry)), startUSDUpkeepBalance);
 
     registry.addFunds(usdUpkeepID, 2);
-    assertEq(registry.getBalance(linkUpkeepID), 1);
-    assertEq(registry.getBalance(usdUpkeepID), 2);
-    assertEq(linkToken.balanceOf(address(registry)), startBalanceLINK + 1);
-    assertEq(usdToken.balanceOf(address(registry)), startBalanceUSDToken + 2);
+    assertEq(registry.getBalance(linkUpkeepID), startBalanceLINK + 1);
+    assertEq(registry.getBalance(usdUpkeepID), startBalanceUSDToken + 2);
+    assertEq(linkToken.balanceOf(address(registry)), startLinkUpkeepBalance + 1);
+    assertEq(usdToken.balanceOf(address(registry)), startUSDUpkeepBalance + 2);
   }
 
   function test_emitsAnEvent() public {
@@ -183,17 +185,20 @@ contract AddFunds is SetUp {
 }
 
 contract Withdraw is SetUp {
-  address internal aMockAddress = address(0x1111111111111111111111111111111111111113);
+  address internal aMockAddress = randomAddress();
 
   function testLinkAvailableForPaymentReturnsLinkBalance() public {
+    uint256 startBalance = linkToken.balanceOf(address(registry));
+    int256 startLinkAvailable = registry.linkAvailableForPayment();
+
     //simulate a deposit of link to the liquidity pool
     _mintLink(address(registry), 1e10);
 
     //check there's a balance
-    assertGt(linkToken.balanceOf(address(registry)), 0);
+    assertEq(linkToken.balanceOf(address(registry)), startBalance + 1e10);
 
-    //check the link available for payment is the link balance
-    assertEq(uint256(registry.linkAvailableForPayment()), linkToken.balanceOf(address(registry)));
+    //check the link available has increased by the same amount
+    assertEq(uint256(registry.linkAvailableForPayment()), uint256(startLinkAvailable) + 1e10);
   }
 
   function testWithdrawLinkRevertsBecauseOnlyFinanceAdminAllowed() public {
@@ -224,9 +229,7 @@ contract Withdraw is SetUp {
   function testWithdrawLinkSuccess() public {
     //simulate a deposit of link to the liquidity pool
     _mintLink(address(registry), 1e10);
-
-    //check there's a balance
-    assertGt(linkToken.balanceOf(address(registry)), 0);
+    uint256 startBalance = linkToken.balanceOf(address(registry));
 
     vm.startPrank(FINANCE_ADMIN);
 
@@ -236,7 +239,7 @@ contract Withdraw is SetUp {
     vm.stopPrank();
 
     assertEq(linkToken.balanceOf(address(aMockAddress)), 1);
-    assertEq(linkToken.balanceOf(address(registry)), 1e10 - 1);
+    assertEq(linkToken.balanceOf(address(registry)), startBalance - 1);
   }
 
   function test_WithdrawERC20Fees_RespectsReserveAmount() public {
@@ -247,21 +250,24 @@ contract Withdraw is SetUp {
   }
 
   function testWithdrawERC20FeeSuccess() public {
-    // simulate a deposit of ERC20 to the liquidity pool
+    // deposit excess USDToken to the registry (this goes to the "finance withdrawable" pool be default)
+    uint256 startReserveAmount = registry.getReserveAmount(address(usdToken));
+    uint256 startAmount = usdToken.balanceOf(address(registry));
     _mintERC20(address(registry), 1e10);
 
-    // check there's a balance
-    assertGt(usdToken.balanceOf(address(registry)), 0);
+    // depositing shouldn't change reserve amount
+    assertEq(registry.getReserveAmount(address(usdToken)), startReserveAmount);
 
     vm.startPrank(FINANCE_ADMIN);
 
-    // try to withdraw 1 link while there is a ton of link available
+    // try to withdraw 1 USDToken
     registry.withdrawERC20Fees(address(usdToken), aMockAddress, 1);
 
     vm.stopPrank();
 
     assertEq(usdToken.balanceOf(address(aMockAddress)), 1);
-    assertEq(usdToken.balanceOf(address(registry)), 1e10 - 1);
+    assertEq(usdToken.balanceOf(address(registry)), startAmount + 1e10 - 1);
+    assertEq(registry.getReserveAmount(address(usdToken)), startReserveAmount);
   }
 }
 

diff --git a/contracts/test/v0.8/automation/AutomationRegistry2_3.test.ts b/contracts/test/v0.8/automation/AutomationRegistry2_3.test.ts
@@ -4636,7 +4636,7 @@ describe('AutomationRegistry2_3', () => {
   describe('#withdrawOwnerFunds', () => {
     it('can only be called by finance admin', async () => {
       await evmRevert(
-        registry.connect(keeper1).withdrawLinkFees(zeroAddress, 1),
+        registry.connect(keeper1).withdrawLink(zeroAddress, 1),
         'OnlyFinanceAdmin()',
       )
     })
@@ -4693,7 +4693,7 @@ describe('AutomationRegistry2_3', () => {
       // Now withdraw
       await registry
         .connect(financeAdmin)
-        .withdrawLinkFees(await owner.getAddress(), ownerRegistryBalance)
+        .withdrawLink(await owner.getAddress(), ownerRegistryBalance)
 
       ownerRegistryBalance = await registry.linkAvailableForPayment()
       const ownerAfter = await linkToken.balanceOf(await owner.getAddress())