Batch GovernorStorage transactions setup

Amxx · Aug 1, 2023 · b2445f8 · b2445f8
1 parent bdddeda
commit b2445f8
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 14 deletions.
diff --git a/test/governance/extensions/GovernorStorage.test.js b/test/governance/extensions/GovernorStorage.test.js
@@ -4,6 +4,7 @@ const { expect } = require('chai');
 const { expectRevertCustomError } = require('../../helpers/customError');
 const Enums = require('../../helpers/enums');
 const { GovernorHelper, timelockSalt } = require('../../helpers/governance');
+const { mineWith } = require('../../helpers/txpool');
 
 const Timelock = artifacts.require('TimelockController');
 const Governor = artifacts.require('$GovernorStorageMock');
@@ -53,19 +54,20 @@ contract('GovernorStorage', function (accounts) {
 
         await web3.eth.sendTransaction({ from: owner, to: this.timelock.address, value });
 
-        // normal setup: governor is proposer, everyone is executor, timelock is its own admin
-        await this.timelock.grantRole(PROPOSER_ROLE, this.mock.address);
-        await this.timelock.grantRole(PROPOSER_ROLE, owner);
-        await this.timelock.grantRole(CANCELLER_ROLE, this.mock.address);
-        await this.timelock.grantRole(CANCELLER_ROLE, owner);
-        await this.timelock.grantRole(EXECUTOR_ROLE, constants.ZERO_ADDRESS);
-        await this.timelock.revokeRole(DEFAULT_ADMIN_ROLE, deployer);
-
-        await this.token.$_mint(owner, tokenSupply);
-        await this.helper.delegate({ token: this.token, to: voter1, value: web3.utils.toWei('10') }, { from: owner });
-        await this.helper.delegate({ token: this.token, to: voter2, value: web3.utils.toWei('7') }, { from: owner });
-        await this.helper.delegate({ token: this.token, to: voter3, value: web3.utils.toWei('5') }, { from: owner });
-        await this.helper.delegate({ token: this.token, to: voter4, value: web3.utils.toWei('2') }, { from: owner });
+        await mineWith(async () => {
+          this.timelock.grantRole(PROPOSER_ROLE, this.mock.address);
+          this.timelock.grantRole(PROPOSER_ROLE, owner);
+          this.timelock.grantRole(CANCELLER_ROLE, this.mock.address);
+          this.timelock.grantRole(CANCELLER_ROLE, owner);
+          this.timelock.grantRole(EXECUTOR_ROLE, constants.ZERO_ADDRESS);
+          this.timelock.revokeRole(DEFAULT_ADMIN_ROLE, deployer);
+
+          this.token.$_mint(owner, tokenSupply);
+          this.helper.delegate({ token: this.token, to: voter1, value: web3.utils.toWei('10') }, { from: owner });
+          this.helper.delegate({ token: this.token, to: voter2, value: web3.utils.toWei('7') }, { from: owner });
+          this.helper.delegate({ token: this.token, to: voter3, value: web3.utils.toWei('5') }, { from: owner });
+          this.helper.delegate({ token: this.token, to: voter4, value: web3.utils.toWei('2') }, { from: owner });
+        });
 
         // default proposal
         this.proposal = this.helper.setProposal(

diff --git a/test/helpers/txpool.js b/test/helpers/txpool.js
@@ -1,3 +1,4 @@
+const { mine } = require('@nomicfoundation/hardhat-network-helpers');
 const { network } = require('hardhat');
 const { promisify } = require('util');
 
@@ -18,7 +19,7 @@ async function batchInBlock(txs) {
       await queue();
     }
     // mine one block
-    await network.provider.send('evm_mine');
+    await mine();
     // fetch receipts
     const receipts = await Promise.all(promises);
     // Sanity check, all tx should be in the same block
@@ -32,7 +33,18 @@ async function batchInBlock(txs) {
   }
 }
 
+async function mineWith(fn) {
+  if (typeof fnOrTxs !== 'function') {
+    throw new Error('Argument to mine with is not a function');
+  }
+  await network.provider.send('evm_setAutomine', [false]);
+  await fn();
+  await network.provider.send('evm_setAutomine', [true]);
+  await mine();
+}
+
 module.exports = {
   countPendingTransactions,
   batchInBlock,
+  mineWith,
 };