From fd4a7e261c8e3dde94d3e875c231a8de0052b740 Mon Sep 17 00:00:00 2001
From: Theo <theo@noreply.org>
Date: Sat, 27 May 2023 14:32:56 +0300
Subject: [PATCH] Added tests that check the correctness of accrual of rewards
 after redeploy Staking contract

---
 test/TestRewards.sol | 104 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 104 insertions(+)
 create mode 100644 test/TestRewards.sol

diff --git a/test/TestRewards.sol b/test/TestRewards.sol
new file mode 100644
index 0000000..3f50d83
--- /dev/null
+++ b/test/TestRewards.sol
@@ -0,0 +1,104 @@
+// 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 / 10e17);
+
+        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 / 10e17);
+
+        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;
+
+        // Remind that we have locked in Governance 25 000 TORN for TEST_ADDRESS_ONE while voting
+        uint256 stakerLockedBalance = governance.lockedBalance(TEST_ADDRESS_ONE);
+        require(stakerLockedBalance == 25_000 ether, "Invalid test staker locked balance");
+
+        uint256 stakerRewardsBeforeBurning = staking.checkReward(TEST_ADDRESS_ONE);
+        console2.log("Staking rewards before burning: %s TORN", stakerRewardsBeforeBurning / 10e17);
+
+        burnTokens(_governanceAddress, toBurn, staking);
+
+        uint256 stakerRewardsAfterBurning = staking.checkReward(TEST_ADDRESS_ONE);
+        console2.log("Staking rewards after burning 10 000 TORN: %s TORN", stakerRewardsAfterBurning / 10e17);
+        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 / 10e17);
+        console2.log("Staker received rewards: %s TORN", receivedReward / 10e17);
+
+        require(receivedReward == expectedRewards, "Expected and received rewards don't match");
+
+        uint256 stakerTORNbalanceBeforeGettingRewards = TORN.balanceOf(TEST_ADDRESS_ONE);
+        console2.log("Staker balance before getting (withdrawal) collected rewards: %s TORN", stakerTORNbalanceBeforeGettingRewards / 10e17);
+
+        vm.startPrank(TEST_ADDRESS_ONE);
+        staking.getReward();
+        vm.stopPrank();
+
+        uint256 stakerTORNbalanceAfterGettingRewards = TORN.balanceOf(TEST_ADDRESS_ONE);
+        console2.log("Staker balance after getting (withdrawal) collected rewards: %s TORN", stakerTORNbalanceAfterGettingRewards / 10e17);
+
+        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");
+    }
+}