feat: add tests for withdraw, fee and make it DRY

foundry/script/Counter.s.sol

@@ -1,19 +0,0 @@
-// SPDX-License-Identifier: UNLICENSED
-pragma solidity ^0.8.13;
-
-import {Script, console} from "forge-std/Script.sol";
-import {Counter} from "../src/Counter.sol";
-
-contract CounterScript is Script {
-    Counter public counter;
-
-    function setUp() public {}
-
-    function run() public {
-        vm.startBroadcast();
-
-        counter = new Counter();
-
-        vm.stopBroadcast();
-    }
-}

foundry/script/Deploy.sol

@@ -0,0 +1,8 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.13;
+
+import {Script, console} from "forge-std/Script.sol";
+
+contract DeployScript is Script {
+// TODO: implement deploy script
+}

foundry/src/Counter.sol

@@ -1,14 +0,0 @@
-// SPDX-License-Identifier: UNLICENSED
-pragma solidity ^0.8.28;
-
-contract Counter {
-    uint256 public number;
-
-    function setNumber(uint256 newNumber) public {
-        number = newNumber;
-    }
-
-    function increment() public {
-        number++;
-    }
-}

foundry/test/Constants.sol

@@ -1,10 +1,13 @@
 // SPDX-License-Identifier: UNLICENSED
-pragma solidity ^0.8.13;
+pragma solidity ^0.8.28;
 
-contract Constants {
+import "forge-std/Test.sol";
-    address ADMIN = address(12395012351212343412541234); //admin=us
-    address BOB = address(123); //bob=someone!=us
 
+contract Constants is Test {
+    address ADMIN = makeAddr("admin"); //admin=us
+    address BOB = makeAddr("bob"); //bob=someone!=us
+    address FUND_RESCUER = makeAddr("fundRescuer");
+    address FEE_SETTER = makeAddr("feeSetter");
     // dummy contracts
-    address DUMMY = address(0x1234);
+    address DUMMY = makeAddr("dummy");
 }

foundry/test/Counter.t.sol

@@ -1,24 +0,0 @@
-// SPDX-License-Identifier: UNLICENSED
-pragma solidity ^0.8.28;
-
-import {Test, console} from "forge-std/Test.sol";
-import {Counter} from "../src/Counter.sol";
-
-contract CounterTest is Test {
-    Counter public counter;
-
-    function setUp() public {
-        counter = new Counter();
-        counter.setNumber(0);
-    }
-
-    function test_Increment() public {
-        counter.increment();
-        assertEq(counter.number(), 1);
-    }
-
-    function testFuzz_SetNumber(uint256 x) public {
-        counter.setNumber(x);
-        assertEq(counter.number(), x);
-    }
-}

foundry/test/TychoRouter.t.sol

@@ -2,9 +2,9 @@
 pragma solidity ^0.8.28;
 
 import {TychoRouter} from "@src/TychoRouter.sol";
-import "./TestTemplate.sol";
+import "./TychoRouterTestSetup.sol";
 
-contract TychoRouterTest is TychoRouterTestTemplate {
+contract TychoRouterTest is TychoRouterTestSetup {
     bytes32 public constant EXECUTOR_SETTER_ROLE =
         0x6a1dd52dcad5bd732e45b6af4e7344fa284e2d7d4b23b5b09cb55d36b0685c87;
     bytes32 public constant FEE_SETTER_ROLE =
@@ -16,15 +16,11 @@ contract TychoRouterTest is TychoRouterTestTemplate {
 
     event ExecutorSet(address indexed executor);
     event CallbackVerifierSet(address indexed callbackVerifier);
-
+    event Withdrawal(
-    function setupTychoRouter() public {
+        address indexed token, uint256 amount, address indexed receiver
-        deployTychoRouter();
+    );
-    }
 
     function testSetValidExecutor() public {
-        setupTychoRouter();
-        deployDummyContract();
-
         vm.startPrank(executorSetter);
         vm.expectEmit();
         // Define the event we expect to be emitted at the next step
@@ -37,9 +33,6 @@ contract TychoRouterTest is TychoRouterTestTemplate {
     }
 
     function testRemoveExecutor() public {
-        setupTychoRouter();
-        deployDummyContract();
-
         vm.startPrank(executorSetter);
         tychoRouter.setSwapExecutor(DUMMY);
         tychoRouter.removeSwapExecutor(DUMMY);
@@ -48,9 +41,6 @@ contract TychoRouterTest is TychoRouterTestTemplate {
     }
 
     function testRemoveUnSetExecutor() public {
-        setupTychoRouter();
-        deployDummyContract();
-
         vm.startPrank(executorSetter);
         tychoRouter.removeSwapExecutor(BOB);
         vm.stopPrank();
@@ -58,24 +48,16 @@ contract TychoRouterTest is TychoRouterTestTemplate {
     }
 
     function testRemoveExecutorMissingSetterRole() public {
-        setupTychoRouter();
-        deployDummyContract();
         vm.expectRevert();
         tychoRouter.removeSwapExecutor(BOB);
     }
 
     function testSetExecutorMissingSetterRole() public {
-        setupTychoRouter();
-        deployDummyContract();
-
         vm.expectRevert();
         tychoRouter.setSwapExecutor(DUMMY);
     }
 
     function testSetExecutorNonContract() public {
-        setupTychoRouter();
-        deployDummyContract();
-
         vm.startPrank(executorSetter);
         vm.expectRevert(
             abi.encodeWithSelector(TychoRouter__NonContractExecutor.selector)
@@ -85,9 +67,6 @@ contract TychoRouterTest is TychoRouterTestTemplate {
     }
 
     function testSetValidVerifier() public {
-        setupTychoRouter();
-        deployDummyContract();
-
         vm.startPrank(executorSetter);
         vm.expectEmit();
         // Define the event we expect to be emitted at the next step
@@ -100,9 +79,6 @@ contract TychoRouterTest is TychoRouterTestTemplate {
     }
 
     function testRemoveVerifier() public {
-        setupTychoRouter();
-        deployDummyContract();
-
         vm.startPrank(executorSetter);
         tychoRouter.setCallbackVerifier(DUMMY);
         tychoRouter.removeCallbackVerifier(DUMMY);
@@ -111,9 +87,6 @@ contract TychoRouterTest is TychoRouterTestTemplate {
     }
 
     function testRemoveUnSetVerifier() public {
-        setupTychoRouter();
-        deployDummyContract();
-
         vm.startPrank(executorSetter);
         tychoRouter.removeCallbackVerifier(BOB);
         vm.stopPrank();
@@ -121,24 +94,16 @@ contract TychoRouterTest is TychoRouterTestTemplate {
     }
 
     function testRemoveVerifierMissingSetterRole() public {
-        setupTychoRouter();
-        deployDummyContract();
         vm.expectRevert();
         tychoRouter.removeCallbackVerifier(BOB);
     }
 
     function testSetVerifierMissingSetterRole() public {
-        setupTychoRouter();
-        deployDummyContract();
-
         vm.expectRevert();
         tychoRouter.setCallbackVerifier(DUMMY);
     }
 
     function testSetVerifierNonContract() public {
-        setupTychoRouter();
-        deployDummyContract();
-
         vm.startPrank(executorSetter);
         vm.expectRevert(
             abi.encodeWithSelector(TychoRouter__NonContractVerifier.selector)
@@ -146,4 +111,91 @@ contract TychoRouterTest is TychoRouterTestTemplate {
         tychoRouter.setCallbackVerifier(BOB);
         vm.stopPrank();
     }
+
+    function testWithdrawNative() public {
+        vm.startPrank(FUND_RESCUER);
+        // Send 100 ether to tychoRouter
+        assertEq(address(tychoRouter).balance, 0);
+        assertEq(FUND_RESCUER.balance, 0);
+        vm.deal(address(tychoRouter), 100 ether);
+        vm.expectEmit();
+        emit Withdrawal(address(0), 100 ether, FUND_RESCUER);
+        tychoRouter.withdrawNative(FUND_RESCUER);
+        assertEq(address(tychoRouter).balance, 0);
+        assertEq(FUND_RESCUER.balance, 100 ether);
+        vm.stopPrank();
+    }
+
+    function testWithdrawNativeOtherCases() public {
+        vm.deal(address(tychoRouter), 100 ether);
+        vm.startPrank(FUND_RESCUER);
+        vm.expectRevert(TychoRouter__AddressZero.selector);
+        tychoRouter.withdrawNative(address(0));
+        vm.stopPrank();
+
+        // Not role FUND_RESCUER
+        vm.startPrank(BOB);
+        vm.expectRevert();
+        tychoRouter.withdrawNative(FUND_RESCUER);
+        vm.stopPrank();
+    }
+
+    function testWithdrawERC20Tokens() public {
+        vm.startPrank(BOB);
+        // Check balances before minting
+        for (uint256 i = 0; i < tokens.length; i++) {
+            // slither-disable-next-line calls-loop
+            assertEq(tokens[i].balanceOf(address(tychoRouter)), 0);
+        }
+        // Mint tokens to tychoRouter
+        for (uint256 i = 0; i < tokens.length; i++) {
+            // slither-disable-next-line calls-loop
+            tokens[i].mint(address(tychoRouter), 100 ether);
+        }
+        // Check balances after minting
+        for (uint256 i = 0; i < tokens.length; i++) {
+            // slither-disable-next-line calls-loop
+            assertEq(tokens[i].balanceOf(address(tychoRouter)), 100 ether);
+        }
+        vm.stopPrank();
+
+        vm.startPrank(FUND_RESCUER);
+        IERC20[] memory tokensArray = new IERC20[](3);
+        tokensArray[0] = IERC20(address(tokens[0]));
+        tokensArray[1] = IERC20(address(tokens[1]));
+        tokensArray[2] = IERC20(address(tokens[2]));
+        tychoRouter.withdraw(tokensArray, FUND_RESCUER);
+
+        // Check balances after withdrawing
+        for (uint256 i = 0; i < tokens.length; i++) {
+            // slither-disable-next-line calls-loop
+            assertEq(tokens[i].balanceOf(address(tychoRouter)), 0);
+            // slither-disable-next-line calls-loop
+            assertEq(tokens[i].balanceOf(FUND_RESCUER), 100 ether);
+        }
+        vm.stopPrank();
+
+        vm.startPrank(FUND_RESCUER);
+        vm.expectRevert();
+        tychoRouter.withdraw(tokensArray, address(0));
+        vm.stopPrank();
+
+        vm.startPrank(BOB);
+        vm.expectRevert();
+        tychoRouter.withdraw(tokensArray, FUND_RESCUER);
+        vm.stopPrank();
+    }
+
+    function testFeeSetterRole() public {
+        vm.startPrank(FEE_SETTER);
+        assertEq(tychoRouter.fee(), 0);
+        tychoRouter.setFee(100);
+        assertEq(tychoRouter.fee(), 100);
+        vm.stopPrank();
+
+        vm.startPrank(BOB);
+        vm.expectRevert();
+        tychoRouter.setFee(200);
+        vm.stopPrank();
+    }
 }

foundry/test/TychoRouterTestSetup.sol

@@ -1,25 +1,30 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity ^0.8.13;
 
-import "forge-std/Test.sol";
 import "@src/TychoRouter.sol";
 import "./Constants.sol";
+import "./mock/MockERC20.sol";
 
-contract TychoRouterTestTemplate is Test, Constants {
+contract TychoRouterTestSetup is Test, Constants {
     TychoRouter tychoRouter;
-    address tychoRouterAddress;
     address executorSetter;
+    address permit2Address = address(0x000000000022D473030F116dDEE9F6B43aC78BA3);
+    MockERC20[] tokens;
 
-    function deployTychoRouter() internal {
+    function setUp() public {
         vm.startPrank(ADMIN);
-
-        address permit2Address =
-            address(0x000000000022D473030F116dDEE9F6B43aC78BA3);
         tychoRouter = new TychoRouter(permit2Address);
-        tychoRouterAddress = address(tychoRouter);
         tychoRouter.grantRole(keccak256("EXECUTOR_SETTER_ROLE"), BOB);
+        tychoRouter.grantRole(keccak256("FUND_RESCUER_ROLE"), FUND_RESCUER);
+        tychoRouter.grantRole(keccak256("FEE_SETTER_ROLE"), FEE_SETTER);
         executorSetter = BOB;
+        deployDummyContract();
+        vm.stopPrank();
 
+        vm.startPrank(BOB);
+        tokens.push(new MockERC20("Token A", "A"));
+        tokens.push(new MockERC20("Token B", "B"));
+        tokens.push(new MockERC20("Token C", "C"));
         vm.stopPrank();
     }
 

foundry/test/mock/MockERC20.sol

@@ -0,0 +1,18 @@
+// SPDX-License-Identifier: Unlicense
+pragma solidity ^0.8.28;
+
+import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
+
+contract MockERC20 is ERC20 {
+    constructor(string memory name_, string memory symbol_)
+        ERC20(name_, symbol_)
+    {}
+
+    function mint(address to, uint256 amount) external {
+        _mint(to, amount);
+    }
+
+    function decimals() public view virtual override returns (uint8) {
+        return 18;
+    }
+}