From 97c0d28cfd35f64026edafbfc8398eafecbc94d5 Mon Sep 17 00:00:00 2001 From: clement-ux Date: Mon, 22 Jul 2024 17:04:38 +0200 Subject: [PATCH] feat: implement asOwner modifier. --- test/fork/concrete/Ownable.t.sol | 6 ++---- test/fork/concrete/Proxy.t.sol | 6 ++---- test/fork/concrete/Transfer.t.sol | 9 +++------ test/fork/concrete/Withdraw.t.sol | 12 +++--------- test/fork/utils/Modifiers.sol | 21 +++++++++++++++++++++ 5 files changed, 31 insertions(+), 23 deletions(-) create mode 100644 test/fork/utils/Modifiers.sol diff --git a/test/fork/concrete/Ownable.t.sol b/test/fork/concrete/Ownable.t.sol index 63bbb1b..41fc695 100644 --- a/test/fork/concrete/Ownable.t.sol +++ b/test/fork/concrete/Ownable.t.sol @@ -27,22 +27,20 @@ contract Fork_Concrete_OethARM_Ownable_Test_ is Fork_Shared_Test_ { ////////////////////////////////////////////////////// /// --- PASSING TESTS ////////////////////////////////////////////////////// - function test_SetOperator() public { + function test_SetOperator() public asOwner { // Assertions before assertEq(oethARM.operator(), address(0)); - vm.prank(oethARM.owner()); oethARM.setOperator(operator); // Assertions after assertEq(oethARM.operator(), operator); } - function test_SetOwner() public { + function test_SetOwner() public asOwner { // Assertions before assertEq(oethARM.owner(), Mainnet.TIMELOCK); - vm.prank(oethARM.owner()); oethARM.setOwner(alice); // Assertions after diff --git a/test/fork/concrete/Proxy.t.sol b/test/fork/concrete/Proxy.t.sol index 30082b4..88f1599 100644 --- a/test/fork/concrete/Proxy.t.sol +++ b/test/fork/concrete/Proxy.t.sol @@ -32,12 +32,11 @@ contract Fork_Concrete_OethARM_Proxy_Test_ is Fork_Shared_Test_ { ////////////////////////////////////////////////////// /// --- PASSING TESTS ////////////////////////////////////////////////////// - function test_Upgrade() public { + function test_Upgrade() public asOwner { address owner = Mainnet.TIMELOCK; // Deploy new implementation OEthARM newImplementation = new OEthARM(); - vm.prank(owner); proxy.upgradeTo(address(newImplementation)); assertEq(proxy.implementation(), address(newImplementation)); @@ -50,14 +49,13 @@ contract Fork_Concrete_OethARM_Proxy_Test_ is Fork_Shared_Test_ { assertEq(address(oethARM.token1()), Mainnet.WETH); } - function test_UpgradeAndCall() public { + function test_UpgradeAndCall() public asOwner { address owner = Mainnet.TIMELOCK; // Deploy new implementation OEthARM newImplementation = new OEthARM(); bytes memory data = abi.encodeWithSignature("setOperator(address)", address(0x123)); - vm.prank(owner); proxy.upgradeToAndCall(address(newImplementation), data); assertEq(proxy.implementation(), address(newImplementation)); diff --git a/test/fork/concrete/Transfer.t.sol b/test/fork/concrete/Transfer.t.sol index 218ba92..aa3c124 100644 --- a/test/fork/concrete/Transfer.t.sol +++ b/test/fork/concrete/Transfer.t.sol @@ -36,10 +36,9 @@ contract Fork_Concrete_OethARM_Transfer_Test_ is Fork_Shared_Test_ { oethARM.transferEth(address(0), 0); } - function test_RevertWhen_TransferETH_Because_ETHTransferFailed() public { + function test_RevertWhen_TransferETH_Because_ETHTransferFailed() public asOwner { shoudRevertOnReceive = true; - vm.prank(oethARM.owner()); vm.expectRevert("ARM: ETH transfer failed"); oethARM.transferEth(address(this), 10 ether); } @@ -47,14 +46,13 @@ contract Fork_Concrete_OethARM_Transfer_Test_ is Fork_Shared_Test_ { ////////////////////////////////////////////////////// /// --- PASSING TESTS ////////////////////////////////////////////////////// - function test_TransferToken() public { + function test_TransferToken() public asOwner { // Assertions before assertEq(weth.balanceOf(address(this)), 0); assertEq(weth.balanceOf(address(oethARM)), 100 ether); vm.expectEmit({emitter: address(weth)}); emit IERC20.Transfer(address(oethARM), address(this), 10 ether); - vm.prank(oethARM.owner()); oethARM.transferToken(address(weth), address(this), 10 ether); // Assertions after @@ -62,12 +60,11 @@ contract Fork_Concrete_OethARM_Transfer_Test_ is Fork_Shared_Test_ { assertEq(weth.balanceOf(address(oethARM)), 90 ether); } - function test_TransferETH() public { + function test_TransferETH() public asOwner { // Assertions before uint256 balanceBefore = address(this).balance; assertEq(address(oethARM).balance, 100 ether); - vm.prank(oethARM.owner()); oethARM.transferEth(address(this), 10 ether); // Assertions after diff --git a/test/fork/concrete/Withdraw.t.sol b/test/fork/concrete/Withdraw.t.sol index 4f155e9..92d222d 100644 --- a/test/fork/concrete/Withdraw.t.sol +++ b/test/fork/concrete/Withdraw.t.sol @@ -49,8 +49,7 @@ contract Fork_Concrete_OethARM_Withdraw_Test_ is Fork_Shared_Test_ { ////////////////////////////////////////////////////// /// --- PASSING TESTS ////////////////////////////////////////////////////// - function test_RequestWithdraw() public { - vm.prank(oethARM.owner()); + function test_RequestWithdraw() public asOwner { vm.expectEmit({emitter: address(oeth)}); emit IERC20.Transfer(address(oethARM), address(0), 1 ether); (uint256 requestId, uint256 queued) = oethARM.requestWithdrawal(1 ether); @@ -61,28 +60,24 @@ contract Fork_Concrete_OethARM_Withdraw_Test_ is Fork_Shared_Test_ { assertEq(oeth.balanceOf(address(oethARM)), 9 ether, "OETH balance should be 99 ether"); } - function test_ClaimWithdraw_() public { + function test_ClaimWithdraw_() public asOwner { // First request withdrawal - vm.prank(oethARM.owner()); (uint256 requestId,) = oethARM.requestWithdrawal(1 ether); vault.addWithdrawalQueueLiquidity(); skip(10 minutes); // Todo: fetch direct value from contract // Then claim withdrawal - vm.prank(oethARM.owner()); oethARM.claimWithdrawal(requestId); // Assertions after assertEq(weth.balanceOf(address(oethARM)), 1 ether, "WETH balance should be 1 ether"); } - function test_ClaimWithdraws() public { + function test_ClaimWithdraws() public asOwner { // First request withdrawal - vm.startPrank(oethARM.owner()); oethARM.requestWithdrawal(1 ether); oethARM.requestWithdrawal(1 ether); - vm.stopPrank(); vault.addWithdrawalQueueLiquidity(); skip(10 minutes); // Todo: fetch direct value from contract @@ -91,7 +86,6 @@ contract Fork_Concrete_OethARM_Withdraw_Test_ is Fork_Shared_Test_ { requestIds[0] = 0; requestIds[1] = 1; // Then claim withdrawal - vm.prank(oethARM.owner()); oethARM.claimWithdrawals(requestIds); // Assertions after diff --git a/test/fork/utils/Modifiers.sol b/test/fork/utils/Modifiers.sol new file mode 100644 index 0000000..bb65c62 --- /dev/null +++ b/test/fork/utils/Modifiers.sol @@ -0,0 +1,21 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.23; + +// Test imports +import {Helpers} from "test/fork/utils/Helpers.sol"; + +abstract contract Modifiers is Helpers { + /// @notice Impersonate the owner of the contract. + modifier asOwner() { + vm.startPrank(oethARM.owner()); + _; + vm.stopPrank(); + } + + /// @notice Impersonate the governor of the vault. + modifier asGovernor() { + vm.startPrank(vault.governor()); + _; + vm.stopPrank(); + } +}