chore: test invariant

liquid-collective · Nov 6, 2023 · 0a30f9f · 0a30f9f
1 parent 6a83f37
commit 0a30f9f
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 4 deletions.
diff --git a/contracts/test/invariants/Base.sol b/contracts/test/invariants/Base.sol
@@ -44,6 +44,7 @@ contract Base is Test, BytesGenerator {
     CoverageFundV1 public coverageFund;
     AllowlistV1 public allowlist;
     OperatorsRegistryV1 public operatorsRegistry;
+    RedeemManagerV1 public redeemManager;
 
     address internal admin;
     address internal newAdmin;
@@ -84,6 +85,7 @@ contract Base is Test, BytesGenerator {
         deployProtocol();
         deployServices();
         addTargetSelectors();
+        excludeDeployedContracts();
     }
 
     function loadBlockState() public {
@@ -144,12 +146,15 @@ contract Base is Test, BytesGenerator {
         LibImplementationUnbricker.unbrick(vm, address(river));
         operatorsRegistry = new OperatorsRegistryV1();
         LibImplementationUnbricker.unbrick(vm, address(operatorsRegistry));
+        redeemManager = new RedeemManagerV1();
+        LibImplementationUnbricker.unbrick(vm, address(redeemManager));
 
         bytes32 withdrawalCredentials = withdraw.getCredentials();
         allowlist.initAllowlistV1(admin, allower);
         operatorsRegistry.initOperatorsRegistryV1(admin, address(river));
         elFeeRecipient.initELFeeRecipientV1(address(river));
         coverageFund.initCoverageFundV1(address(river));
+        redeemManager.initializeRedeemManagerV1(address(river));
 
         river.initRiverV1(
             address(deposit),
@@ -195,14 +200,32 @@ contract Base is Test, BytesGenerator {
 
     function deployServices() internal {
         stakerService = new StakerService(this);
+        address[] memory stakerServiceArray = new address[](1);
+        stakerServiceArray[0] = address(stakerService);
+        uint256[] memory stakerServiceMask = new uint256[](1);
+        stakerServiceMask[0] = 5;
+        vm.prank(allower);
+        allowlist.allow(stakerServiceArray, stakerServiceMask);
     }
 
-    function dealETH(address _to, uint256 _amount) external {
+    function dealETH(address _to, uint256 _amount) public {
+        console.log("dealing");
         vm.deal(_to, _amount);
     }
 
     function addTargetSelectors() internal virtual {
-        StdInvariant.FuzzSelector memory selectors = stakerService.getTargetSelectors();
-        targetSelector(selectors);
+        targetSelector(stakerService.getTargetSelectors());
+    }
+
+    function excludeDeployedContracts() internal virtual {
+        excludeContract(address(river));
+        excludeContract(address(deposit));
+        excludeContract(address(withdraw));
+        excludeContract(address(oracle));
+        excludeContract(address(elFeeRecipient));
+        excludeContract(address(coverageFund));
+        excludeContract(address(allowlist));
+        excludeContract(address(operatorsRegistry));
+        excludeContract(address(redeemManager));
     }
 }
diff --git a/contracts/test/invariants/Invariant_river.sol b/contracts/test/invariants/Invariant_river.sol
@@ -1,12 +1,23 @@
 pragma solidity 0.8.10;
 
 import "forge-std/Test.sol";
+import "forge-std/console.sol";
 import {Base} from "./Base.sol";
 
 contract Invariant_River is Base {
-    Base private base;
+    function setUp() public override {
+        super.setUp();
+    }
 
     function invariant_test() public {
         assert(1 == 1);
     }
+
+    function invariant_setBalanceToDeposit() public {
+        uint256 balanceToDepositBefore = river.getBalanceToDeposit();
+        stakerService.action_stakeAmount(10 ether);
+        uint256 balanceToDepositAfter = river.getBalanceToDeposit();
+
+        assertEq(balanceToDepositAfter, balanceToDepositBefore + 10 ether);
+    }
 }
diff --git a/contracts/test/invariants/handlers/StakerService.sol b/contracts/test/invariants/handlers/StakerService.sol
@@ -4,6 +4,7 @@ pragma solidity 0.8.10;
 
 import "forge-std/Test.sol";
 import "forge-std/console.sol";
+
 import {BaseService} from "./BaseService.sol";
 import {Base} from "../Base.sol";
 
@@ -21,14 +22,24 @@ contract StakerService is BaseService {
     }
 
     function action_stakeAll() public recordBlockData {
+        base.dealETH(address(this), 1 ether);
         console.log("Staking all funds");
+        base.river().deposit{value: 1 ether}();
     }
 
     function action_stakePercent(uint256 pct) public recordBlockData {
+        base.dealETH(address(this), 1 ether);
         console.log("Staking percentage of funds");
+        base.river().deposit{value: 1 ether}();
     }
 
     function action_unstakeAll() public recordBlockData {
         console.log("Unstaking all funds");
     }
+
+    function action_stakeAmount(uint256 amount) public recordBlockData {
+        console.log("Staking amount of funds");
+        base.dealETH(address(this), amount);
+        base.river().deposit{value: amount}();
+    }
 }