test: bound rps to realistic ranges (#263)

* test: bound rps to realistic ranges

* fix: deposit amount in fuzz and invariant tests

* test(fuzz): add over streaming tests for wider ranges

* test: polish test names

---------

Co-authored-by: andreivladbrg <andreivladbrg@gmail.com>

src/SablierFlow.sol

@@ -92,8 +92,8 @@ contract SablierFlow is
                 uint128 scaleFactor = (10 ** (18 - tokenDecimals)).toUint128();
                 solvencyAmount = (balance - snapshotDebt + 1) * scaleFactor;
             }
-            uint128 solvencyPeriod = solvencyAmount / _streams[streamId].ratePerSecond.unwrap();
-            return _streams[streamId].snapshotTime + uint40(solvencyPeriod);
+            uint256 solvencyPeriod = solvencyAmount / _streams[streamId].ratePerSecond.unwrap();
+            return _streams[streamId].snapshotTime + solvencyPeriod.toUint40();
         }
     }
 

tests/integration/fuzz/Fuzz.t.sol

@@ -54,7 +54,7 @@ abstract contract Shared_Integration_Fuzz_Test is Integration_Test {
             // Hash the next stream ID and the decimal to generate a seed.
             uint128 amountSeed = uint128(uint256(keccak256(abi.encodePacked(flow.nextStreamId(), decimals))));
             // Bound the amount between a realistic range.
-            uint128 amount = boundUint128(amountSeed, 1, 1_000_000_000e18);
+            uint128 amount = boundUint128(amountSeed, 1e18, 200_000e18);
             uint128 depositAmount = getDescaledAmount(amount, decimals);
 
             // Deposit into the stream.

tests/integration/fuzz/create.t.sol

@@ -30,7 +30,7 @@ contract Create_Integration_Fuzz_Test is Shared_Integration_Fuzz_Test {
         vm.assume(sender != address(0) && recipient != address(0));
 
         // Bound the variables.
-        ratePerSecond = ud21x18(boundUint128(ratePerSecond.unwrap(), 1, UINT128_MAX - 1));
+        ratePerSecond = boundRatePerSecond(ratePerSecond);
         decimals = boundUint8(decimals, 0, 18);
 
         // Create a new token.

tests/integration/fuzz/restart.t.sol

@@ -27,7 +27,7 @@ contract Restart_Integration_Fuzz_Test is Shared_Integration_Fuzz_Test {
     {
         (streamId,,) = useFuzzedStreamOrCreate(streamId, decimals);
 
-        ratePerSecond = ud21x18(boundUint128(ratePerSecond.unwrap(), 1, UINT128_MAX));
+        ratePerSecond = boundRatePerSecond(ratePerSecond);
 
         // Pause the stream.
         flow.pause(streamId);

tests/integration/fuzz/withdrawMultiple.t.sol

@@ -3,21 +3,21 @@ pragma solidity >=0.8.22;
 
 import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import { ud21x18 } from "@prb/math/src/UD21x18.sol";
+import { ISablierFlow } from "src/interfaces/ISablierFlow.sol";
 
 import { Shared_Integration_Fuzz_Test } from "./Fuzz.t.sol";
 
 contract WithdrawMultiple_Delay_Fuzz_Test is Shared_Integration_Fuzz_Test {
     /// @dev Checklist:
-    /// - It should test multiple withdrawals from the stream.
+    /// - It should test multiple withdrawals from the stream using `withdrawMax`.
     /// - It should assert that the actual amount withdrawn is less than the desired amount.
     /// - It should check that stream delay and deviation are within acceptable limits for realistic values of rps.
     ///
     /// Given enough runs, all of the following scenarios should be fuzzed for USDC:
     /// - Multiple values for realistic rps.
     /// - Multiple withdrawal counts on the same stream at multiple points in time.
-    function testFuzz_WithdrawMultiple_Delay(uint128 rps, uint256 withdrawCount, uint40 timeJump) external {
-        // Bound the rps to a reasonable range [$100/month, $1000/month].
-        rps = boundUint128(rps, 38_580_246_913_580, 385_802_469_135_800);
+    function testFuzz_WithdrawMultiple_Usdc_SmallDelay(uint128 rps, uint256 withdrawCount, uint40 timeJump) external {
+        rps = boundRatePerSecond(ud21x18(rps)).unwrap();
 
         IERC20 token = createToken(DECIMALS);
         uint256 streamId = createDefaultStream(ud21x18(rps), token);
@@ -64,4 +64,107 @@ contract WithdrawMultiple_Delay_Fuzz_Test is Shared_Integration_Fuzz_Test {
         // Assert that actual amount withdrawn is always less than the desired amount.
         assertLe(actualTotalAmountWithdrawn, desiredTotalAmountWithdrawn);
     }
+
+    /// @dev Checklist:
+    /// - It should test multiple withdrawals from the stream using `withdrawMax`.
+    /// - It should assert that the actual amount withdrawn is always less than the desired amount.
+    ///
+    /// Given enough runs, all of the following scenarios should be fuzzed:
+    /// - Multiple values for decimals
+    /// - Multiple values for wide ranged rps.
+    /// - Multiple withdrawal counts on the same stream at multiple points in time.
+    function testFuzz_WithdrawMaxMultiple_RpsWideRange(
+        uint128 rps,
+        uint256 withdrawCount,
+        uint40 timeJump,
+        uint8 decimals
+    )
+        external
+    {
+        _test_WithdrawMultiple(rps, withdrawCount, timeJump, decimals, ISablierFlow.withdrawMax.selector, 0);
+    }
+
+    /// @dev Checklist:
+    /// - It should test multiple withdrawals from the stream using `withdraw`.
+    /// - It should assert that the actual amount withdrawn is always less than the desired amount.
+    ///
+    /// Given enough runs, all of the following scenarios should be fuzzed:
+    /// - Multiple values for decimals
+    /// - Multiple values for wide ranged rps.
+    /// - Multiple amounts to withdraw.
+    /// - Multiple withdrawal counts on the same stream at multiple points in time.
+    function testFuzz_WithdrawMultiple_RpsWideRange(
+        uint128 rps,
+        uint128 withdrawAmount,
+        uint256 withdrawCount,
+        uint40 timeJump,
+        uint8 decimals
+    )
+        external
+    {
+        _test_WithdrawMultiple(rps, withdrawCount, timeJump, decimals, ISablierFlow.withdraw.selector, withdrawAmount);
+    }
+
+    // Private helper function.
+    function _test_WithdrawMultiple(
+        uint128 rps,
+        uint256 withdrawCount,
+        uint40 timeJump,
+        uint8 decimals,
+        bytes4 selector,
+        uint128 withdrawAmount
+    )
+        private
+    {
+        decimals = boundUint8(decimals, 0, 18);
+        IERC20 token = createToken(decimals);
+
+        // Bound rate per second to a wider range for 18 decimals.
+        if (decimals == 18) {
+            rps = boundUint128(rps, 0.0000000001e18, 2e18);
+        }
+        // For all other decimals, choose the minimum rps such that it takes 1 minute to stream 1 token.
+        else {
+            rps = boundUint128(rps, getScaledAmount(1, decimals) / 60 + 1, 1e18);
+        }
+
+        uint256 streamId = createDefaultStream(ud21x18(rps), token);
+
+        withdrawCount = _bound(withdrawCount, 100, 200);
+
+        // Deposit the sufficient amount.
+        uint256 sufficientDepositAmount = getDescaledAmount(uint128(rps * 1 days * withdrawCount), decimals);
+        deposit(streamId, uint128(sufficientDepositAmount));
+
+        // Actual total amount withdrawn in a given run.
+        uint256 actualTotalAmountWithdrawn;
+
+        uint40 timeBeforeFirstWithdraw = getBlockTimestamp();
+
+        for (uint256 i; i < withdrawCount; ++i) {
+            // Warp the time.
+            timeJump = boundUint40(timeJump, 1 hours, 1 days);
+            vm.warp({ newTimestamp: getBlockTimestamp() + timeJump });
+
+            // Withdraw the tokens based on the selector value.
+            // ISablierFlow.withdraw
+            if (selector == ISablierFlow.withdraw.selector) {
+                withdrawAmount = boundUint128(withdrawAmount, 1, flow.withdrawableAmountOf(streamId));
+                flow.withdraw(streamId, users.recipient, withdrawAmount);
+            }
+            // ISablierFlow.withdrawMax
+            else if (selector == ISablierFlow.withdrawMax.selector) {
+                withdrawAmount = flow.withdrawMax(streamId, users.recipient);
+            }
+
+            // Update the actual total amount withdrawn.
+            actualTotalAmountWithdrawn += withdrawAmount;
+        }
+
+        uint40 totalStreamPeriod = getBlockTimestamp() - timeBeforeFirstWithdraw;
+        uint256 desiredTotalAmountWithdrawn = getDescaledAmount(rps * totalStreamPeriod, decimals);
+
+        // Assert that actual amount withdrawn is always less than the desired amount.
+        assertLe(actualTotalAmountWithdrawn, desiredTotalAmountWithdrawn);
+    }
 }

tests/invariant/handlers/FlowCreateHandler.sol

@@ -88,9 +88,10 @@ contract FlowCreateHandler is BaseHandler {
 
         // Calculate the upper bound, based on the token decimals, for the deposit amount.
         uint128 upperBound = getDescaledAmount(1_000_000e18, IERC20Metadata(address(currentToken)).decimals());
+        uint128 lowerBound = getDescaledAmount(1e18, IERC20Metadata(address(currentToken)).decimals());
 
         // Make sure the deposit amount is non-zero and less than values that could cause an overflow.
-        vm.assume(params.depositAmount >= 100 && params.depositAmount <= upperBound);
+        vm.assume(params.depositAmount >= lowerBound && params.depositAmount <= upperBound);
 
         // Mint enough tokens to the Sender.
         deal({
@@ -139,10 +140,12 @@ contract FlowCreateHandler is BaseHandler {
         // Calculate the minimum value in scaled version that can be withdrawn for this token.
         uint128 mvt = getScaledAmount(1, decimals);
 
-        // Check the rate per second is within a realistic range such that it can also be smaller than mvt.
+        // For 18 decimal, check the rate per second is within a realistic range.
         if (decimals == 18) {
             vm.assume(params.ratePerSecond > 0.00001e18 && params.ratePerSecond <= 1e18);
-        } else {
+        }
+        // For all other decimals, choose the minimum rps such that it takes 100 seconds to stream 1 token.
+        else {
             vm.assume(params.ratePerSecond > mvt / 100 && params.ratePerSecond <= 1e18);
         }
     }

tests/invariant/handlers/FlowHandler.sol

@@ -124,9 +124,10 @@ contract FlowHandler is BaseHandler {
 
         // Calculate the upper bound, based on the token decimals, for the deposit amount.
         uint128 upperBound = getDescaledAmount(1_000_000e18, flow.getTokenDecimals(currentStreamId));
+        uint128 lowerBound = getDescaledAmount(1e18, flow.getTokenDecimals(currentStreamId));
 
         // Make sure the deposit amount is non-zero and less than values that could cause an overflow.
-        vm.assume(depositAmount >= 100 && depositAmount <= upperBound);
+        vm.assume(depositAmount >= lowerBound && depositAmount <= upperBound);
 
         IERC20 token = flow.getToken(currentStreamId);
 

tests/utils/Utils.sol

@@ -29,9 +29,9 @@ abstract contract Utils is CommonBase, Constants, PRBMathUtils {
         depositAmount = boundUint128(amount, 1, maxDepositAmount - 1);
     }
 
-    /// @dev Bounds the rate per second between a realistic range.
+    /// @dev Bounds the rate per second between a realistic range i.e. for USDC [$50/month $5000/month].
     function boundRatePerSecond(UD21x18 ratePerSecond) internal pure returns (UD21x18) {
-        return ud21x18(boundUint128(ratePerSecond.unwrap(), 0.0000000001e18, 10e18));
+        return ud21x18(boundUint128(ratePerSecond.unwrap(), 0.00002e18, 0.002e18));
     }
 
     /// @dev Bounds a `uint128` number.