fluxswap/test/StakingGauge.t.sol

// SPDX-License-Identifier: MIT
pragma solidity 0.8.36;

import {Test} from "forge-std/Test.sol";
import {StdInvariant} from "forge-std/StdInvariant.sol";
import {Vm} from "forge-std/Vm.sol";
import {StakingGauge} from "../src/StakingGauge.sol";
import {FluxToken} from "../src/FluxToken.sol";
import {TestErc20} from "../src/TestErc20.sol";

contract StakingGaugeTest is Test {
    StakingGauge private gauge;
    TestErc20 private stakingToken;
    FluxToken private rewardToken;

    address private alice = address(0xA11CE);
    address private bob = address(0xB0B);

    uint256 private constant funded = 7000e18;

    function setUp() public {
        stakingToken = new TestErc20("Liquidity", "LP", 0);
        rewardToken = new FluxToken(address(this));
        gauge = new StakingGauge(address(stakingToken), address(rewardToken), address(this));

        rewardToken.mint(address(gauge), funded);
        gauge.notifyRewardAmount(funded);

        stakingToken.mint(alice, 1e24);
        stakingToken.mint(bob, 1e24);
        vm.prank(alice);
        stakingToken.approve(address(gauge), type(uint256).max);
        vm.prank(bob);
        stakingToken.approve(address(gauge), type(uint256).max);
    }

    function testSingleStakerEarnsWholeReward() public {
        vm.prank(alice);
        gauge.stake(100e18);

        vm.warp(block.timestamp + 7 days);

        uint256 earned = gauge.earned(alice);
        assertApproxEqRel(earned, funded, 1e15);

        vm.prank(alice);
        gauge.getReward();
        assertApproxEqRel(rewardToken.balanceOf(alice), funded, 1e15);
    }

    function testTwoStakersSplitProportionally() public {
        vm.prank(alice);
        gauge.stake(100e18);
        vm.prank(bob);
        gauge.stake(300e18);

        vm.warp(block.timestamp + 7 days);

        uint256 aliceEarned = gauge.earned(alice);
        uint256 bobEarned = gauge.earned(bob);

        assertApproxEqRel(aliceEarned, funded / 4, 1e15);
        assertApproxEqRel(bobEarned, (funded * 3) / 4, 1e15);
        assertApproxEqRel(aliceEarned + bobEarned, funded, 1e15);
    }

    function testStakeTransfersIn() public {
        vm.prank(alice);
        gauge.stake(100e18);
        assertEq(stakingToken.balanceOf(address(gauge)), 100e18);
        assertEq(gauge.totalSupply(), 100e18);
        assertEq(gauge.balanceOf(alice), 100e18);
    }

    function testWithdrawReturnsStake() public {
        vm.prank(alice);
        gauge.stake(100e18);
        vm.warp(block.timestamp + 1 days);
        vm.prank(alice);
        gauge.withdraw(40e18);
        assertEq(gauge.balanceOf(alice), 60e18);
        assertEq(stakingToken.balanceOf(alice), 1e24 - 60e18);
    }

    function testExitWithdrawsAndClaims() public {
        vm.prank(alice);
        gauge.stake(100e18);
        vm.warp(block.timestamp + 7 days);
        vm.prank(alice);
        gauge.exit();
        assertEq(gauge.balanceOf(alice), 0);
        assertEq(stakingToken.balanceOf(alice), 1e24);
        assertApproxEqRel(rewardToken.balanceOf(alice), funded, 1e15);
    }

    function testLateStakerEarnsLess() public {
        vm.prank(alice);
        gauge.stake(100e18);
        vm.warp(block.timestamp + 3 days + 12 hours);
        vm.prank(bob);
        gauge.stake(100e18);
        vm.warp(block.timestamp + 3 days + 12 hours);

        assertGt(gauge.earned(alice), gauge.earned(bob));
    }

    function testNotifyRevertsWhenUnderfunded() public {
        vm.warp(block.timestamp + 8 days);
        vm.expectRevert(bytes("StakingGauge: REWARD_TOO_HIGH"));
        gauge.notifyRewardAmount(1e30);
    }

    function testOnlyOwnerNotifies() public {
        vm.prank(alice);
        vm.expectRevert(bytes("StakingGauge: FORBIDDEN"));
        gauge.notifyRewardAmount(1e18);
    }

    function testFuzzEarnedNeverExceedsFunded(uint96 stakeAmount, uint32 elapsed) public {
        uint256 amount = bound(uint256(stakeAmount), 1e6, 1e24);
        uint256 dt = bound(uint256(elapsed), 0, 14 days);

        vm.prank(alice);
        gauge.stake(amount);
        vm.warp(block.timestamp + dt);

        assertLe(gauge.earned(alice), funded);
    }
}

contract StakingHandler {
    Vm private constant vm = Vm(address(uint160(uint256(keccak256("hevm cheat code")))));

    StakingGauge public gauge;
    TestErc20 public stakingToken;
    FluxToken public rewardToken;
    uint256 public paidOut;

    constructor(StakingGauge gaugeAddress, TestErc20 staking, FluxToken reward) {
        gauge = gaugeAddress;
        stakingToken = staking;
        rewardToken = reward;
        stakingToken.mint(address(this), 1e30);
        stakingToken.approve(address(gauge), type(uint256).max);
    }

    function bound(uint256 x, uint256 min, uint256 max) private pure returns (uint256) {
        if (max <= min) return min;
        return min + (x % (max - min + 1));
    }

    function doStake(uint256 amountSeed) external {
        uint256 amount = bound(amountSeed, 1, 1e24);
        gauge.stake(amount);
    }

    function doWithdraw(uint256 amountSeed) external {
        uint256 have = gauge.balanceOf(address(this));
        if (have == 0) return;
        gauge.withdraw(bound(amountSeed, 1, have));
    }

    function doClaim() external {
        uint256 before = rewardToken.balanceOf(address(this));
        gauge.getReward();
        paidOut += rewardToken.balanceOf(address(this)) - before;
    }

    function doWarp(uint32 dtSeed) external {
        vm.warp(block.timestamp + bound(uint256(dtSeed), 1, 2 days));
    }
}

contract StakingGaugeInvariantTest is StdInvariant, Test {
    StakingGauge private gauge;
    TestErc20 private stakingToken;
    FluxToken private rewardToken;
    StakingHandler private handler;

    uint256 private constant funded = 7000e18;

    function setUp() public {
        stakingToken = new TestErc20("Liquidity", "LP", 0);
        rewardToken = new FluxToken(address(this));
        gauge = new StakingGauge(address(stakingToken), address(rewardToken), address(this));

        rewardToken.mint(address(gauge), funded);
        gauge.notifyRewardAmount(funded);

        handler = new StakingHandler(gauge, stakingToken, rewardToken);

        bytes4[] memory selectors = new bytes4[](4);
        selectors[0] = StakingHandler.doStake.selector;
        selectors[1] = StakingHandler.doWithdraw.selector;
        selectors[2] = StakingHandler.doClaim.selector;
        selectors[3] = StakingHandler.doWarp.selector;
        targetSelector(FuzzSelector({addr: address(handler), selectors: selectors}));
        targetContract(address(handler));
    }

    function invariantRewardsNeverExceedFunding() public view {
        assertLe(handler.paidOut(), funded);
        assertEq(rewardToken.balanceOf(address(gauge)) + handler.paidOut(), funded);
    }

    function invariantStakingTokenFullyBacked() public view {
        assertEq(stakingToken.balanceOf(address(gauge)), gauge.totalSupply());
    }
}