feat(scripts): add upgrade script

This introduces a simple upgrade script template that can be used to upgrade the stake manager.
vacp2p · Dec 9, 2024 · dcc0807 · dcc0807
1 parent a565dbb
commit dcc0807
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 26 deletions.
diff --git a/.gas-report b/.gas-report
@@ -14,6 +14,14 @@
 | activeNetworkConfig                                     | 454             | 454 | 454    | 454 | 114     |
 
 
+| script/UpgradeRewardsStreamerMP.s.sol:UpgradeRewardsStreamerMPScript contract |                 |         |         |         |         |
+|-------------------------------------------------------------------------------|-----------------|---------|---------|---------|---------|
+| Deployment Cost                                                               | Deployment Size |         |         |         |         |
+| 2845816                                                                       | 14079           |         |         |         |         |
+| Function Name                                                                 | min             | avg     | median  | max     | # calls |
+| run                                                                           | 2383391         | 2383391 | 2383391 | 2383391 | 3       |
+
+
 | src/RewardsStreamer.sol:RewardsStreamer contract |                 |        |        |        |         |
 |--------------------------------------------------|-----------------|--------|--------|--------|---------|
 | Deployment Cost                                  | Deployment Size |        |        |        |         |
@@ -66,7 +74,7 @@
 | unstake                                              | 60382           | 60919  | 60382  | 63877  | 13      |
 | updateAccountMP                                      | 15396           | 18474  | 17898  | 34998  | 21      |
 | updateGlobalState                                    | 11088           | 28045  | 25249  | 110284 | 21      |
-| upgradeToAndCall                                     | 3225            | 9387   | 10926  | 10936  | 5       |
+| upgradeToAndCall                                     | 3225            | 7887   | 8426   | 10936  | 5       |
 
 
 | src/StakeManagerProxy.sol:StakeManagerProxy contract |                 |       |        |        |         |
@@ -101,7 +109,7 @@
 | totalStaked                                          | 823             | 823   | 823    | 823    | 82      |
 | updateAccountMP                                      | 41755           | 44833 | 44257  | 61357  | 21      |
 | updateGlobalState                                    | 37076           | 54033 | 51237  | 136272 | 21      |
-| upgradeToAndCall                                     | 29868           | 36025 | 37562  | 37572  | 5       |
+| upgradeToAndCall                                     | 29868           | 33720 | 33720  | 37572  | 2       |
 
 
 | src/StakeVault.sol:StakeVault contract |                 |        |        |        |         |

diff --git a/.gas-snapshot b/.gas-snapshot
@@ -7,9 +7,9 @@ EmergencyExitTest:test_EmergencyExitWithLock() (gas: 391784)
 EmergencyExitTest:test_EmergencyExitWithRewards() (gas: 377316)
 EmergencyExitTest:test_OnlyOwnerCanEnableEmergencyMode() (gas: 39430)
 IntegrationTest:testStakeFoo() (gas: 1178499)
-LeaveTest:test_LeaveShouldProperlyUpdateAccounting() (gas: 2960876)
+LeaveTest:test_LeaveShouldProperlyUpdateAccounting() (gas: 5626814)
 LeaveTest:test_RevertWhenStakeManagerIsTrusted() (gas: 294826)
-LeaveTest:test_TrustNewStakeManager() (gas: 3038518)
+LeaveTest:test_TrustNewStakeManager() (gas: 5699451)
 LockTest:test_LockFailsWithInvalidPeriod() (gas: 309911)
 LockTest:test_LockFailsWithNoStake() (gas: 63708)
 LockTest:test_LockWithoutPriorLock() (gas: 385937)
@@ -61,7 +61,7 @@ UnstakeTest:test_UnstakeOneAccountAndAccruedMP() (gas: 494565)
 UnstakeTest:test_UnstakeOneAccountAndRewards() (gas: 404122)
 UnstakeTest:test_UnstakeOneAccountWithLockUpAndAccruedMP() (gas: 516307)
 UpgradeTest:test_RevertWhenNotOwner() (gas: 2602182)
-UpgradeTest:test_UpgradeStakeManager() (gas: 2875621)
+UpgradeTest:test_UpgradeStakeManager() (gas: 5541546)
 VaultRegistrationTest:test_VaultRegistration() (gas: 62035)
 WithdrawTest:test_CannotWithdrawStakedFunds() (gas: 310550)
 XPNFTTokenTest:testApproveNotAllowed() (gas: 10500)

diff --git a/script/UpgradeRewardsStreamerMP.s.sol b/script/UpgradeRewardsStreamerMP.s.sol
@@ -0,0 +1,22 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.26;
+
+import { UUPSUpgradeable } from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
+import { BaseScript } from "./Base.s.sol";
+import { RewardsStreamerMP } from "../src/RewardsStreamerMP.sol";
+import { IStakeManagerProxy } from "../src/interfaces/IStakeManagerProxy.sol";
+
+contract UpgradeRewardsStreamerMPScript is BaseScript {
+    function run(address admin, IStakeManagerProxy currentImplProxy) public {
+        address deployer = broadcaster;
+        if (admin != address(0)) {
+            deployer = admin;
+        }
+        vm.startBroadcast(deployer);
+        // Replace this with actual new version of the contract
+        address nextImpl = address(new RewardsStreamerMP());
+        bytes memory initializeData;
+        UUPSUpgradeable(address(currentImplProxy)).upgradeToAndCall(nextImpl, initializeData);
+        vm.stopBroadcast();
+    }
+}
diff --git a/test/RewardsStreamerMP.t.sol b/test/RewardsStreamerMP.t.sol
@@ -4,6 +4,7 @@ pragma solidity ^0.8.26;
 import { Test } from "forge-std/Test.sol";
 import { Math } from "@openzeppelin/contracts/utils/math/Math.sol";
 import { DeployRewardsStreamerMPScript } from "../script/DeployRewardsStreamerMP.s.sol";
+import { UpgradeRewardsStreamerMPScript } from "../script/UpgradeRewardsStreamerMP.s.sol";
 import { DeploymentConfig } from "../script/DeploymentConfig.s.sol";
 import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol";
 import { UUPSUpgradeable } from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
@@ -146,6 +147,11 @@ contract RewardsStreamerMPTest is Test {
         uint256 mpPerYear = amount * streamer.MP_RATE_PER_YEAR();
         return maxMP * 365 days / mpPerYear;
     }
+
+    function _upgradeStakeManager() internal {
+        UpgradeRewardsStreamerMPScript upgrade = new UpgradeRewardsStreamerMPScript();
+        upgrade.run(admin, IStakeManagerProxy(address(streamer)));
+    }
 }
 
 contract VaultRegistrationTest is RewardsStreamerMPTest {
@@ -1758,13 +1764,6 @@ contract EmergencyExitTest is RewardsStreamerMPTest {
 }
 
 contract UpgradeTest is RewardsStreamerMPTest {
-    function _upgradeStakeManager() internal {
-        address newImpl = address(new RewardsStreamerMP());
-        bytes memory initializeData;
-        vm.prank(admin);
-        UUPSUpgradeable(streamer).upgradeToAndCall(newImpl, initializeData);
-    }
-
     function setUp() public override {
         super.setUp();
     }
@@ -1811,13 +1810,6 @@ contract UpgradeTest is RewardsStreamerMPTest {
 }
 
 contract LeaveTest is RewardsStreamerMPTest {
-    function _upgradeStakeManager() internal {
-        address newImpl = address(new RewardsStreamerMP());
-        bytes memory initializeData;
-        vm.prank(admin);
-        UUPSUpgradeable(streamer).upgradeToAndCall(newImpl, initializeData);
-    }
-
     function setUp() public override {
         super.setUp();
     }
@@ -1910,13 +1902,6 @@ contract LeaveTest is RewardsStreamerMPTest {
 }
 
 contract MaliciousUpgradeTest is RewardsStreamerMPTest {
-    function _upgradeStakeManager() internal {
-        address newImpl = address(new RewardsStreamerMP());
-        bytes memory initializeData;
-        vm.prank(admin);
-        UUPSUpgradeable(streamer).upgradeToAndCall(newImpl, initializeData);
-    }
-
     function setUp() public override {
         super.setUp();
     }