test: stata claim reward snapshot-diffs

Co-authored-by: Lukas <[email protected]>

Author: brotherlymite

snapshots/StataTokenV2.json

@@ -1,4 +1,5 @@
 {
+  "claimRewards": "359669",
   "deposit": "280955",
   "depositATokens": "219311",
   "redeem": "205837",

tests/gas/StataToken.gas.t.sol

@@ -5,7 +5,8 @@ import 'forge-std/Test.sol';
 
 import {Errors} from '../../src/contracts/protocol/libraries/helpers/Errors.sol';
 import {UserConfiguration} from '../../src/contracts/protocol/libraries/configuration/UserConfiguration.sol';
-import {Testhelpers, IERC20} from './Testhelpers.sol';
+import {PullRewardsTransferStrategy} from '../../src/contracts/rewards/transfer-strategies/PullRewardsTransferStrategy.sol';
+import {Testhelpers, TestnetERC20, IERC20} from './Testhelpers.sol';
 import {StataTokenFactory} from '../../src/contracts/extensions/stata-token/StataTokenFactory.sol';
 import {StataTokenV2} from '../../src/contracts/extensions/stata-token/StataTokenV2.sol';
 import {DataTypes} from '../../src/contracts/protocol/libraries/types/DataTypes.sol';
@@ -16,6 +17,11 @@ import {DataTypes} from '../../src/contracts/protocol/libraries/types/DataTypes.
 contract StataToken_gas_Tests is Testhelpers {
   StataTokenV2 public stataToken;
 
+  address public rewardTokenOne;
+  address public rewardTokenTwo;
+  address public emissionAdmin;
+  PullRewardsTransferStrategy strategy;
+
   function setUp() public override {
     super.setUp();
     StataTokenFactory(report.staticATokenFactoryProxy).createStataTokens(
@@ -24,6 +30,20 @@ contract StataToken_gas_Tests is Testhelpers {
     stataToken = StataTokenV2(
       StataTokenFactory(report.staticATokenFactoryProxy).getStataToken(tokenList.usdx)
     );
+
+    emissionAdmin = vm.addr(1024);
+    rewardTokenOne = address(new TestnetERC20('LM Reward ERC20 One', 'RWD_1', 18, poolAdmin));
+    rewardTokenTwo = address(new TestnetERC20('LM Reward ERC20 Two', 'RWD_2', 18, poolAdmin));
+    strategy = new PullRewardsTransferStrategy(
+      report.rewardsControllerProxy,
+      emissionAdmin,
+      emissionAdmin
+    );
+
+    vm.startPrank(poolAdmin);
+    contracts.emissionManager.setEmissionAdmin(rewardTokenOne, emissionAdmin);
+    contracts.emissionManager.setEmissionAdmin(rewardTokenTwo, emissionAdmin);
+    vm.stopPrank();
   }
 
   function test_deposit() external {
@@ -52,4 +72,41 @@ contract StataToken_gas_Tests is Testhelpers {
     stataToken.redeemATokens(shares, address(this), address(this));
     vm.snapshotGasLastCall('StataTokenV2', 'redeemAToken');
   }
+
+  function test_claimRewards() external {
+    uint256 amountToDeposit = 1000e8;
+    _supplyOnReserve(address(this), amountToDeposit, tokenList.usdx);
+    DataTypes.ReserveDataLegacy memory reserveData = contracts.poolProxy.getReserveData(
+      tokenList.usdx
+    );
+    IERC20(reserveData.aTokenAddress).approve(address(stataToken), amountToDeposit);
+    stataToken.depositATokens(amountToDeposit, address(this));
+
+    uint32 endTimestamp = uint32(vm.getBlockTimestamp() + 15 days);
+    _setupEmission(
+      rewardTokenOne,
+      reserveData.aTokenAddress,
+      endTimestamp,
+      1e14,
+      emissionAdmin,
+      address(strategy)
+    );
+    _setupEmission(
+      rewardTokenTwo,
+      reserveData.aTokenAddress,
+      endTimestamp,
+      1e14,
+      emissionAdmin,
+      address(strategy)
+    );
+    stataToken.refreshRewardTokens();
+    vm.warp(endTimestamp);
+
+    address[] memory rewardTokens = new address[](2);
+    rewardTokens[0] = rewardTokenOne;
+    rewardTokens[1] = rewardTokenTwo;
+
+    stataToken.claimRewards(address(this), rewardTokens);
+    vm.snapshotGasLastCall('StataTokenV2', 'claimRewards');
+  }
 }

tests/gas/Testhelpers.sol

@@ -3,7 +3,10 @@ pragma solidity ^0.8.0;
 
 import {IERC20} from 'openzeppelin-contracts/contracts/interfaces/IERC20.sol';
 import {IPriceOracleGetter} from '../../src/contracts/interfaces/IPriceOracleGetter.sol';
-import {TestnetProcedures} from '../utils/TestnetProcedures.sol';
+import {RewardsDataTypes} from '../../src/contracts/rewards/libraries/RewardsDataTypes.sol';
+import {TestnetProcedures, TestnetERC20} from '../utils/TestnetProcedures.sol';
+import {ITransferStrategyBase} from '../../src/contracts/rewards/transfer-strategies/PullRewardsTransferStrategy.sol';
+import {AggregatorInterface} from '../../src/contracts/dependencies/chainlink/AggregatorInterface.sol';
 
 contract Testhelpers is TestnetProcedures {
   address rando = makeAddr('randomUser');
@@ -12,7 +15,7 @@ contract Testhelpers is TestnetProcedures {
     initTestEnvironment(false);
 
     // supply and borrow some on reserve with a random user as "some" interest accrual
-    // is the realisitc usecase we want to check in gas snapshots
+    // is the realistic use case we want to check in gas snapshots
     _supplyOnReserve(rando, 100 ether, tokenList.weth);
     _supplyOnReserve(rando, 1_000_000e6, tokenList.usdx);
     _supplyOnReserve(rando, 100e8, tokenList.wbtc);
@@ -43,7 +46,7 @@ contract Testhelpers is TestnetProcedures {
       abi.encodeWithSelector(IPriceOracleGetter.getAssetPrice.selector, address(asset)),
       abi.encode(0)
     );
-    // borrow the full emount of the asset
+    // borrow the full amount of the asset
     vm.prank(borrower);
     contracts.poolProxy.borrow(asset, amount, 2, 0, borrower);
     // revert the oracle price
@@ -58,4 +61,36 @@ contract Testhelpers is TestnetProcedures {
     vm.warp(vm.getBlockTimestamp() + amount * 12);
     vm.roll(vm.getBlockNumber() + amount);
   }
+
+  function _setupEmission(
+    address rewardToken,
+    address asset,
+    uint32 emissionEnd,
+    uint88 emissionPerSecond,
+    address emissionAdmin,
+    address strategy
+  ) internal {
+    RewardsDataTypes.RewardsConfigInput[] memory config = new RewardsDataTypes.RewardsConfigInput[](
+      1
+    );
+    config[0] = RewardsDataTypes.RewardsConfigInput(
+      emissionPerSecond,
+      0,
+      emissionEnd,
+      asset,
+      rewardToken,
+      ITransferStrategyBase(strategy),
+      AggregatorInterface(address(2))
+    );
+
+    // configure asset
+    vm.prank(emissionAdmin);
+    contracts.emissionManager.configureAssets(config);
+
+    // fund admin & approve transfers to allow claiming
+    uint256 fundsToEmit = (emissionEnd - vm.getBlockTimestamp()) * emissionPerSecond;
+    deal(rewardToken, emissionAdmin, fundsToEmit, true);
+    vm.prank(emissionAdmin);
+    IERC20(rewardToken).approve(address(strategy), fundsToEmit);
+  }
 }