Simplify testing scenario

immutable · Feb 27, 2024 · 5a47e0f · 5a47e0f
1 parent 1784c91
commit 5a47e0f
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 59 deletions.
diff --git a/test/invariant/InvariantBridge.t.sol b/test/invariant/InvariantBridge.t.sol
@@ -232,6 +232,27 @@ contract InvariantBridge is Test {
         vm.selectFork(resetId);
     }
 
+    /// forge-config: default.invariant.runs = 256	
+    /// forge-config: default.invariant.depth = 15	
+    /// forge-config: default.invariant.fail-on-revert = true	
+    function invariant_IndividualERC20TokenBalanced() external {	
+        for (uint256 i = 0; i < NO_OF_TOKENS; i++) {	
+            address rootToken = rootTokens[i];	
+            for (uint256 j = 0; j < NO_OF_USERS; j++) {	
+                address user = users[j];	
+
+                vm.selectFork(rootId);	
+                uint256 balanceL1 = ChildERC20(rootToken).balanceOf(user);	
+                address childToken = rootBridge.rootTokenToChildToken(rootToken);	
+
+                vm.selectFork(childId);	
+                uint256 balanceL2 = ChildERC20(childToken).balanceOf(user);	
+
+                assertEq(balanceL1 + balanceL2, MAX_AMOUNT);	
+            }	
+        }	
+    }
+
     /// forge-config: default.invariant.runs = 256
     /// forge-config: default.invariant.depth = 15
     /// forge-config: default.invariant.fail-on-revert = true

diff --git a/test/invariant/child/ChildERC20BridgeHandler.sol b/test/invariant/child/ChildERC20BridgeHandler.sol
@@ -69,7 +69,9 @@ contract ChildERC20BridgeHandler is Test {
 
         if (currentBalance < amount) {
             // Fund difference
-            fund(userIndexSeed, rootToken, childToken, amount - currentBalance);
+            vm.selectFork(rootId);
+            rootHelper.deposit(user, rootToken, amount - currentBalance, gasAmt);
+            vm.selectFork(childId);
         }
 
         vm.selectFork(rootId);
@@ -112,7 +114,9 @@ contract ChildERC20BridgeHandler is Test {
 
         if (currentBalance < amount) {
             // Fund difference
-            fund(userIndexSeed, rootToken, childToken, amount - currentBalance);
+            vm.selectFork(rootId);
+            rootHelper.deposit(user, rootToken, amount - currentBalance, gasAmt);
+            vm.selectFork(childId);
         }
 
         vm.selectFork(rootId);
@@ -123,36 +127,13 @@ contract ChildERC20BridgeHandler is Test {
 
         vm.selectFork(rootId);
         rootHelper.finaliseWithdrawal(recipient, previousLen);
+        // If recipient is different, transfer back
+        if (user != recipient) {
+            vm.prank(recipient);
+            ChildERC20(rootToken).transfer(user, amount);
+        }
         vm.selectFork(childId);
 
         vm.selectFork(original);
     }
-
-    function fund(uint256 userIndexSeed, address rootToken, address childToken, uint256 diff) public {
-        uint256 offset = bound(userIndexSeed, 0, users.length - 1);
-        address user = users[offset];
-        address from = findFrom(offset, childToken, diff);
-        if (from != address(0)) {
-            vm.prank(from);
-            ChildERC20(childToken).transfer(user, diff);
-        } else {
-            vm.selectFork(rootId);
-            from = findFrom(offset, rootToken, diff);
-            rootHelper.depositTo(from, user, rootToken, diff, 1);
-            vm.selectFork(childId);
-        }
-    }
-
-    function findFrom(uint256 offset, address token, uint256 requiredAmt) public view returns (address from) {
-        for (uint256 i = 0; i < users.length; i++) {
-            uint256 index = i + offset;
-            if (index >= users.length) {
-                index -= users.length;
-            }
-            if (ChildERC20(token).balanceOf(users[index]) >= requiredAmt) {
-                from = users[index];
-                break;
-            }
-        }
-    }
 }
diff --git a/test/invariant/root/RootERC20BridgeFlowRateHandler.sol b/test/invariant/root/RootERC20BridgeFlowRateHandler.sol
@@ -53,7 +53,11 @@ contract RootERC20BridgeFlowRateHandler is Test {
 
         if (currentBalance < amount) {
             // Fund difference
-            fund(userIndexSeed, rootToken, childToken, amount - currentBalance);
+            uint256 previousLen = rootHelper.getQueueSize(user);
+            vm.selectFork(childId);
+            childHelper.withdraw(user, childToken, amount - currentBalance, gasAmt);
+            vm.selectFork(rootId);
+            rootHelper.finaliseWithdrawal(user, previousLen);
         }
 
         rootHelper.deposit(user, rootToken, amount, gasAmt);
@@ -88,41 +92,23 @@ contract RootERC20BridgeFlowRateHandler is Test {
 
         if (currentBalance < amount) {
             // Fund difference
-            fund(userIndexSeed, rootToken, childToken, amount - currentBalance);
+            uint256 previousLen = rootHelper.getQueueSize(user);
+            vm.selectFork(childId);
+            childHelper.withdraw(user, childToken, amount - currentBalance, gasAmt);
+            vm.selectFork(rootId);
+            rootHelper.finaliseWithdrawal(user, previousLen);
         }
 
         rootHelper.depositTo(user, recipient, rootToken, amount, gasAmt);
 
-        vm.selectFork(original);
-    }
-
-    function fund(uint256 userIndexSeed, address rootToken, address childToken, uint256 diff) public {
-        uint256 offset = bound(userIndexSeed, 0, users.length - 1);
-        address user = users[offset];
-        address from = findFrom(offset, rootToken, diff);
-        if (from != address(0)) {
-            vm.prank(from);
-            ChildERC20(rootToken).transfer(user, diff);
-        } else {
-            uint256 previousLen = rootHelper.getQueueSize(user);
+        // If recipient is different, transfer back
+        if (user != recipient) {
             vm.selectFork(childId);
-            from = findFrom(offset, childToken, diff);
-            childHelper.withdrawTo(from, user, childToken, diff, 1);
+            vm.prank(recipient);
+            ChildERC20(childToken).transfer(user, amount);
             vm.selectFork(rootId);
-            rootHelper.finaliseWithdrawal(user, previousLen);
         }
-    }
 
-    function findFrom(uint256 offset, address token, uint256 requiredAmt) public view returns (address from) {
-        for (uint256 i = 0; i < users.length; i++) {
-            uint256 index = i + offset;
-            if (index >= users.length) {
-                index -= users.length;
-            }
-            if (ChildERC20(token).balanceOf(users[index]) >= requiredAmt) {
-                from = users[index];
-                break;
-            }
-        }
+        vm.selectFork(original);
     }
 }