Pre-deploy extra StakedTokenHandler contracts

to use with StakedTokens deployed during fuzzing so the handlers get access to cheat codes see [this foundry issue](foundry-rs/foundry#6652)
Increment-Finance · Jan 2, 2024 · 9db8789 · 9db8789
1 parent 1ca5903
commit 9db8789
Show file tree

Hide file tree

Showing 2 changed files with 80 additions and 4 deletions.
diff --git a/test/invariant/SafetyModuleInvariantTest.sol b/test/invariant/SafetyModuleInvariantTest.sol
@@ -78,7 +78,15 @@ contract SafetyModuleInvariantTest is Test {
     SMRDHandler public smrdHandler;
     StakedTokenHandler public stakedTokenHandler1;
     StakedBPTHandler public stakedTokenHandler2;
+    StakedTokenHandler public stakedTokenHandler3;
+    StakedTokenHandler public stakedTokenHandler4;
+    StakedTokenHandler public stakedTokenHandler5;
+    StakedTokenHandler public stakedTokenHandler6;
+    StakedTokenHandler public stakedTokenHandler7;
+    StakedTokenHandler public stakedTokenHandler8;
+    StakedTokenHandler public stakedTokenHandler9;
     StakedTokenHandler[] public stakedTokenHandlers;
+    uint256 public numStakedTokenHandlers = 2;
 
     // Invariant ghost variables
     mapping(address => mapping(address => uint256))
@@ -275,6 +283,35 @@ contract SafetyModuleInvariantTest is Test {
             poolAssets,
             poolId
         );
+        // Deploy extra staked token handlers for later use
+        stakedTokenHandler3 = new StakedTokenHandler(
+            StakedToken(address(0)),
+            stakers
+        );
+        stakedTokenHandler4 = new StakedTokenHandler(
+            StakedToken(address(0)),
+            stakers
+        );
+        stakedTokenHandler5 = new StakedTokenHandler(
+            StakedToken(address(0)),
+            stakers
+        );
+        stakedTokenHandler6 = new StakedTokenHandler(
+            StakedToken(address(0)),
+            stakers
+        );
+        stakedTokenHandler7 = new StakedTokenHandler(
+            StakedToken(address(0)),
+            stakers
+        );
+        stakedTokenHandler8 = new StakedTokenHandler(
+            StakedToken(address(0)),
+            stakers
+        );
+        stakedTokenHandler9 = new StakedTokenHandler(
+            StakedToken(address(0)),
+            stakers
+        );
         stakedTokenHandlers = [stakedTokenHandler1, stakedTokenHandler2];
 
         // Set handlers as target contracts
@@ -447,12 +484,46 @@ contract SafetyModuleInvariantTest is Test {
     function addStakedToken(
         StakedToken newStakedToken
     ) external onlySafetyModuleHandler {
+        IERC20 underlying = newStakedToken.getUnderlyingToken();
+
+        // Use pre-deployed StakedTokenHandler
+        // since deploying a new one doesn't give it access to cheats
+        StakedTokenHandler newStakedTokenHandler;
+        if (numStakedTokenHandlers == 2)
+            newStakedTokenHandler = stakedTokenHandler3;
+        else if (numStakedTokenHandlers == 3)
+            newStakedTokenHandler = stakedTokenHandler4;
+        else if (numStakedTokenHandlers == 4)
+            newStakedTokenHandler = stakedTokenHandler5;
+        else if (numStakedTokenHandlers == 5)
+            newStakedTokenHandler = stakedTokenHandler6;
+        else if (numStakedTokenHandlers == 6)
+            newStakedTokenHandler = stakedTokenHandler7;
+        else if (numStakedTokenHandlers == 7)
+            newStakedTokenHandler = stakedTokenHandler8;
+        else if (numStakedTokenHandlers == 8)
+            newStakedTokenHandler = stakedTokenHandler9;
+        else revert("too many staked token handlers");
+        numStakedTokenHandlers++;
+        newStakedTokenHandler.setStakedToken(newStakedToken);
+        // newStakedTokenHandler = new StakedTokenHandler(newStakedToken, stakers);
+
+        // Add staked token and its handler to lists
         stakedTokens.push(newStakedToken);
-        StakedTokenHandler newStakedTokenHandler = new StakedTokenHandler(
-            newStakedToken,
-            stakers
-        );
         stakedTokenHandlers.push(newStakedTokenHandler);
+
+        // Register new target and exclude base contracts
         targetContract(address(newStakedTokenHandler));
+        excludeContract(address(newStakedToken));
+        excludeContract(address(underlying));
+
+        // Approve new staked token for users
+        vm.startPrank(stakerOne);
+        underlying.approve(address(newStakedToken), type(uint256).max);
+        vm.startPrank(stakerTwo);
+        underlying.approve(address(newStakedToken), type(uint256).max);
+        vm.startPrank(stakerThree);
+        underlying.approve(address(newStakedToken), type(uint256).max);
+        vm.stopPrank();
     }
 }
diff --git a/test/invariant/handlers/SafetyModuleHandler.sol b/test/invariant/handlers/SafetyModuleHandler.sol
@@ -15,6 +15,8 @@ import {strings} from "stringutils/strings.sol";
 import {PRBMathUD60x18} from "prb-math/contracts/PRBMathUD60x18.sol";
 
 interface ITestContract {
+    function numStakedTokenHandlers() external view returns (uint256);
+
     function addStakedToken(StakedToken stakedToken) external;
 }
 
@@ -93,6 +95,9 @@ contract SafetyModuleHandler is Test {
     //     uint256 unstakeWindowSeconds,
     //     uint256 maxStakeAmount
     // ) external useGovernance {
+    //     if (testContract.numStakedTokenHandlers() == 9) {
+    //         return;
+    //     }
     //     cooldownSeconds = bound(cooldownSeconds, 1 hours, 1 weeks);
     //     unstakeWindowSeconds = bound(unstakeWindowSeconds, 1 hours, 1 weeks);
     //     maxStakeAmount = bound(maxStakeAmount, 10_000e18, 1_000_000e18);