From 4e300274cd28f79f0d61088be5b84f65ffcc6c3a Mon Sep 17 00:00:00 2001 From: Rajath Alex Date: Mon, 8 Apr 2024 19:33:21 -0400 Subject: [PATCH] test_ClaimRewards --- .../LlamaRewardsClaimAccountExtension.sol | 2 +- test/mock/MockRewardsContract.sol | 5 ++ .../LlamaRewardsClaimAccountExtension.t.sol | 84 ++++++++++++++++++- 3 files changed, 89 insertions(+), 2 deletions(-) diff --git a/src/rewards-claimer/LlamaRewardsClaimAccountExtension.sol b/src/rewards-claimer/LlamaRewardsClaimAccountExtension.sol index 71dfbb7..e767c3d 100644 --- a/src/rewards-claimer/LlamaRewardsClaimAccountExtension.sol +++ b/src/rewards-claimer/LlamaRewardsClaimAccountExtension.sol @@ -39,7 +39,7 @@ contract LlamaRewardsClaimAccountExtension is LlamaBaseAccountExtension { function claimRewards(TargetData[] memory targetData) external onlyDelegateCall returns (bytes[] memory returnData) { uint256 length = targetData.length; returnData = new bytes[](length); - for (uint256 i = 0; i < length; LlamaUtils.uncheckedIncrement(i)) { + for (uint256 i = 0; i < length; i = LlamaUtils.uncheckedIncrement(i)) { address target = targetData[i].target; uint256 value = targetData[i].value; bytes memory callData = targetData[i].data; diff --git a/test/mock/MockRewardsContract.sol b/test/mock/MockRewardsContract.sol index 01823c8..e170e2c 100644 --- a/test/mock/MockRewardsContract.sol +++ b/test/mock/MockRewardsContract.sol @@ -15,6 +15,9 @@ contract MockRewardsContract { _; } + event ethWithdrawn(address indexed from, address indexed to, uint256 amount); + event erc20Withdrawn(IERC20 indexed token, address indexed from, address indexed to, uint256 amount); + address public immutable REWARD_CLAIMER; constructor(address rewardClaimer) { @@ -22,10 +25,12 @@ contract MockRewardsContract { } function withdrawETH() external payable onlyRewardClaimer { + emit ethWithdrawn(address(this), REWARD_CLAIMER, address(this).balance); Address.sendValue(payable(REWARD_CLAIMER), address(this).balance); } function withdrawERC20(IERC20 token) external onlyRewardClaimer { + emit erc20Withdrawn(token, address(this), REWARD_CLAIMER, token.balanceOf(address(this))); token.safeTransfer(REWARD_CLAIMER, token.balanceOf(address(this))); } } diff --git a/test/rewards-claimer/LlamaRewardsClaimAccountExtension.t.sol b/test/rewards-claimer/LlamaRewardsClaimAccountExtension.t.sol index 5a24265..4312d7c 100644 --- a/test/rewards-claimer/LlamaRewardsClaimAccountExtension.t.sol +++ b/test/rewards-claimer/LlamaRewardsClaimAccountExtension.t.sol @@ -3,14 +3,20 @@ pragma solidity ^0.8.23; import {Test, console2} from "forge-std/Test.sol"; +import {IERC20} from "@openzeppelin/token/ERC20/IERC20.sol"; + import {MockRewardsContract} from "test/mock/MockRewardsContract.sol"; import {LlamaRewardsClaimTestSetup} from "test/rewards-claimer/LlamaRewardsClaimTestSetup.sol"; import {ILlamaAccount} from "src/interfaces/ILlamaAccount.sol"; +import {ActionInfo} from "src/lib/Structs.sol"; import {LlamaRewardsClaimAccountExtension} from "src/rewards-claimer/LlamaRewardsClaimAccountExtension.sol"; import {LlamaRewardsClaimGuard} from "src/rewards-claimer/LlamaRewardsClaimGuard.sol"; contract LlamaRewardsClaimAccountExtensionTest is LlamaRewardsClaimTestSetup { + event ethWithdrawn(address indexed from, address indexed to, uint256 amount); + event erc20Withdrawn(IERC20 indexed token, address indexed from, address indexed to, uint256 amount); + function setUp() public override { LlamaRewardsClaimTestSetup.setUp(); } @@ -22,4 +28,80 @@ contract Constructor is LlamaRewardsClaimAccountExtensionTest { } } -contract ClaimRewards is LlamaRewardsClaimAccountExtensionTest {} +contract ClaimRewards is LlamaRewardsClaimAccountExtensionTest { + function _createAndQueueActionClaimRewards() public returns (ActionInfo memory actionInfo) { + LlamaRewardsClaimAccountExtension.TargetData[] memory targetData = _setupClaimRewardsData(); + bytes memory accountExtensionData = abi.encodeCall(LlamaRewardsClaimAccountExtension.claimRewards, (targetData)); + bytes memory data = + abi.encodeCall(ILlamaAccount.execute, (address(rewardsClaimAccountExtension), true, 0, accountExtensionData)); + + // Create an action to claim rewards. + vm.prank(coreTeam1); + uint256 actionId = CORE.createAction(CORE_TEAM_ROLE, STRATEGY, address(ACCOUNT), 0, data, ""); + actionInfo = ActionInfo(actionId, coreTeam1, CORE_TEAM_ROLE, STRATEGY, address(ACCOUNT), 0, data); + + // Approve and queue the action. + vm.prank(coreTeam1); + CORE.castApproval(CORE_TEAM_ROLE, actionInfo, ""); + vm.prank(coreTeam2); + CORE.castApproval(CORE_TEAM_ROLE, actionInfo, ""); + vm.prank(coreTeam3); + CORE.castApproval(CORE_TEAM_ROLE, actionInfo, ""); + } + + function test_ClaimRewards() public { + ActionInfo memory actionInfo = _createAndQueueActionClaimRewards(); + + assertEq(address(rewardsContract1).balance, 1 ether); + assertEq(USDC.balanceOf(address(rewardsContract1)), 1 ether); + assertEq(address(rewardsContract2).balance, 1 ether); + assertEq(USDC.balanceOf(address(rewardsContract2)), 1 ether); + assertEq(address(rewardsContract3).balance, 1 ether); + assertEq(USDC.balanceOf(address(rewardsContract3)), 1 ether); + assertEq(address(rewardsContract4).balance, 1 ether); + assertEq(USDC.balanceOf(address(rewardsContract4)), 1 ether); + assertEq(address(rewardsContract5).balance, 1 ether); + assertEq(USDC.balanceOf(address(rewardsContract5)), 1 ether); + + uint256 initialLlamaAccountETHBalance = address(ACCOUNT).balance; + uint256 initialLlamaAccountUSDCBalance = USDC.balanceOf(address(ACCOUNT)); + + // Execute the action. + vm.expectEmit(); + emit ethWithdrawn(address(rewardsContract1), address(ACCOUNT), 1 ether); + vm.expectEmit(); + emit erc20Withdrawn(USDC, address(rewardsContract1), address(ACCOUNT), 1 ether); + vm.expectEmit(); + emit ethWithdrawn(address(rewardsContract2), address(ACCOUNT), 1 ether); + vm.expectEmit(); + emit erc20Withdrawn(USDC, address(rewardsContract2), address(ACCOUNT), 1 ether); + vm.expectEmit(); + emit ethWithdrawn(address(rewardsContract3), address(ACCOUNT), 1 ether); + vm.expectEmit(); + emit erc20Withdrawn(USDC, address(rewardsContract3), address(ACCOUNT), 1 ether); + vm.expectEmit(); + emit ethWithdrawn(address(rewardsContract4), address(ACCOUNT), 1 ether); + vm.expectEmit(); + emit erc20Withdrawn(USDC, address(rewardsContract4), address(ACCOUNT), 1 ether); + vm.expectEmit(); + emit ethWithdrawn(address(rewardsContract5), address(ACCOUNT), 1 ether); + vm.expectEmit(); + emit erc20Withdrawn(USDC, address(rewardsContract5), address(ACCOUNT), 1 ether); + + CORE.executeAction(actionInfo); + + assertEq(address(rewardsContract1).balance, 0); + assertEq(USDC.balanceOf(address(rewardsContract1)), 0); + assertEq(address(rewardsContract2).balance, 0); + assertEq(USDC.balanceOf(address(rewardsContract2)), 0); + assertEq(address(rewardsContract3).balance, 0); + assertEq(USDC.balanceOf(address(rewardsContract3)), 0); + assertEq(address(rewardsContract4).balance, 0); + assertEq(USDC.balanceOf(address(rewardsContract4)), 0); + assertEq(address(rewardsContract5).balance, 0); + assertEq(USDC.balanceOf(address(rewardsContract5)), 0); + + assertEq(address(ACCOUNT).balance, initialLlamaAccountETHBalance + 5 ether); + assertEq(USDC.balanceOf(address(ACCOUNT)), initialLlamaAccountUSDCBalance + 5 ether); + } +}