Fix test

rocket-pool · Mar 13, 2024 · 8e92558 · 8e92558
1 parent cf1c9fa
commit 8e92558
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/test/minipool/scenario-scrub.js b/test/minipool/scenario-scrub.js
@@ -2,20 +2,21 @@
 import {
     RocketDAONodeTrusted,
     RocketDAONodeTrustedSettingsMinipool, RocketDAOProtocolSettingsNode, RocketNetworkPrices,
-    RocketNodeStaking,
+    RocketNodeStaking, RocketNodeStakingNew,
     RocketTokenRPL,
-    RocketVault
+    RocketVault,
 } from '../_utils/artifacts';
 import { assertBN } from '../_helpers/bn';
 import { minipoolStates } from '../_helpers/minipool';
+import { upgradeExecuted } from '../_utils/upgrade';
 
 
 export async function voteScrub(minipool, txOptions) {
     // Get minipool owner
     const nodeAddress = await minipool.getNodeAddress.call();
 
     // Get contracts
-    const rocketNodeStaking = await RocketNodeStaking.deployed();
+    const rocketNodeStaking = (await upgradeExecuted()) ? await RocketNodeStakingNew.deployed() : await RocketNodeStaking.deployed();
     const rocketVault = await RocketVault.deployed();
     const rocketTokenRPL = await RocketTokenRPL.deployed();
     const rocketDAONodeTrustedSettingsMinipool = await RocketDAONodeTrustedSettingsMinipool.deployed();