harb/onchain/test/Harberg.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 {TwabController} from "pt-v5-twab-controller/TwabController.sol";
import "../src/Harberg.sol";

contract HarbergTest is Test {
    TwabController tc;
    Harberg harberg;
    address stakingPool;
    address liquidityPool;
    address liquidityManager;
    address taxPool;

    function setUp() public {
        tc = new TwabController(60 * 60, uint32(block.timestamp));
        harberg = new Harberg("HARB", "HARB", tc);
        taxPool = harberg.TAX_POOL();
        stakingPool = makeAddr("stakingPool");
        harberg.setStakingPool(stakingPool);
        liquidityPool = makeAddr("liquidityPool");
        harberg.setLiquidityPool(liquidityPool);
        liquidityManager = makeAddr("liquidityManager");
        harberg.setLiquidityManager(liquidityManager);
    }

    // Simulates staking by transferring tokens to the stakingPool address.
    function simulateStake(uint256 amount) internal {
        // the amount of token has to be available on the balance
        // of the test contract
        harberg.transfer(stakingPool, amount);
    }

    // Simulates unstaking by transferring tokens from the stakingPool back to a given address.
    function simulateUnstake(uint256 amount) internal {
        // Direct transfer from the stakingPool to 'to' address to simulate unstaking
        vm.prank(stakingPool);  // Assuming 'stake' contract would allow this in an actual scenario
        harberg.transfer(address(this), amount);
    }

    function testHarbergConstructor() public view {
        // Check if the token details are set as expected
        assertEq(harberg.name(), "HARB");
        assertEq(harberg.symbol(), "HARB");

        // Confirm that the TwabController address is correctly set
        (address _tc, address _lm, address _sp, address _lp) = harberg.peripheryContracts();
        assertEq(_tc, address(tc));
        assertEq(_lm, liquidityManager);
        assertEq(_sp, stakingPool);
        assertEq(_lp, liquidityPool);
    }

    function testMintWithEmptyStakingPool() public {
        uint256 initialSupply = harberg.totalSupply();
        uint256 mintAmount = 1000 * 1e18; // 1000 HARB tokens

        vm.prank(address(liquidityManager));
        harberg.mint(mintAmount);

        // Check if the total supply has increased correctly
        assertEq(harberg.totalSupply(), initialSupply + mintAmount);
        // Check if the staking pool balance is still 0, as before
        assertEq(harberg.balanceOf(stakingPool), 0);
    }

    function testBurnWithEmptyStakingPool() public {
        uint256 initialSupply = harberg.totalSupply();
        uint256 burnAmount = 500 * 1e18; // 500 HARB tokens

        // First, mint some tokens to burn
        vm.prank(address(liquidityManager));
        harberg.mint(burnAmount);

        vm.prank(address(liquidityManager));
        harberg.burn(burnAmount);

        // Check if the total supply has decreased correctly
        assertEq(harberg.totalSupply(), initialSupply);
        // Check if the staking pool balance has decreased correctly
        assertEq(harberg.balanceOf(stakingPool), 0);
    }

    function testMintImpactOnSimulatedStaking() public {
        uint256 initialStakingPoolBalance = harberg.balanceOf(stakingPool);
        uint256 mintAmount = 1000 * 1e18; // 1000 HARB tokens
        // Ensure the test contract has enough tokens to simulate staking
        vm.prank(address(liquidityManager));
        harberg.mint(mintAmount);
        vm.prank(address(liquidityManager));
        harberg.transfer(address(this), mintAmount);

        // Simulate staking of the minted amount
        simulateStake(mintAmount);

        // Check balances after simulated staking
        assertEq(harberg.balanceOf(stakingPool), initialStakingPoolBalance + mintAmount);
    }

    function testUnstakeImpactOnTotalSupply() public {
        uint256 stakeAmount = 500 * 1e18; // 500 HARB tokens
        // Ensure the test contract has enough tokens to simulate staking
        vm.prank(address(liquidityManager));
        harberg.mint(stakeAmount);
        vm.prank(address(liquidityManager));
        harberg.transfer(address(this), stakeAmount);

        uint256 initialTotalSupply = harberg.totalSupply();

        // Simulate staking and then unstaking
        simulateStake(stakeAmount);
        simulateUnstake(stakeAmount);

        // Check total supply remains unchanged after unstake
        assertEq(harberg.totalSupply(), initialTotalSupply);
    }

    // Fuzz test for mint function with varying stake amounts
    function testMintWithStake(uint8 _stakePercentage, uint256 mintAmount) public {
        uint256 initialAmount = 500 * 1e18;
        // Ensure the test contract has enough tokens to simulate staking
        vm.prank(address(liquidityManager));
        harberg.mint(initialAmount);
        vm.prank(address(liquidityManager));
        harberg.transfer(address(this), initialAmount);

        // Limit fuzzing input to 0% - 20%
        uint8 effectiveStakePercentage = _stakePercentage % 21;
        uint256 stakeAmount = (initialAmount * effectiveStakePercentage) / 100;
        simulateStake(stakeAmount);

        uint256 initialTotalSupply = harberg.totalSupply();
        uint256 initialStakingPoolBalance = harberg.balanceOf(stakingPool);

        mintAmount = bound(mintAmount, 0, 500 * 1e18);
        uint256 expectedNewStake = initialStakingPoolBalance * mintAmount / (initialTotalSupply - initialStakingPoolBalance);

        vm.prank(address(liquidityManager));
        harberg.mint(mintAmount);

        uint256 expectedStakingPoolBalance = initialStakingPoolBalance + expectedNewStake;
        uint256 expectedTotalSupply = initialTotalSupply + mintAmount + expectedNewStake;

        assertEq(harberg.totalSupply(), expectedTotalSupply, "Total supply did not match expected after mint.");
        assertEq(harberg.balanceOf(stakingPool), expectedStakingPoolBalance, "Staking pool balance did not adjust correctly after mint.");
    }

    // Fuzz test for burn function with varying stake amounts
    function testBurnWithStake(uint8 _stakePercentage, uint256 burnAmount) public {
        uint256 mintAmount = 500 * 1e18;
        // Ensure the test contract has enough tokens to simulate staking
        vm.prank(address(liquidityManager));
        harberg.mint(mintAmount);

        // Limit fuzzing input to 0% - 20%
        uint8 effectiveStakePercentage = _stakePercentage % 21;
        uint256 stakeAmount = (mintAmount * effectiveStakePercentage) / 100;
        vm.prank(address(liquidityManager));
        harberg.transfer(address(this), stakeAmount);
        simulateStake(stakeAmount);

        burnAmount = bound(burnAmount, 0, 200 * 1e18);
        uint256 initialTotalSupply = harberg.totalSupply();
        uint256 initialStakingPoolBalance = harberg.balanceOf(stakingPool);
        uint256 expectedExcessStake = initialStakingPoolBalance * burnAmount / (initialTotalSupply - initialStakingPoolBalance);

        vm.prank(address(liquidityManager));
        harberg.burn(burnAmount);

        uint256 expectedStakingPoolBalance = initialStakingPoolBalance - expectedExcessStake;
        uint256 expectedTotalSupply = initialTotalSupply - burnAmount - expectedExcessStake;

        assertEq(harberg.totalSupply(), expectedTotalSupply, "Total supply did not match expected after burn.");
        assertEq(harberg.balanceOf(stakingPool), expectedStakingPoolBalance, "Staking pool balance did not adjust correctly after burn.");
    }

    function testTaxAccumulation() public {
        uint256 taxAmount = 100 * 1e18; // 100 HARB tokens

        vm.prank(address(liquidityManager));
        harberg.mint(taxAmount);

        // Initial tax collected should be zero
        assertEq(harberg.sumTaxCollected(), 0, "Initial tax collected should be zero.");

        // Simulate sending tokens to the taxPool
        vm.prank(address(liquidityManager));
        harberg.transfer(taxPool, taxAmount);

        // Check that sumTaxCollected has been updated correctly
        assertEq(harberg.sumTaxCollected(), taxAmount, "Tax collected not updated correctly after transfer to taxPool.");
    }

    function testUBIClaimBySingleAccountOverTime() public {
        uint256 initialSupply = 1000 * 1e18; // 1000 HARB tokens
        uint256 taxAmount = 200 * 1e18; // 200 HARB tokens to be collected as tax
        address user = makeAddr("alice");

        // Setup initial supply and distribute to user
        vm.prank(address(liquidityManager));
        harberg.mint(initialSupply + taxAmount);
        vm.prank(address(liquidityManager));
        harberg.transfer(user, initialSupply);

        // Simulate tax collection
        vm.prank(address(liquidityManager));
        harberg.transfer(taxPool, taxAmount);

        // Simulate time passage to ensure TWAB is recorded over time
        vm.warp(block.timestamp + 30 days);

        // Assert initial user balance and sumTaxCollected before claiming UBI
        assertEq(harberg.balanceOf(user), initialSupply, "User should hold the entire initial supply.");
        assertEq(harberg.sumTaxCollected(), taxAmount, "Tax collected should match the tax amount transferred.");

        // User claims UBI
        vm.prank(user);
        harberg.claimUbi(user);

        // Compute expected UBI
        // Assume the user is the only one holding tokens, they get all the collected taxes.
        uint256 expectedUbiAmount = taxAmount;

        // Verify UBI claim
        uint256 postClaimBalance = harberg.balanceOf(user);
        assertEq(postClaimBalance, initialSupply + expectedUbiAmount, "User's balance after claiming UBI is incorrect.");

        // Ensure that claiming doesn't affect the total supply
        uint256 expectedTotalSupply = initialSupply + taxAmount; // Include minted tax
        assertEq(harberg.totalSupply(), expectedTotalSupply, "Total supply should be unchanged after UBI claim.");
    }

    function testUBIClaimBySingleAccountWithWraparound() public {
        uint256 initialSupply = 1000 * 1e18; // 1000 HARB tokens
        uint256 taxAmount = 200 * 1e18; // 200 HARB tokens to be collected as tax
        uint256 nearMaxUint = type(uint256).max - 10;
        address user = makeAddr("alice");

        // Set sumTaxCollected to near max value to simulate wrap-around
        vm.store(
            address(harberg),
            bytes32(uint256(9)),
            bytes32(nearMaxUint)
        );

        // Read the value back to confirm it's set correctly
        assertEq(harberg.sumTaxCollected(), nearMaxUint, "Initial sumTaxCollected should be near max uint256");

        // Setup initial supply and distribute to user
        vm.prank(address(liquidityManager));
        harberg.mint(initialSupply + taxAmount);
        vm.prank(address(liquidityManager));
        harberg.transfer(user, initialSupply);

        // Simulate tax collection to cause overflow
        vm.prank(address(liquidityManager));
        harberg.transfer(taxPool, taxAmount);

        // Simulate time passage to ensure TWAB is recorded over time
        vm.warp(block.timestamp + 30 days);

        // Verify the new value of sumTaxCollected after overflow
        uint256 newSumTaxCollected = harberg.sumTaxCollected();
        assertGt(taxAmount, newSumTaxCollected, "sumTaxCollected should have wrapped around and be less than taxAmount");

        // User claims UBI
        vm.prank(user);
        harberg.claimUbi(user);

        // Compute expected UBI
        // Assume the user is the only one holding tokens, they get all the collected taxes.
        uint256 expectedUbiAmount = taxAmount;

        // Verify UBI claim
        uint256 postClaimBalance = harberg.balanceOf(user);
        assertApproxEqRel(postClaimBalance, initialSupply + expectedUbiAmount, 1 * 1e14, "User's balance after claiming UBI is incorrect.");

        // Ensure that claiming doesn't affect the total supply
        uint256 expectedTotalSupply = initialSupply + taxAmount; // Include minted tax
        assertEq(harberg.totalSupply(), expectedTotalSupply, "Total supply should be unchanged after UBI claim.");
    }

    function testUBIClaimByMultipleAccountsWithDifferentHoldingPeriods() public {
        uint256 initialSupply = 1000 * 1e18; // 1000 HARB tokens
        uint256 taxAmount = 300 * 1e18; // 300 HARB tokens to be collected as tax
        address account1 = makeAddr("alice");
        address account2 = makeAddr("bob");
        address account3 = makeAddr("charly");

        // Setup initial supply and distribute to users
        vm.startPrank(address(liquidityManager));
        harberg.mint(initialSupply);
        harberg.transfer(account1, 400 * 1e18); // Account 1 gets 400 tokens
        harberg.transfer(account2, 300 * 1e18); // Account 2 gets 300 tokens
        harberg.transfer(account3, 300 * 1e18); // Account 3 gets 300 tokens
        vm.stopPrank();

        uint256 startTime = block.timestamp;

        // Simulate different holding periods
        vm.warp(block.timestamp + 10 days); // Fast forward 10 days
        vm.prank(account1);
        harberg.transfer(account2, 100 * 1e18); // Account 1 transfers 100 tokens to Account 2

        vm.warp(block.timestamp + 20 days); // Fast forward another 20 days
        vm.prank(account2);
        harberg.transfer(account3, 100 * 1e18); // Account 2 transfers 100 tokens to Account 3

        // Simulate tax collection after the transactions
        vm.startPrank(address(liquidityManager));
        harberg.mint(taxAmount);
        harberg.transfer(taxPool, taxAmount);
        vm.stopPrank();

        // Assert sumTaxCollected before claiming UBI
        assertEq(harberg.sumTaxCollected(), taxAmount, "Tax collected should match the tax amount transferred.");

        // Each account claims UBI
        vm.prank(account1);
        harberg.claimUbi(account1);
        vm.prank(account2);
        harberg.claimUbi(account2);
        vm.prank(account3);
        harberg.claimUbi(account3);


        // Assert the post-claim balances reflect the TWAB calculations
        {
            uint256 totalDistributed = harberg.balanceOf(account1) + harberg.balanceOf(account2) + harberg.balanceOf(account3) - initialSupply;
            // Tolerance setup: 0.01% of the total tax amount
            uint256 lowerBound = taxAmount - (taxAmount / 10000);
            assertTrue(totalDistributed >= lowerBound && totalDistributed <= totalDistributed, "Total distributed UBI does not match the total tax collected within an acceptable tolerance range.");
        }

        // Calculate expected UBI amounts based on simplified TWAB assumptions
        // These should be replaced with actual TWAB calculations from your contract
        uint256 totalPeriod = block.timestamp - startTime;
        uint256 account1TWAB = (400 * 1e18 * 10 days + 300 * 1e18 * (totalPeriod - 10 days)) / totalPeriod;
        uint256 account2TWAB = (300 * 1e18 * 10 days + 400 * 1e18 * 20 days + 300 * 1e18 * (totalPeriod - 30 days)) / totalPeriod;
        uint256 account3TWAB = (300 * 1e18 * 30 days + 400 * 1e18 * (totalPeriod - 30 days)) / totalPeriod;

        uint256 totalTWAB = account1TWAB + account2TWAB + account3TWAB;

        // Calculate exact expected UBI payouts
        uint256 expectedBalance1 = (taxAmount * account1TWAB) / totalTWAB;
        uint256 expectedBalance2 = (taxAmount * account2TWAB) / totalTWAB;
        uint256 expectedBalance3 = (taxAmount * account3TWAB) / totalTWAB;

        // Assert the post-claim balances reflect the TWAB calculations with a smaller rounding tolerance
        // 1 * 1e14; // 0.0001 HARB token tolerance for rounding errors
        assertApproxEqRel(harberg.balanceOf(account1) - 300 * 1e18, expectedBalance1, 1 * 1e14, "Account 1's balance after claiming UBI is incorrect.");
        assertApproxEqRel(harberg.balanceOf(account2) - 300 * 1e18, expectedBalance2, 1 * 1e14, "Account 2's balance after claiming UBI is incorrect.");
        assertApproxEqRel(harberg.balanceOf(account3) - 400 * 1e18, expectedBalance3, 1 * 1e14, "Account 3's balance after claiming UBI is incorrect.");
    }

    function testUBIClaimWithoutAnyTaxCollected() public {
        uint256 initialSupply = 1000 * 1e18; // 1000 HARB tokens
        address user = makeAddr("alice");

        // Setup initial supply and allocate to user
        vm.startPrank(address(liquidityManager));
        harberg.mint(initialSupply);
        harberg.transfer(user, initialSupply);
        vm.stopPrank();

        // Ensure no tax has been collected yet
        assertEq(harberg.sumTaxCollected(), 0, "Initial tax collected should be zero.");

        // Simulate time passage to ensure TWAB is recorded
        vm.warp(block.timestamp + 30 days);

        // User attempts to claim UBI
        vm.prank(user);
        vm.expectRevert("No UBI to claim."); // Assuming your contract reverts with a message when there's no UBI to claim
        harberg.claimUbi(user);

        // Ensure the user's balance remains unchanged as no UBI should be distributed
        assertEq(harberg.balanceOf(user), initialSupply, "User's balance should not change after attempting to claim UBI without any taxes collected.");

        // Check if sumTaxCollected remains zero after the claim attempt
        assertEq(harberg.sumTaxCollected(), 0, "No tax should be collected, and sumTaxCollected should remain zero after the claim attempt.");
    }

    function testEdgeCaseWithMaximumTaxCollection() public {
        uint256 initialSupply = 1e24; // Large number of HARB tokens to simulate realistic large-scale deployment
        uint256 maxTaxAmount = type(uint96).max - initialSupply; // Setting max tax just below overflow threshold when added to total supply

        address account1 = makeAddr("alice");

        // Setup initial supply and allocate to user
        vm.startPrank(address(liquidityManager));
        harberg.mint(initialSupply + maxTaxAmount);
        harberg.transfer(account1, initialSupply);
        harberg.transfer(taxPool, maxTaxAmount); // Simulate tax collection at the theoretical maximum
        vm.stopPrank();

        // Assert that maximum tax was collected
        assertEq(harberg.sumTaxCollected(), maxTaxAmount, "Max tax collected should match the max tax amount transferred.");

        // Simulate time passage and UBI claim
        vm.warp(block.timestamp + 30 days);

        // Account 1 claims UBI
        vm.prank(account1);
        harberg.claimUbi(account1);

        // Check if the account's balance increased correctly
        uint256 expectedBalance = initialSupply + maxTaxAmount; // This assumes the entire tax pool goes to one account, simplify as needed
        assertEq(harberg.balanceOf(account1), expectedBalance, "Account 1's balance after claiming UBI with max tax collection is incorrect.");

        // Verify that no taxes are left unclaimed
        assertEq(harberg.balanceOf(taxPool), 0, "All taxes should be claimed after the UBI claim.");
    }


    // TODO: why is this test passing even though it exceeds MAX_CARDINALITY?
    function testTwabBeyondBuffer() public {
        uint256 initialSupply = 1000 * 1e18; // 1000 HARB tokens
        uint256 taxAmount = 300 * 1e18; // 300 HARB tokens to be collected as tax

        address account1 = makeAddr("alice");

        // Setup initial supply and allocate to user
        vm.startPrank(address(liquidityManager));
        harberg.mint(initialSupply + taxAmount);
        harberg.transfer(account1, initialSupply / 800);
        harberg.transfer(taxPool, taxAmount); // Simulate tax collection at the theoretical maximum
        harberg.transfer(liquidityPool, harberg.balanceOf(address(liquidityManager)));
        vm.stopPrank();
        vm.warp(block.timestamp + 1 hours);


        // Simulate updates over a longer period, e.g., enough to potentially wrap the buffer.
        uint numHours = 399;  // More than 365 to potentially test buffer wrapping (MAX_CARDINALITY)
        for (uint i = 0; i < numHours; i++) {
            vm.prank(liquidityPool);
            harberg.transfer(account1, initialSupply / 800);
            vm.warp(block.timestamp + 1 hours);  // Fast-forward time by one hour.
        }

        // Account 1 claims UBI
        vm.prank(account1);
        uint256 ubiCollected = harberg.claimUbi(account1);

        // Check if the account's balance increased correctly
        uint256 expectedBalance = (initialSupply / 2) + ubiCollected; // This assumes the entire tax pool goes to one account, simplify as needed
        assertApproxEqRel(harberg.balanceOf(account1), expectedBalance, 1 * 1e18, "Account 1's balance after claiming UBI with max tax collection is incorrect.");

        // Verify that no taxes are left unclaimed
        assertEq(harberg.balanceOf(taxPool), 0, "All taxes should be claimed after the UBI claim.");
    }

}