test: Add rate setting logic (SC-483)

test/invariant/Invariants.t.sol

@@ -5,15 +5,17 @@ import "forge-std/Test.sol";
 
 import { PSMTestBase } from "test/PSMTestBase.sol";
 
-import { LpHandler }       from "test/invariant/handlers/LpHandler.sol";
-import { SwapperHandler }  from "test/invariant/handlers/SwapperHandler.sol";
-import { TransferHandler } from "test/invariant/handlers/TransferHandler.sol";
+import { LpHandler }         from "test/invariant/handlers/LpHandler.sol";
+import { RateSetterHandler } from "test/invariant/handlers/RateSetterHandler.sol";
+import { SwapperHandler }    from "test/invariant/handlers/SwapperHandler.sol";
+import { TransferHandler }   from "test/invariant/handlers/TransferHandler.sol";
 
 abstract contract PSMInvariantTestBase is PSMTestBase {
 
-    LpHandler       public lpHandler;
-    SwapperHandler  public swapperHandler;
-    TransferHandler public transferHandler;
+    LpHandler         public lpHandler;
+    RateSetterHandler public rateSetterHandler;
+    SwapperHandler    public swapperHandler;
+    TransferHandler   public transferHandler;
 
     address BURN_ADDRESS = makeAddr("burn-address");
 
@@ -45,7 +47,7 @@ abstract contract PSMInvariantTestBase is PSMTestBase {
         assertApproxEqAbs(
             psm.getPsmTotalValue(),
             psm.convertToAssetValue(psm.totalShares()),
-            2
+            3
         );
     }
 
@@ -69,12 +71,14 @@ abstract contract PSMInvariantTestBase is PSMTestBase {
         console.log("withdrawCount   ", lpHandler.withdrawCount());
         console.log("swapCount       ", swapperHandler.swapCount());
         console.log("zeroBalanceCount", swapperHandler.zeroBalanceCount());
+        console.log("setRateCount    ", rateSetterHandler.setRateCount());
         console.log(
             "sum             ",
             lpHandler.depositCount() +
             lpHandler.withdrawCount() +
             swapperHandler.swapCount() +
-            swapperHandler.zeroBalanceCount()
+            swapperHandler.zeroBalanceCount() +
+            rateSetterHandler.setRateCount()
         );
     }
 
@@ -196,8 +200,6 @@ contract PSMInvariants_ConstantRate_NoTransfer is PSMInvariantTestBase {
         lpHandler      = new LpHandler(psm, dai, usdc, sDai, 3);
         swapperHandler = new SwapperHandler(psm, dai, usdc, sDai, 3);
 
-        rateProvider.__setConversionRate(1.25e27);
-
         targetContract(address(lpHandler));
         targetContract(address(swapperHandler));
     }
@@ -229,9 +231,73 @@ contract PSMInvariants_ConstantRate_WithTransfers is PSMInvariantTestBase {
         swapperHandler  = new SwapperHandler(psm, dai, usdc, sDai, 3);
         transferHandler = new TransferHandler(psm, dai, usdc, sDai);
 
-        rateProvider.__setConversionRate(1.25e27);
+        targetContract(address(lpHandler));
+        targetContract(address(swapperHandler));
+        targetContract(address(transferHandler));
+    }
+
+    function invariant_A() public view {
+        _checkInvariant_A();
+    }
+
+    function invariant_B() public view {
+        _checkInvariant_B();
+    }
+
+    function invariant_C() public view {
+        _checkInvariant_C();
+    }
+
+    function afterInvariant() public {
+        _withdrawAllPositions();
+    }
+
+}
+
+contract PSMInvariants_RateSetting_NoTransfer is PSMInvariantTestBase {
+
+    function setUp() public override {
+        super.setUp();
+
+        lpHandler         = new LpHandler(psm, dai, usdc, sDai, 3);
+        rateSetterHandler = new RateSetterHandler(rateProvider, 1.25e27);
+        swapperHandler    = new SwapperHandler(psm, dai, usdc, sDai, 3);
+
+        targetContract(address(lpHandler));
+        targetContract(address(rateSetterHandler));
+        targetContract(address(swapperHandler));
+    }
+
+    function invariant_A() public view {
+        _checkInvariant_A();
+    }
+
+    function invariant_B() public view {
+        _checkInvariant_B();
+    }
+
+    function invariant_C() public view {
+        _checkInvariant_C();
+    }
+
+    function afterInvariant() public {
+        _withdrawAllPositions();
+    }
+
+}
+
+contract PSMInvariants_RateSetting_WithTransfers is PSMInvariantTestBase {
+
+    function setUp() public override {
+        super.setUp();
+
+        lpHandler         = new LpHandler(psm, dai, usdc, sDai, 3);
+        rateSetterHandler = new RateSetterHandler(rateProvider, 1.25e27);
+        swapperHandler    = new SwapperHandler(psm, dai, usdc, sDai, 3);
+        transferHandler   = new TransferHandler(psm, dai, usdc, sDai);
 
         targetContract(address(lpHandler));
+        targetContract(address(rateSetterHandler));
         targetContract(address(swapperHandler));
         targetContract(address(transferHandler));
     }

test/invariant/handlers/RateSetterHandler.sol

@@ -0,0 +1,30 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+pragma solidity ^0.8.13;
+
+import { StdUtils } from "forge-std/StdUtils.sol";
+
+import { MockRateProvider } from "test/mocks/MockRateProvider.sol";
+
+contract RateSetterHandler is StdUtils {
+
+    uint256 public rate;
+
+    MockRateProvider public rateProvider;
+
+    uint256 public setRateCount;
+
+    constructor(MockRateProvider rateProvider_, uint256 initialRate) {
+        rateProvider = rateProvider_;
+        rate         = initialRate;
+    }
+
+    function setRate(uint256 rateIncrease) external {
+        // Increase the rate by up to 100%
+        rate += _bound(rateIncrease, 0, 1e27);
+
+        rateProvider.__setConversionRate(rate);
+
+        setRateCount++;
+    }
+
+}