Fix failing test and correctly handle legacy RPL stake

rocket-pool · Mar 13, 2024 · 786ebfb · 786ebfb
1 parent 8e92558
commit 786ebfb
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/contracts/contract/node/RocketNodeStaking.sol b/contracts/contract/node/RocketNodeStaking.sol
@@ -87,6 +87,9 @@ contract RocketNodeStaking is RocketBase, RocketNodeStakingInterface {
         bytes32 key = keccak256(abi.encodePacked("rpl.staked.node.amount", _nodeAddress));
         RocketNetworkSnapshotsInterface rocketNetworkSnapshots = RocketNetworkSnapshotsInterface(getContractAddress("rocketNetworkSnapshots"));
         uint224 value = rocketNetworkSnapshots.latestValue(key);
+        if (value == 0){
+            value = uint224(getUint(key));
+        }
         rocketNetworkSnapshots.push(key, uint32(block.number), value + uint224(_amount));
     }
 
@@ -96,6 +99,9 @@ contract RocketNodeStaking is RocketBase, RocketNodeStakingInterface {
         bytes32 key = keccak256(abi.encodePacked("rpl.staked.node.amount", _nodeAddress));
         RocketNetworkSnapshotsInterface rocketNetworkSnapshots = RocketNetworkSnapshotsInterface(getContractAddress("rocketNetworkSnapshots"));
         uint224 value = rocketNetworkSnapshots.latestValue(key);
+        if (value == 0){
+            value = uint224(getUint(key));
+        }
         rocketNetworkSnapshots.push(key, uint32(block.number), value - uint224(_amount));
     }
 

diff --git a/test/rewards/scenario-claim-and-stake-rewards.js b/test/rewards/scenario-claim-and-stake-rewards.js
@@ -1,16 +1,19 @@
 import {
     RocketMerkleDistributorMainnet,
-    RocketNodeManager, RocketNodeStaking,
+    RocketNodeManager, RocketNodeStaking, RocketNodeStakingNew,
     RocketRewardsPool,
-    RocketStorage, RocketTokenRPL
+    RocketStorage, RocketTokenRPL,
 } from '../_utils/artifacts';
 import { parseRewardsMap } from '../_utils/merkle-tree';
 import { assertBN } from '../_helpers/bn';
+import { upgradeExecuted } from '../_utils/upgrade';
 
 
 // Submit network prices
 export async function claimAndStakeRewards(nodeAddress, indices, rewards, stakeAmount, txOptions) {
 
+    const upgraded = await upgradeExecuted();
+
     // Load contracts
     const [
         rocketRewardsPool,
@@ -22,7 +25,7 @@ export async function claimAndStakeRewards(nodeAddress, indices, rewards, stakeA
     ] = await Promise.all([
         RocketRewardsPool.deployed(),
         RocketNodeManager.deployed(),
-        RocketNodeStaking.deployed(),
+        upgraded ? RocketNodeStakingNew.deployed() : RocketNodeStaking.deployed(),
         RocketMerkleDistributorMainnet.deployed(),
         RocketStorage.deployed(),
         RocketTokenRPL.deployed(),