diff --git a/src/EtherFiRewardsRouter.sol b/src/EtherFiRewardsRouter.sol index 4531eb35..3a7a9f15 100644 --- a/src/EtherFiRewardsRouter.sol +++ b/src/EtherFiRewardsRouter.sol @@ -19,14 +19,16 @@ contract EtherFiRewardsRouter is OwnableUpgradeable, UUPSUpgradeable { event EthReceived(address indexed from, uint256 value); event EthSent(address indexed from, address indexed to, uint256 value); + event UpdatedTreasury(address indexed treasury); event Erc20Sent(address indexed caller, address indexed token, uint256 amount); event Erc721Sent(address indexed caller, address indexed token, uint256 tokenId); error IncorrectRole(); - constructor(address _liquidityPool, address _roleRegistry) { + constructor(address _liquidityPool, address _treasury, address _roleRegistry) { _disableInitializers(); liquidityPool = _liquidityPool; + treasury = _treasury; roleRegistry = RoleRegistry(_roleRegistry); } diff --git a/test/EtherFiRewardsRouter.t.sol b/test/EtherFiRewardsRouter.t.sol index 37a0f165..4b59dc9e 100644 --- a/test/EtherFiRewardsRouter.t.sol +++ b/test/EtherFiRewardsRouter.t.sol @@ -1,30 +1,36 @@ import "./TestSetup.sol"; import "../src/EtherFiRewardsRouter.sol"; import "../src/LiquidityPool.sol"; +import "forge-std/console2.sol"; contract EtherFiRewardsRouterTest is TestSetup { - address liquidityPoolAddress; EtherFiRewardsRouter etherfiRewardsRouterImplementation; UUPSProxy etherfiRewardsRouterProxy; EtherFiRewardsRouter etherfiRewardsRouterInstance; + function get_eeth() public { + vm.startPrank(address(etherfiRewardsRouterInstance)); + vm.deal(address(etherfiRewardsRouterInstance), 10 ether); + liquidityPoolInstance.deposit{value: 2 ether}(); + vm.stopPrank(); + } function setUp() public { setUpTests(); initializeRealisticFork(MAINNET_FORK); - vm.startPrank(owner); - etherfiRewardsRouterImplementation = new EtherFiRewardsRouter(address(liquidityPoolInstance), address(roleRegistry)); + vm.startPrank(superAdmin); + etherfiRewardsRouterImplementation = new EtherFiRewardsRouter(address(liquidityPoolInstance), address(treasuryInstance), address(roleRegistry)); etherfiRewardsRouterProxy = new UUPSProxy(address(etherfiRewardsRouterImplementation), ""); etherfiRewardsRouterInstance = EtherFiRewardsRouter(payable(address(etherfiRewardsRouterProxy))); etherfiRewardsRouterInstance.initialize(); vm.startPrank(superAdmin); roleRegistry.grantRole(etherfiRewardsRouterInstance.ETHERFI_ROUTER_ADMIN(), admin); + etherfiRewardsRouterInstance.transferOwnership(owner); vm.stopPrank(); - liquidityPoolAddress = address(liquidityPoolInstance); vm.deal(address(etherfiRewardsRouterInstance), 10 ether); - vm.stopPrank(); + vm.stopPrank(); } function test_withdrawToLiquidityPool() public { @@ -38,7 +44,7 @@ contract EtherFiRewardsRouterTest is TestSetup { function test_elRouterUpgrade() public { vm.startPrank(owner); - EtherFiRewardsRouter newEtherfiRewardsRouterImplementation = new EtherFiRewardsRouter(address(liquidityPoolInstance), address(roleRegistry)); + EtherFiRewardsRouter newEtherfiRewardsRouterImplementation = new EtherFiRewardsRouter(address(liquidityPoolInstance), address(treasuryInstance), address(roleRegistry)); address oldImplementation = etherfiRewardsRouterInstance.getImplementation(); etherfiRewardsRouterInstance.upgradeTo(address(newEtherfiRewardsRouterImplementation)); address newImplementation = etherfiRewardsRouterInstance.getImplementation(); @@ -46,7 +52,7 @@ contract EtherFiRewardsRouterTest is TestSetup { assert(newImplementation == address(newEtherfiRewardsRouterImplementation)); test_withdrawToLiquidityPool(); vm.stopPrank(); - } + } function test_checkEventEmitted() public { vm.startPrank(admin); @@ -54,5 +60,14 @@ contract EtherFiRewardsRouterTest is TestSetup { emit EtherFiRewardsRouter.EthSent(address(etherfiRewardsRouterInstance), liquidityPoolAddress, 10 ether); etherfiRewardsRouterInstance.withdrawToLiquidityPool(); vm.stopPrank(); - } -} \ No newline at end of file + } + + function test_recoverERC20() public { + get_eeth(); + vm.startPrank(admin); + uint256 balanceBefore = eETHInstance.balanceOf(address(treasuryInstance)); + etherfiRewardsRouterInstance.recoverERC20(address(eETHInstance), 1 ether); + uint256 balanceAfter = eETHInstance.balanceOf(address(treasuryInstance)); + assertApproxEqAbs(balanceAfter, balanceBefore + 1 ether, 1); + } +}