update test to use deployer not address(this) (#1865)

river-build · Dec 18, 2024 · fc78b1b · fc78b1b
1 parent 535fc7d
commit fc78b1b
Showing 1 changed file with 30 additions and 15 deletions.
diff --git a/contracts/test/base/registry/DelegationProxy.t.sol b/contracts/test/base/registry/DelegationProxy.t.sol
@@ -1,61 +1,67 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.23;
 
-import "forge-std/Test.sol";
+import {TestUtils} from "contracts/test/utils/TestUtils.sol";
 import {ERC20Votes} from "@openzeppelin/contracts/token/ERC20/extensions/ERC20Votes.sol";
 import {LibClone} from "solady/utils/LibClone.sol";
 import {UpgradeableBeacon} from "solady/utils/UpgradeableBeacon.sol";
 import {Initializable_AlreadyInitialized} from "@river-build/diamond/src/facets/initializable/Initializable.sol";
 import {DeployRiverBase} from "contracts/scripts/deployments/utils/DeployRiverBase.s.sol";
 import {DelegationProxy} from "contracts/src/base/registry/facets/distribution/v2/DelegationProxy.sol";
 
-contract DelegationProxyTest is Test {
+contract DelegationProxyTest is TestUtils {
   DeployRiverBase internal deployRiverTokenBase = new DeployRiverBase();
+
+  address internal deployer;
   address internal river;
   address internal beacon;
 
   function setUp() public {
-    river = deployRiverTokenBase.deploy(address(this));
+    deployer = getDeployer();
+    river = deployRiverTokenBase.deploy(deployer);
     beacon = address(
-      new UpgradeableBeacon(address(this), address(new DelegationProxy()))
+      new UpgradeableBeacon(deployer, address(new DelegationProxy()))
     );
   }
 
   function test_initialize_revertIf_alreadyInitialized() public {
     address proxy = LibClone.deployERC1967BeaconProxy(beacon);
-    DelegationProxy(proxy).initialize(river, address(this));
+    DelegationProxy(proxy).initialize(river, deployer);
 
     vm.expectRevert(
       abi.encodeWithSelector(
         Initializable_AlreadyInitialized.selector,
         uint32(1)
       )
     );
-    DelegationProxy(proxy).initialize(river, address(this));
+
+    vm.prank(deployer);
+    DelegationProxy(proxy).initialize(river, deployer);
   }
 
   function test_fuzz_initialize(address delegatee) public {
     vm.assume(delegatee != address(0));
     address proxy = LibClone.deployERC1967BeaconProxy(beacon);
+
+    vm.prank(deployer);
     DelegationProxy(proxy).initialize(river, delegatee);
 
-    assertEq(DelegationProxy(proxy).factory(), address(this));
+    assertEq(DelegationProxy(proxy).factory(), deployer);
     assertEq(DelegationProxy(proxy).stakeToken(), river);
     assertEq(ERC20Votes(river).delegates(proxy), delegatee);
-    assertEq(
-      ERC20Votes(river).allowance(proxy, address(this)),
-      type(uint256).max
-    );
+    assertEq(ERC20Votes(river).allowance(proxy, deployer), type(uint256).max);
   }
 
   function test_fuzz_redelegate_revertIf_notFactory(address caller) public {
-    vm.assume(caller != address(this));
+    vm.assume(caller != deployer);
     address proxy = LibClone.deployERC1967BeaconProxy(beacon);
-    DelegationProxy(proxy).initialize(river, address(this));
+
+    vm.prank(deployer);
+    DelegationProxy(proxy).initialize(river, deployer);
 
     vm.prank(caller);
     vm.expectRevert();
-    DelegationProxy(proxy).redelegate(address(this));
+    DelegationProxy(proxy).redelegate(deployer);
   }
 
   function test_fuzz_redelegate(
@@ -65,20 +71,29 @@ contract DelegationProxyTest is Test {
     vm.assume(delegatee != address(0));
     vm.assume(delegatee != newDelegatee);
     address proxy = LibClone.deployERC1967BeaconProxy(beacon);
+
+    vm.prank(deployer);
     DelegationProxy(proxy).initialize(river, delegatee);
 
+    vm.prank(deployer);
     DelegationProxy(proxy).redelegate(newDelegatee);
+
     assertEq(ERC20Votes(river).delegates(proxy), newDelegatee);
   }
 
   function test_upgradeBeacon() public {
     address proxy = LibClone.deployERC1967BeaconProxy(beacon);
-    DelegationProxy(proxy).initialize(river, address(this));
+
+    vm.prank(deployer);
+    DelegationProxy(proxy).initialize(river, deployer);
 
     DelegationProxy newImpl = new DelegationProxy();
+
+    vm.prank(deployer);
     UpgradeableBeacon(beacon).upgradeTo(address(newImpl));
 
     // Verify proxy still works with new implementation
+    vm.prank(deployer);
     DelegationProxy(proxy).redelegate(address(1));
   }
 }