hyperlane-monorepo/solidity/test/lib/RateLimited.t.sol

// SPDX-License-Identifier: MIT or Apache-2.0
pragma solidity ^0.8.13;
import {Test} from "forge-std/Test.sol";
import {RateLimited} from "../../contracts/libs/RateLimited.sol";

contract RateLimitLibTest is Test {
    RateLimited rateLimited;
    uint256 constant MAX_CAPACITY = 1 ether;
    uint256 constant ONE_PERCENT = 0.01 ether; // Used for assertApproxEqRel
    address HOOK = makeAddr("HOOK");

    function setUp() public {
        rateLimited = new RateLimited(MAX_CAPACITY);
    }

    function testRateLimited_setsNewLimit() external {
        rateLimited.setRefillRate(2 ether);
        assertApproxEqRel(rateLimited.maxCapacity(), 2 ether, ONE_PERCENT);
        assertEq(rateLimited.refillRate(), uint256(2 ether) / 1 days); // 2 ether / 1 day
    }

    function testRateLimited_revertsIfMaxNotSet() external {
        rateLimited.setRefillRate(0);
        vm.expectRevert();
        rateLimited.calculateCurrentLevel();
    }

    function testRateLimited_returnsCurrentFilledLevel_anyDay(
        uint40 time
    ) external {
        bound(time, 1 days, 2 days);
        vm.warp(time);

        // Using approx because division won't be exact
        assertApproxEqRel(
            rateLimited.calculateCurrentLevel(),
            MAX_CAPACITY,
            ONE_PERCENT
        );
    }

    function testRateLimited_onlyOwnerCanSetTargetLimit() external {
        vm.prank(address(0));
        vm.expectRevert();
        rateLimited.setRefillRate(1 ether);
    }

    function testRateLimited_neverReturnsGtMaxLimit(
        uint256 _newAmount,
        uint40 _newTime
    ) external {
        vm.warp(_newTime);
        vm.assume(_newAmount <= rateLimited.calculateCurrentLevel());
        rateLimited.validateAndConsumeFilledLevel(_newAmount);
        assertLe(
            rateLimited.calculateCurrentLevel(),
            rateLimited.maxCapacity()
        );
    }

    function testRateLimited_decreasesLimitWithinSameDay() external {
        vm.warp(1 days);
        uint256 currentTargetLimit = rateLimited.calculateCurrentLevel();
        uint256 amount = 0.4 ether;
        uint256 newLimit = rateLimited.validateAndConsumeFilledLevel(amount);
        assertEq(newLimit, currentTargetLimit - amount);

        // Consume the same amount
        currentTargetLimit = rateLimited.calculateCurrentLevel();
        newLimit = rateLimited.validateAndConsumeFilledLevel(amount);
        assertEq(newLimit, currentTargetLimit - amount);

        // One more to exceed limit
        vm.expectRevert();
        rateLimited.validateAndConsumeFilledLevel(amount);
    }

    function testRateLimited_replinishesWithinSameDay() external {
        vm.warp(1 days);
        uint256 amount = 0.95 ether;
        uint256 newLimit = rateLimited.validateAndConsumeFilledLevel(amount);
        uint256 currentTargetLimit = rateLimited.calculateCurrentLevel();
        assertApproxEqRel(currentTargetLimit, 0.05 ether, ONE_PERCENT);

        // Warp to near end-of-day
        vm.warp(block.timestamp + 0.99 days);
        newLimit = rateLimited.validateAndConsumeFilledLevel(amount);
        assertApproxEqRel(newLimit, 0.05 ether, ONE_PERCENT);
    }

    function testRateLimited_shouldResetLimit_ifDurationExceeds(
        uint256 _amount
    ) external {
        // Transfer less than the limit
        vm.warp(0.5 days);
        uint256 currentTargetLimit = rateLimited.calculateCurrentLevel();
        vm.assume(_amount < currentTargetLimit);

        uint256 newLimit = rateLimited.validateAndConsumeFilledLevel(_amount);
        assertApproxEqRel(newLimit, currentTargetLimit - _amount, ONE_PERCENT);

        // Warp to a new cycle
        vm.warp(10 days);
        currentTargetLimit = rateLimited.calculateCurrentLevel();
        assertApproxEqRel(currentTargetLimit, MAX_CAPACITY, ONE_PERCENT);
    }
}