proposal-22-governance-and-rewards-patch/test/forge/TestRewards.sol

// SPDX-License-Identifier: MIT
pragma solidity ^0.6.12;
pragma experimental ABIEncoderV2;

import { console2 } from "@forge-std/console2.sol";
import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";

import { MockProposal } from "./MockProposal.sol";
import { GovernancePatchUpgrade } from "@root/v4-patch/GovernancePatchUpgrade.sol";
import { RelayerRegistry } from "@root/v4-patch/RelayerRegistry.sol";
import { TornadoStakingRewards } from "@root/v4-patch/TornadoStakingRewards.sol";

contract TestGovernanceStakingRewards is MockProposal {
    modifier replenishGovernanceStakingBalanceBefore() {
        vm.startPrank(_governanceAddress);
        IERC20(_tokenAddress).transfer(getStakingProxyAddress(), 100_000 ether);
        vm.stopPrank();
        _;
    }

    function burnTokens(address caller, uint256 amount, TornadoStakingRewards staking) internal {
        vm.startPrank(caller);
        staking.addBurnRewards(amount);
        vm.stopPrank();
    }

    function testAccumulatedRewardCanBeUpdated()
        public
        executeAttackerProposalBefore
        executeCurrentProposalBefore
    {
        TornadoStakingRewards staking = TornadoStakingRewards(getStakingProxyAddress());

        uint256 accumulatedRewardPerTornBeforeBurning =
            staking.accumulatedRewardPerTorn() / _tornMaximumSupply;

        console2.log(
            "Accumulated reward per TORN right after proposal execution: %s TORN",
            accumulatedRewardPerTornBeforeBurning / _tornDecimals
        );

        burnTokens(_governanceAddress, 10_000_000 ether, staking);

        uint256 accumulatedRewardPerTornAfterBurning = staking.accumulatedRewardPerTorn() / _tornMaximumSupply;

        console2.log(
            "Accumulated reward per TORN after burning 10 000 000 TORN: ~ %s TORN",
            accumulatedRewardPerTornAfterBurning / _tornDecimals
        );

        require(
            accumulatedRewardPerTornAfterBurning > accumulatedRewardPerTornBeforeBurning,
            "Staking rewards isn't updated"
        );
    }

    function testStakerCanGetRewards()
        public
        executeAttackerProposalBefore
        executeCurrentProposalBefore
        replenishGovernanceStakingBalanceBefore
    {
        TornadoStakingRewards staking = TornadoStakingRewards(getStakingProxyAddress());
        IERC20 TORN = IERC20(_tokenAddress);

        uint256 toBurn = 10_000 ether;

        retrieveAndLockBalance(TEST_STAKER_PRIVATE_KEY, TEST_STAKER_ADDRESS, PROPOSAL_THRESHOLD);
        uint256 stakerLockedBalance = governance.lockedBalance(TEST_STAKER_ADDRESS);
        require(stakerLockedBalance == PROPOSAL_THRESHOLD, "Invalid test staker locked balance");

        uint256 stakerRewardsBeforeBurning = staking.checkReward(TEST_STAKER_ADDRESS);
        console2.log("Staking rewards before burning: %s TORN", stakerRewardsBeforeBurning / _tornDecimals);

        burnTokens(_governanceAddress, toBurn, staking);

        uint256 stakerRewardsAfterBurning = staking.checkReward(TEST_STAKER_ADDRESS);
        console2.log(
            "Staking rewards after burning 10 000 TORN: %s TORN\n", stakerRewardsAfterBurning / _tornDecimals
        );
        require(stakerRewardsAfterBurning > stakerRewardsBeforeBurning, "Rewards isn't changed after burning");

        // All TORN, locked by users in Governance, is on the userVault contract balance
        uint256 governanceLockedAmount = TORN.balanceOf(address(governance.userVault()));
        uint256 receivedReward = stakerRewardsAfterBurning - stakerRewardsBeforeBurning;
        uint256 expectedRewards = stakerLockedBalance * toBurn / governanceLockedAmount;

        console2.log("Expected staking rewards: %s TORN", expectedRewards / _tornDecimals);
        console2.log("Staker received rewards: %s TORN\n", receivedReward / _tornDecimals);

        require(receivedReward == expectedRewards, "Expected and received rewards don't match");

        uint256 stakerTORNbalanceBeforeGettingRewards = TORN.balanceOf(TEST_STAKER_ADDRESS);
        console2.log(
            "Staker balance before getting (withdrawal) collected rewards: %s TORN",
            stakerTORNbalanceBeforeGettingRewards / _tornDecimals
        );

        vm.startPrank(TEST_STAKER_ADDRESS);
        staking.getReward();
        vm.stopPrank();

        uint256 stakerTORNbalanceAfterGettingRewards = TORN.balanceOf(TEST_STAKER_ADDRESS);
        console2.log(
            "Staker balance after getting (withdrawal) collected rewards: %s TORN",
            stakerTORNbalanceAfterGettingRewards / _tornDecimals
        );

        require(
            stakerTORNbalanceAfterGettingRewards > stakerTORNbalanceBeforeGettingRewards,
            "Rewards isn't withdrawed"
        );
        require(
            stakerTORNbalanceAfterGettingRewards - stakerTORNbalanceBeforeGettingRewards == receivedReward,
            "Incorrect rewards amount withdrawed"
        );
    }

    function testCanBurnFromRelayerRegistry() public {
        TornadoStakingRewards staking = TornadoStakingRewards(getStakingProxyAddress());
        address relayerRegistryAddress = getRelayerRegistryProxyAddress();

        uint256 accumulatedRewardPerTornBeforeBurning = staking.accumulatedRewardPerTorn();

        burnTokens(relayerRegistryAddress, 10_000 ether, staking);

        require(
            staking.accumulatedRewardPerTorn() > accumulatedRewardPerTornBeforeBurning,
            "Burn from relayer registry failed"
        );
    }
}