fix: tests after rename

test/helpers/deployer.js

@@ -52,7 +52,7 @@ class Deployer {
       ])
     )
 
-    this.stakeManager = await contractFactories.StakeManager.attach(stakeManagerProxy.address)
+    this.stakeManager = contractFactories.StakeManager.attach(stakeManagerProxy.address)
     // TODO cannot alter functions like we used to here, replace usage with actual impl like below
     // this.buildStakeManagerObject(this.stakeManager, this.governance)
     await this.governance.update(
@@ -84,7 +84,7 @@ class Deployer {
     return _contracts
   }
 
-  async deployStakeManager(wallets, legacy = false) {
+  async deployStakeManager(wallets, pol = false) {
     this.governance = await this.deployGovernance()
     this.registry = await contractFactories.Registry.deploy(this.governance.address)
 
@@ -93,15 +93,7 @@ class Deployer {
     this.validatorShare = await contractFactories.ValidatorShare.deploy()
     this.rootChain = await this.deployRootChain()
     this.stakingInfo = await contractFactories.StakingInfo.deploy(this.registry.address)
-
-    let token 
-    if (legacy){
-      this.legacyToken = await contractFactories.TestToken.deploy('Legacy Token', 'LEG')
-      token = this.legacyToken
-    } else {
-      this.stakeToken = await contractFactories.TestToken.deploy('Stake Token', 'STAKE')
-      token = this.stakeToken
-    }
+    this.stakeToken = await contractFactories.TestToken.deploy('Stake Token', 'STAKE')
 
     this.stakingNFT = await contractFactories.StakingNFT.deploy('Matic Validator', 'MV')
 
@@ -114,7 +106,7 @@ class Deployer {
       stakeManager.interface.encodeFunctionData('initialize', [
         this.registry.address,
         rootChainOwner.getAddressString(),
-        token.address,
+        this.stakeToken.address,
         this.stakingNFT.address,
         this.stakingInfo.address,
         this.validatorShareFactory.address,
@@ -124,40 +116,43 @@ class Deployer {
       ])
     )
 
-    this.stakeManager = await contractFactories.StakeManagerTestable.attach(proxy.address)
+    this.stakeManager = contractFactories.StakeManagerTestable.attach(proxy.address)
     this.slashingManager = await contractFactories.SlashingManager.deploy(
       this.registry.address,
       this.stakingInfo.address,
       'heimdall-P5rXwg'
     )
 
-    if (legacy){
-      await token.mint(this.stakeManager.address, web3.utils.toWei('10000000'))
+    await this.stakingNFT.transferOwnership(this.stakeManager.address)
+    await this.updateContractMap(ethUtils.keccak256('stakeManager'), this.stakeManager.address)
+    await this.updateContractMap(ethUtils.keccak256('validatorShare'), this.validatorShare.address)
+    await this.updateContractMap(ethUtils.keccak256('slashingManager'), this.slashingManager.address)
+
+    await this.stakeToken.mint(this.stakeManager.address, web3.utils.toWei('10000000'))
 
-      this.stakeToken = await contractFactories.TestToken.deploy('POL', 'POL')
+    if (pol){
+      this.polToken = await contractFactories.ERC20Permit.deploy('POL', 'POL', '1.1.0')
 
-      this.migration = await contractFactories.PolygonMigration.deploy(this.legacyToken.address, this.stakeToken.address)
+      this.migration = await contractFactories.PolygonMigration.deploy(this.stakeToken.address, this.polToken.address)
 
       await this.stakeToken.mint(this.migration.address, web3.utils.toWei('50000000'))
-      await this.legacyToken.mint(this.migration.address, web3.utils.toWei('50000000'))
+      await this.polToken.mint(this.migration.address, web3.utils.toWei('50000000'))
 
       await this.governance.update(
           this.stakeManager.address,
-          this.stakeManager.interface.encodeFunctionData('initializeLegacy', [this.stakeToken.address, this.migration.address])
+          this.stakeManager.interface.encodeFunctionData('initializePOL', [this.polToken.address, this.migration.address])
         )
-    }
-
-    await this.stakingNFT.transferOwnership(this.stakeManager.address)
-    await this.updateContractMap(ethUtils.keccak256('stakeManager'), this.stakeManager.address)
-    await this.updateContractMap(ethUtils.keccak256('validatorShare'), this.validatorShare.address)
-    await this.updateContractMap(ethUtils.keccak256('slashingManager'), this.slashingManager.address)
+
+      await this.updateContractMap(ethUtils.keccak256('pol'), this.polToken.address)
+    } 
+
     let _contracts = {
       rootChainOwner: rootChainOwner,
       registry: this.registry,
       rootChain: this.rootChain,
       stakeManager: this.stakeManager,
       stakeToken: this.stakeToken,
-      legacyToken: this.legacyToken,
+      polToken: this.polToken,
       slashingManager: this.slashingManager,
       stakingInfo: this.stakingInfo,
       governance: this.governance,

test/units/root/RootChain.test.js

@@ -44,15 +44,11 @@ describe('RootChain', async function (accounts) {
   })
 
   async function freshDeploy() {
-    const contracts = await deployer.deployStakeManager(wallets)
+    const contracts = await deployer.deployStakeManager(wallets, true)
     governance = contracts.governance
     rootChain = contracts.rootChain
     stakeManager = contracts.stakeManager
-    stakeToken = await TestToken.deploy('Stake Token', 'STAKE')
-    await contracts.governance.update(
-      stakeManager.address,
-      stakeManager.interface.encodeFunctionData('setStakingToken', [stakeToken.address])
-    )
+    stakeToken = contracts.stakeToken
     await stakeManager.changeRootChain(rootChain.address)
     await contracts.governance.update(
       stakeManager.address,

test/units/staking/DrainStakeManager.test.js

@@ -2,6 +2,7 @@
 
 import chai, { assert } from 'chai'
 import chaiAsPromised from 'chai-as-promised'
+import { BN } from '@openzeppelin/test-helpers'
 
 import { DrainStakeManager } from '../../helpers/artifacts.js'
 
@@ -12,21 +13,25 @@ import { generateFirstWallets, mnemonics } from '../../helpers/wallets.js'
 
 chai.use(chaiAsPromised).should()
 
-describe('DrainStakeManager', async function (accounts) {
-  accounts = await ethers.getSigners()
-  accounts = accounts.map((account) => {
-    return account.address
+// todo: is this even needed anymore?
+describe('DrainStakeManager', function (accounts) {
+  let owner
+  before(async function() {
+    accounts = await ethers.getSigners()
+    accounts = accounts.map((account) => {
+      return account.address
+    })
+    owner = accounts[0]
   })
-
-  const owner = accounts[0]
-  describe('Upgrade and drain staking contract', async function () {
+
+  describe('Upgrade and drain staking contract', function () {
     before(async function () {
       this.wallets = generateFirstWallets(mnemonics, 10)
 
-      let contracts = await deployer.deployStakeManager(this.wallets)
+      let contracts = await deployer.deployStakeManager(this.wallets, true)
 
       this.governance = contracts.governance
-      this.stakeToken = contracts.stakeToken
+      this.polToken = contracts.polToken
       this.stakeManager = contracts.stakeManager
       this.proxy = contracts.stakeManagerProxy
       this.stakeManagerImpl = contracts.stakeManagerImpl
@@ -36,9 +41,10 @@ describe('DrainStakeManager', async function (accounts) {
     })
 
     it('must have some tokens', async function () {
+      const initalAmount = web3.utils.toWei('10000000')
       const amount = web3.utils.toWei('90000')
-      await this.stakeToken.mint(this.stakeManager.address, amount)
-      ;(await this.stakeToken.balanceOf(this.stakeManager.address)).toString().should.be.equal(amount.toString())
+      await this.polToken.mint(this.stakeManager.address, amount)
+      ;(await this.polToken.balanceOf(this.stakeManager.address)).toString().should.be.equal(BN(amount).add(BN(initalAmount)).toString())
     })
 
     it('must lock stake manager', async function () {
@@ -60,7 +66,7 @@ describe('DrainStakeManager', async function (accounts) {
     })
 
     it('must fail draining when not drained owner', async function () {
-      const balance = await this.stakeToken.balanceOf(this.stakeManager.address)
+      const balance = await this.polToken.balanceOf(this.stakeManager.address)
       try {
         await this.stakeManagerDrainable.drain(owner, balance)
         assert.fail('Funds should not be drained')
@@ -70,10 +76,10 @@ describe('DrainStakeManager', async function (accounts) {
     })
 
     it('must drain all funds when drained by owner (Gnosis safe)', async function () {
-      const balance = await this.stakeToken.balanceOf(this.stakeManager.address)
+      const balance = await this.polToken.balanceOf(this.stakeManager.address)
       const data = this.stakeManagerDrainable.interface.encodeFunctionData('drain', [owner, balance.toString()])
       await execSafe(this.gSafe, this.stakeManager.address, data, [accounts[1], accounts[2]])
-      assert.equal((await this.stakeToken.balanceOf(this.stakeManager.address)).toString(), '0')
+      assert.equal((await this.polToken.balanceOf(this.stakeManager.address)).toString(), '0')
     })
 
     it('must swap back to normal implementaion', async function () {