Update InvariantBridge.t.sol

immutable · Feb 20, 2024 · 0375afd · 0375afd
1 parent ced0691
commit 0375afd
Showing 1 changed file with 15 additions and 9 deletions.
diff --git a/test/invariant/InvariantBridge.t.sol b/test/invariant/InvariantBridge.t.sol
@@ -34,6 +34,7 @@ contract InvariantBridge is Test {
     address[] users;
     address[] rootTokens;
 
+    ChildERC20BridgeHandler childBridgeHandler;
     RootERC20BridgeFlowRateHandler rootBridgeHandler;
 
     function setUp() public {
@@ -46,9 +47,8 @@ contract InvariantBridge is Test {
         ChildERC20 childTokenTemplate = new ChildERC20();
         childTokenTemplate.initialize(address(123), "Test", "TST", 18);
         childAdaptor = new MockAdaptor();
-        childBridge = new ChildERC20Bridge(address(this));
-        new RootERC20BridgeFlowRate(address(this));
         vm.stopPrank();
+        childBridge = new ChildERC20Bridge(address(this));
         WIMX wIMX = new WIMX();
 
         // Deploy contracts on root chain.
@@ -57,9 +57,8 @@ contract InvariantBridge is Test {
         ChildERC20 rootTokenTemplate = new ChildERC20();
         rootTokenTemplate.initialize(address(123), "Test", "TST", 18);
         rootAdaptor = new MockAdaptor();
-        new ChildERC20Bridge(address(this));
-        rootBridge = new RootERC20BridgeFlowRate(address(this));
         vm.stopPrank();
+        rootBridge = new RootERC20BridgeFlowRate(address(this));
         ChildERC20 rootIMXToken = new ChildERC20();
         rootIMXToken.initialize(address(123), "Immutable X", "IMX", 18);
         WIMX wETH = new WIMX();
@@ -68,12 +67,12 @@ contract InvariantBridge is Test {
         vm.selectFork(childId);
         vm.startPrank(ADMIN);
         childHelper = new ChildHelper(childId, payable(childBridge));
-        address(new RootHelper(rootId, ADMIN, payable(rootBridge)));
+        new RootHelper(rootId, ADMIN, payable(rootBridge));
         vm.stopPrank();
 
         vm.selectFork(rootId);
         vm.startPrank(ADMIN);
-        address(new ChildHelper(childId, payable(childBridge)));
+        new ChildHelper(childId, payable(childBridge));
         rootHelper = new RootHelper(rootId, ADMIN, payable(rootBridge));
         vm.stopPrank();
 
@@ -149,14 +148,20 @@ contract InvariantBridge is Test {
             rootTokens.push(address(rootToken));
         }
 
-        // Deploy handlers
+        // Deploy handlers on both chains.
         vm.selectFork(childId);
-        ChildERC20BridgeHandler childBridgeHandler =
+        vm.startPrank(ADMIN);
+        childBridgeHandler =
             new ChildERC20BridgeHandler(childId, users, rootTokens, address(childHelper), address(rootHelper));
+        new RootERC20BridgeFlowRateHandler(rootId, users, rootTokens, address(childHelper), address(rootHelper));
+        vm.stopPrank();
 
         vm.selectFork(rootId);
+        vm.startPrank(ADMIN);
+        new ChildERC20BridgeHandler(childId, users, rootTokens, address(childHelper), address(rootHelper));
         rootBridgeHandler =
             new RootERC20BridgeFlowRateHandler(rootId, users, rootTokens, address(childHelper), address(rootHelper));
+        vm.stopPrank();
 
         // Target contracts
         bytes4[] memory childSelectors = new bytes4[](1);
@@ -171,5 +176,6 @@ contract InvariantBridge is Test {
         targetContract(address(rootBridgeHandler));
     }
 
-    function invariant_A() external {}
+    function invariant_A() external {
+    }
 }