harb/onchain/test/Stake.t.sol

// SPDX-License-Identifier: GPL-3.0-or-later
pragma solidity ^0.8.19;

import "forge-std/Test.sol";
import "forge-std/console.sol";
import "../src/Kraiken.sol";
import {TooMuchSnatch, Stake} from "../src/Stake.sol";

contract StakeTest is Test {
    Kraiken kraiken;
    Stake stakingPool;
    address liquidityPool;
    address liquidityManager;

    event PositionCreated(
        uint256 indexed positionId, address indexed owner, uint256 kraikenDeposit, uint256 share, uint32 taxRate
    );
    event PositionRemoved(uint256 indexed positionId, address indexed owner, uint256 kraikenPayout);

    function setUp() public {
        kraiken = new Kraiken("KRAIKEN", "KRK");
        stakingPool = new Stake(address(kraiken), makeAddr("taxRecipient"));
        kraiken.setStakingPool(address(stakingPool));
        liquidityManager = makeAddr("liquidityManager");
        kraiken.setLiquidityManager(liquidityManager);
    }

    function assertPosition(uint256 positionId, uint256 expectedShares, uint32 expectedTaxRate) private view {
        (uint256 shares,,,, uint32 taxRate) = stakingPool.positions(positionId);
        assertEq(shares, expectedShares, "Incorrect share amount for new position");
        assertEq(taxRate, expectedTaxRate, "Incorrect tax rate for new position");
    }

    function verifyPositionShrunkOrRemoved(uint256 positionId, uint256 initialStake) private view {
        (uint256 remainingShare,,,,) = stakingPool.positions(positionId);
        uint256 expectedInitialShares = stakingPool.assetsToShares(initialStake);
        bool positionRemoved = remainingShare == 0;
        bool positionShrunk = remainingShare < expectedInitialShares;

        assertTrue(positionRemoved || positionShrunk, "Position was not correctly shrunk or removed");
    }

    function testBasicStaking() public {
        // Setup
        uint256 stakeAmount = 1 ether;
        address staker = makeAddr("staker");

        vm.startPrank(liquidityManager);
        kraiken.mint(stakeAmount * 5);
        kraiken.transfer(staker, stakeAmount);
        vm.stopPrank();

        vm.startPrank(staker);

        // Approve and stake
        kraiken.approve(address(stakingPool), stakeAmount);
        uint256[] memory empty;
        uint256 sharesExpected = stakingPool.assetsToShares(stakeAmount);
        vm.expectEmit(address(stakingPool));
        emit PositionCreated(654321, staker, stakeAmount, sharesExpected, 1);
        uint256 positionId = stakingPool.snatch(stakeAmount, staker, 1, empty);

        // Check results
        assertEq(
            stakingPool.outstandingStake(),
            stakingPool.assetsToShares(stakeAmount),
            "Outstanding stake did not update correctly"
        );
        (uint256 share, address owner, uint32 creationTime,, uint32 taxRate) = stakingPool.positions(positionId);
        assertEq(stakingPool.sharesToAssets(share), stakeAmount, "Stake amount in position is incorrect");
        assertEq(owner, staker, "Stake owner is incorrect");
        assertEq(creationTime, uint32(block.timestamp), "Creation time is incorrect");
        assertEq(taxRate, 1, "Tax rate should be initialized to 1");

        vm.stopPrank();
    }

    function testUnstaking() public {
        // Setup: Create a staking position first
        uint256 stakeAmount = 1 ether;
        address staker = makeAddr("staker");

        vm.startPrank(liquidityManager);
        kraiken.mint(stakeAmount * 5); // Ensuring the staker has enough balance
        kraiken.transfer(staker, stakeAmount);
        vm.stopPrank();

        // Staker stakes tokens
        vm.startPrank(staker);
        kraiken.approve(address(stakingPool), stakeAmount);
        uint256[] memory empty;
        uint256 positionId = stakingPool.snatch(stakeAmount, staker, 1, empty);

        // Simulate time passage to accumulate tax liability
        vm.warp(block.timestamp + 3 days);

        // Calculate expected tax due at the moment of unstaking
        uint256 taxAmount = stakingPool.taxDue(positionId, 0);
        uint256 assetsAfterTax = stakeAmount - taxAmount;

        // Expect the PositionRemoved event with the expected parameters
        vm.expectEmit(true, true, true, true);
        emit PositionRemoved(positionId, staker, assetsAfterTax);

        // Perform unstaking
        stakingPool.exitPosition(positionId);

        // Check results after unstaking
        assertEq(kraiken.balanceOf(staker), assetsAfterTax, "Assets after tax not returned correctly");
        assertEq(stakingPool.outstandingStake(), 0, "Outstanding stake not updated correctly");

        // Ensure the position is cleared
        (, address owner, uint32 time,,) = stakingPool.positions(positionId);
        assertEq(time, 0, "Position time not cleared");
        assertEq(owner, address(0), "Position owner not cleared");

        vm.stopPrank();
    }

    function testSnatchFunction() public {
        uint256 initialStake1 = 1 ether;
        uint256 initialStake2 = 2 ether;
        address firstStaker = makeAddr("firstStaker");
        address secondStaker = makeAddr("secondStaker");
        address newStaker = makeAddr("newStaker");
        uint256 snatchAmount = 1.5 ether;

        // Mint and distribute tokens
        vm.startPrank(liquidityManager);
        kraiken.mint((initialStake1 + initialStake2) * 5);
        kraiken.transfer(firstStaker, initialStake1);
        kraiken.transfer(secondStaker, initialStake2);
        kraiken.transfer(newStaker, snatchAmount);
        vm.stopPrank();

        // Setup initial stakers
        uint256 positionId1 = doSnatch(firstStaker, initialStake1, 1);
        uint256 positionId2 = doSnatch(secondStaker, initialStake2, 5);

        // Snatch setup
        vm.startPrank(newStaker);
        kraiken.approve(address(stakingPool), snatchAmount);
        uint256 snatchShares = stakingPool.assetsToShares(snatchAmount);
        uint256[] memory targetPositions = new uint256[](2);
        targetPositions[0] = positionId1;
        targetPositions[1] = positionId2;

        // Perform snatch
        uint256 newPositionId = stakingPool.snatch(snatchAmount, newStaker, 10, targetPositions);

        // Verify new position
        assertPosition(newPositionId, snatchShares, 10);

        // Check old positions
        verifyPositionShrunkOrRemoved(1, initialStake1);
        verifyPositionShrunkOrRemoved(2, initialStake2);

        vm.stopPrank();
    }

    function doSnatch(address staker, uint256 amount, uint32 taxRate) private returns (uint256 positionId) {
        vm.startPrank(staker);
        kraiken.approve(address(stakingPool), amount);
        uint256[] memory empty;
        positionId = stakingPool.snatch(amount, staker, taxRate, empty);
        vm.stopPrank();
    }

    function bp(uint256 val) internal pure returns (uint256) {
        return val / 1e15;
    }

    function denormTR(uint256 normalizedTaxRate) internal pure returns (uint256) {
        return normalizedTaxRate * 97;
    }

    function testAvgTaxRateAndPercentageStaked() public {
        uint256 smallstake = 0.3e17;
        uint256 stakeOneThird = 1 ether;
        uint256 stakeTwoThird = 2 ether;
        address staker = makeAddr("staker");

        // Mint and distribute tokens
        vm.startPrank(liquidityManager);
        // mint all the tokens we will need in the test
        kraiken.mint((smallstake + stakeOneThird + stakeTwoThird) * 5);
        // send 20% of that to staker
        kraiken.transfer(staker, (smallstake + stakeOneThird + stakeTwoThird) * 2);
        vm.stopPrank();

        // Setup initial stakers
        uint256 positionId1 = doSnatch(staker, smallstake, 0);

        uint256 avgTaxRate;
        uint256 percentageStaked;
        avgTaxRate = stakingPool.getAverageTaxRate();
        percentageStaked = stakingPool.getPercentageStaked();

        // let this be about 10 basis points of tax rate
        assertApproxEqRel(bp(denormTR(avgTaxRate)), 10, 1e17);
        assertApproxEqRel(bp(percentageStaked), 10, 1e17);

        vm.prank(staker);
        stakingPool.exitPosition(positionId1);
        uint256 positionId2 = doSnatch(staker, stakeOneThird, 2);

        avgTaxRate = stakingPool.getAverageTaxRate();
        percentageStaked = stakingPool.getPercentageStaked();

        assertApproxEqRel(bp(denormTR(avgTaxRate)), 50, 1e17);
        assertApproxEqRel(bp(percentageStaked), 300, 1e17);

        vm.prank(staker);
        stakingPool.exitPosition(positionId2);
        positionId1 = doSnatch(staker, stakeOneThird, 10);
        positionId2 = doSnatch(staker, stakeTwoThird, 11);

        avgTaxRate = stakingPool.getAverageTaxRate();
        percentageStaked = stakingPool.getPercentageStaked();

        assertApproxEqRel(bp(denormTR(avgTaxRate)), 730, 1e17);
        assertApproxEqRel(bp(percentageStaked), 1000, 1e17);

        vm.startPrank(staker);
        stakingPool.exitPosition(positionId1);
        stakingPool.exitPosition(positionId2);
        vm.stopPrank();
        positionId1 = doSnatch(staker, stakeOneThird, 29);
        positionId2 = doSnatch(staker, stakeTwoThird, 29);

        avgTaxRate = stakingPool.getAverageTaxRate();
        assertApproxEqRel(bp(denormTR(avgTaxRate)), 97000, 1e17);

        vm.startPrank(staker);
        stakingPool.exitPosition(positionId1);
        stakingPool.exitPosition(positionId2);
        vm.stopPrank();
        positionId2 = doSnatch(staker, stakeTwoThird, 15);

        avgTaxRate = stakingPool.getAverageTaxRate();
        percentageStaked = stakingPool.getPercentageStaked();

        assertApproxEqRel(bp(denormTR(avgTaxRate)), 2500, 1e17);
        assertApproxEqRel(bp(percentageStaked), 660, 1e17);

        vm.startPrank(staker);
        stakingPool.exitPosition(positionId2);
        vm.stopPrank();

        positionId1 = doSnatch(staker, stakeOneThird, 15);

        avgTaxRate = stakingPool.getAverageTaxRate();
        percentageStaked = stakingPool.getPercentageStaked();

        assertApproxEqRel(bp(denormTR(avgTaxRate)), 2500, 1e17);
        assertApproxEqRel(bp(percentageStaked), 330, 1e17);
    }

    function testRevert_SharesTooLow() public {
        address staker = makeAddr("staker");
        vm.startPrank(liquidityManager);
        kraiken.mint(10 ether);
        uint256 tooSmallStake = kraiken.previousTotalSupply() / 4000; // Less than minStake calculation
        kraiken.transfer(staker, tooSmallStake);
        vm.stopPrank();

        vm.startPrank(staker);
        kraiken.approve(address(stakingPool), tooSmallStake);

        uint256[] memory empty;
        vm.expectRevert(
            abi.encodeWithSelector(
                Stake.StakeTooLow.selector, staker, tooSmallStake, kraiken.previousTotalSupply() / 3000
            )
        );
        stakingPool.snatch(tooSmallStake, staker, 1, empty);
        vm.stopPrank();
    }

    function testRevert_TaxTooLow() public {
        address existingStaker = makeAddr("existingStaker");
        address newStaker = makeAddr("newStaker");
        vm.startPrank(liquidityManager);
        kraiken.mint(10 ether);
        kraiken.transfer(existingStaker, 1 ether);
        kraiken.transfer(newStaker, 1 ether);
        vm.stopPrank();

        uint256 positionId = doSnatch(existingStaker, 1 ether, 5); // Existing staker with tax rate 5

        vm.startPrank(newStaker);
        kraiken.transfer(newStaker, 1 ether);
        kraiken.approve(address(stakingPool), 1 ether);

        uint256[] memory positions = new uint256[](1);
        positions[0] = positionId; // Assuming position ID 1 has tax rate 5

        vm.expectRevert(abi.encodeWithSelector(Stake.TaxTooLow.selector, newStaker, 5, 5, positionId));
        stakingPool.snatch(1 ether, newStaker, 5, positions); // Same tax rate should fail
        vm.stopPrank();
    }

    function testRevert_TooMuchSnatch() public {
        address staker = makeAddr("staker");
        address ambitiousStaker = makeAddr("ambitiousStaker");

        vm.startPrank(liquidityManager);
        kraiken.mint(20 ether);
        kraiken.transfer(staker, 2 ether);
        kraiken.transfer(ambitiousStaker, 1 ether);
        vm.stopPrank();

        uint256 positionId = doSnatch(staker, 2 ether, 10);

        vm.startPrank(ambitiousStaker);
        kraiken.approve(address(stakingPool), 1 ether);

        uint256[] memory positions = new uint256[](1);
        positions[0] = positionId;
        vm.expectRevert(
            abi.encodeWithSelector(TooMuchSnatch.selector, ambitiousStaker, 500000 ether, 1000000 ether, 1000000 ether)
        );
        stakingPool.snatch(1 ether, ambitiousStaker, 20, positions);
        vm.stopPrank();
    }

    function testRevert_PositionNotFound() public {
        address staker = makeAddr("staker");

        vm.startPrank(liquidityManager);
        kraiken.mint(10 ether);
        kraiken.transfer(staker, 1 ether);
        vm.stopPrank();

        vm.startPrank(staker);
        kraiken.approve(address(stakingPool), 1 ether);

        uint256[] memory nonExistentPositions = new uint256[](1);
        nonExistentPositions[0] = 999; // Assumed non-existent position ID

        vm.expectRevert(abi.encodeWithSelector(Stake.PositionNotFound.selector, 999, staker));
        stakingPool.snatch(1 ether, staker, 15, nonExistentPositions);
        vm.stopPrank();
    }

    function testSuccessfulChangeTaxAndReverts() public {
        uint32 newTaxRate = 3; // New valid tax rate

        address staker = makeAddr("staker");

        vm.startPrank(liquidityManager);
        kraiken.mint(10 ether);
        kraiken.transfer(staker, 1 ether);
        vm.stopPrank();

        vm.startPrank(staker);
        kraiken.approve(address(stakingPool), 1 ether);
        uint256[] memory empty;
        uint256 positionId = stakingPool.snatch(1 ether, staker, 1, empty);

        uint32 invalidTaxRate = 9999; // Assuming this is out of bounds
        vm.expectRevert(bytes("tax rate out of bounds"));
        stakingPool.changeTax(positionId, invalidTaxRate);

        stakingPool.changeTax(positionId, newTaxRate);
        vm.stopPrank();

        // Verify the change
        (,,,, uint32 taxRate) = stakingPool.positions(positionId);
        assertEq(taxRate, newTaxRate, "Tax rate did not update correctly");

        // notOwner tries to change tax rate
        address notOwner = makeAddr("notOwner");
        vm.prank(notOwner);
        vm.expectRevert(abi.encodeWithSelector(Stake.NoPermission.selector, notOwner, staker));
        stakingPool.changeTax(positionId, newTaxRate * 2);
    }

    function testNormalTaxPayment() public {
        address staker = makeAddr("staker");

        vm.startPrank(liquidityManager);
        kraiken.mint(10 ether);
        kraiken.transfer(staker, 1 ether);
        vm.stopPrank();

        vm.startPrank(staker);
        kraiken.approve(address(stakingPool), 1 ether);
        uint256[] memory empty;
        uint256 positionId = stakingPool.snatch(1 ether, staker, 5, empty); // Using tax rate index 5, which is 18% per year
        (uint256 shareBefore,,,,) = stakingPool.positions(positionId);

        // Immediately after staking, no tax due
        stakingPool.payTax(positionId);
        // Verify no change in position
        (uint256 share,,,,) = stakingPool.positions(positionId);
        assertEq(share, shareBefore, "Share should not change when no tax is due");

        // Move time forward 30 days
        vm.warp(block.timestamp + 30 days);
        stakingPool.payTax(positionId);
        vm.stopPrank();

        // Check that the tax was paid and position updated
        (share,,,,) = stakingPool.positions(positionId);
        uint256 daysElapsed = 30;
        uint256 taxRate = 18; // Corresponding to 18% annually
        uint256 daysInYear = 365;
        uint256 taxBase = 100;

        uint256 taxFractionForTime = taxRate * daysElapsed * 1 ether / daysInYear / taxBase;
        uint256 expectedShareAfterTax = (1 ether - taxFractionForTime) * 1000000;

        assertTrue(share < shareBefore, "Share should decrease correctly after tax payment 1");
        assertEq(share, expectedShareAfterTax, "Share should decrease correctly after tax payment 2");
    }

    function testLiquidation() public {
        address staker = makeAddr("staker");

        vm.startPrank(liquidityManager);
        kraiken.mint(10 ether);
        kraiken.transfer(staker, 1 ether);
        vm.stopPrank();

        vm.startPrank(staker);
        kraiken.approve(address(stakingPool), 1 ether);
        uint256[] memory empty;
        uint256 positionId = stakingPool.snatch(1 ether, staker, 12, empty); // Using tax rate index 5, which is 100% per year
        vm.warp(block.timestamp + 365 days); // Move time forward to ensure maximum tax due
        stakingPool.payTax(positionId);
        vm.stopPrank();

        // Verify position is liquidated
        (uint256 share,,,,) = stakingPool.positions(positionId);
        assertEq(share, 0, "Share should be zero after liquidation");
    }
}