From 9c6b78c59f7f5958c1227bbc7febbaa352acc314 Mon Sep 17 00:00:00 2001
From: wcgcyx <wcgcyx@gmail.com>
Date: Tue, 23 Jul 2024 10:39:43 +1000
Subject: [PATCH] Fix CI

---
 test/invariant/IChainManager.sol     |  6 ++++++
 test/invariant/InvariantBridge.t.sol | 11 ++++++++---
 test/invariant/MockAdaptor.sol       |  9 ++++++---
 3 files changed, 20 insertions(+), 6 deletions(-)
 create mode 100644 test/invariant/IChainManager.sol

diff --git a/test/invariant/IChainManager.sol b/test/invariant/IChainManager.sol
new file mode 100644
index 00000000..c49d16e1
--- /dev/null
+++ b/test/invariant/IChainManager.sol
@@ -0,0 +1,6 @@
+// SPDX-License-Identifier: Apache 2.0
+pragma solidity 0.8.19;
+
+interface IChainManager {
+    function switchToChain(uint256 chainId) external;
+}
\ No newline at end of file
diff --git a/test/invariant/InvariantBridge.t.sol b/test/invariant/InvariantBridge.t.sol
index 12ec8695..d0d997ec 100644
--- a/test/invariant/InvariantBridge.t.sol
+++ b/test/invariant/InvariantBridge.t.sol
@@ -11,11 +11,12 @@ import {RootERC20BridgeFlowRate} from "../../src/root/flowrate/RootERC20BridgeFl
 import {MockAdaptor} from "./MockAdaptor.sol";
 import {ChildHelper} from "./child/ChildHelper.sol";
 import {RootHelper} from "./root/RootHelper.sol";
+import {IChainManager} from "./IChainManager.sol";
 import {ChildERC20BridgeHandler} from "./child/ChildERC20BridgeHandler.sol";
 import {RootERC20BridgeFlowRateHandler} from "./root/RootERC20BridgeFlowRateHandler.sol";
 import "forge-std/console.sol";
 
-contract InvariantBridge is Test {
+contract InvariantBridge is Test, IChainManager {
     string public constant CHAIN_URL = "http://127.0.0.1:8500";
     uint256 public constant IMX_DEPOSIT_LIMIT = 10000 ether;
     uint256 public constant MAX_AMOUNT = 10000;
@@ -40,6 +41,10 @@ contract InvariantBridge is Test {
 
     uint256 mappingGas;
 
+    function switchToChain(uint256 chainId) external {
+        vm.selectFork(chainId);
+    }
+
     function setUp() public {
         childId = vm.createFork(CHAIN_URL);
         rootId = vm.createFork(CHAIN_URL);
@@ -77,7 +82,7 @@ contract InvariantBridge is Test {
 
         // Configure contracts on child chain.
         vm.selectFork(childId);
-        childAdaptor.initialize(rootId, address(childBridge));
+        childAdaptor.initialize(rootId, address(childBridge), address(this));
         IChildERC20Bridge.InitializationRoles memory childRoles = IChildERC20Bridge.InitializationRoles({
             defaultAdmin: address(this),
             pauser: address(this),
@@ -93,7 +98,7 @@ contract InvariantBridge is Test {
 
         // Configure contracts on root chain.
         vm.selectFork(rootId);
-        rootAdaptor.initialize(childId, address(rootBridge));
+        rootAdaptor.initialize(childId, address(rootBridge), address(this));
         IRootERC20Bridge.InitializationRoles memory rootRoles = IRootERC20Bridge.InitializationRoles({
             defaultAdmin: address(this),
             pauser: address(this),
diff --git a/test/invariant/MockAdaptor.sol b/test/invariant/MockAdaptor.sol
index 0a5bfe83..ae4ab769 100644
--- a/test/invariant/MockAdaptor.sol
+++ b/test/invariant/MockAdaptor.sol
@@ -4,6 +4,7 @@ pragma solidity 0.8.19;
 import {Test} from "forge-std/Test.sol";
 import {IChildBridgeAdaptor} from "../../src/interfaces/child/IChildBridgeAdaptor.sol";
 import {IRootBridgeAdaptor} from "../../src/interfaces/root/IRootBridgeAdaptor.sol";
+import {IChainManager} from "./IChainManager.sol";
 
 interface MessageReceiver {
     function onMessageReceive(bytes calldata data) external;
@@ -12,12 +13,14 @@ interface MessageReceiver {
 contract MockAdaptor is Test, IChildBridgeAdaptor, IRootBridgeAdaptor {
     uint256 otherChainId;
     MessageReceiver messageReceiver;
+    IChainManager chainManager;
 
     constructor() {}
 
-    function initialize(uint256 _otherChainId, address _messageReceiver) public {
+    function initialize(uint256 _otherChainId, address _messageReceiver, address _chainManager) public {
         otherChainId = _otherChainId;
         messageReceiver = MessageReceiver(_messageReceiver);
+        chainManager = IChainManager(_chainManager);
     }
 
     function sendMessage(bytes calldata payload, address /*refundRecipient*/ )
@@ -28,10 +31,10 @@ contract MockAdaptor is Test, IChildBridgeAdaptor, IRootBridgeAdaptor {
         uint256 original = vm.activeFork();
 
         // Switch to the other chain.
-        vm.selectFork(otherChainId);
+        chainManager.switchToChain(otherChainId);
         onMessageReceive(payload);
 
-        vm.selectFork(original);
+        chainManager.switchToChain(original);
     }
 
     function onMessageReceive(bytes calldata data) public {