// 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()); } }