feat: implement asOwner modifier.

OriginProtocol · Jul 22, 2024 · 97c0d28 · 97c0d28
1 parent 381f2db
commit 97c0d28
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 23 deletions.
diff --git 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
@@ -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
@@ -36,38 +36,35 @@ 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);
     }
 
     //////////////////////////////////////////////////////
     /// --- 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
         assertEq(weth.balanceOf(address(this)), 10 ether);
         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
@@ -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
@@ -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();
+    }
+}