From 38fe7f25d99e012ac70ae2445d2cc480e7d56a9a Mon Sep 17 00:00:00 2001
From: Joey Santoro <joseph.santoro@duke.edu>
Date: Sun, 24 Oct 2021 18:34:26 -0700
Subject: [PATCH] cleanup

---
 proposals/dao/fip_34.ts                       |  4 ++-
 scripts/utils/checkProposal.ts                | 34 +++++++++++--------
 scripts/utils/exec.ts                         |  2 +-
 .../feirari/RewardsDistributorAdmin.test.ts   | 10 ++----
 4 files changed, 25 insertions(+), 25 deletions(-)

diff --git a/proposals/dao/fip_34.ts b/proposals/dao/fip_34.ts
index f60e91e1e..91a7a301d 100644
--- a/proposals/dao/fip_34.ts
+++ b/proposals/dao/fip_34.ts
@@ -9,6 +9,7 @@ import {
   ValidateUpgradeFunc
 } from '../../types/types';
 import { FeiDAOTimelock } from '@custom-types/contracts';
+import { getImpersonatedSigner } from '@test/helpers';
 
 chai.use(CBN(ethers.BigNumber));
 
@@ -51,7 +52,8 @@ export const deploy: DeployUpgradeFunc = async (deployAddress, addresses, loggin
 };
 
 export const setup: SetupUpgradeFunc = async (addresses, oldContracts, contracts, logging) => {
-  await (await (contracts.feiDAOTimelock as FeiDAOTimelock).rollback()).wait();
+  const timelock: FeiDAOTimelock = contracts.feiDAOTimelock as FeiDAOTimelock;
+  await (await timelock.connect(await getImpersonatedSigner(addresses.multisig)).rollback()).wait();
 };
 
 export const teardown: TeardownUpgradeFunc = async (addresses, oldContracts, contracts, logging) => {
diff --git a/scripts/utils/checkProposal.ts b/scripts/utils/checkProposal.ts
index b39f7fe89..8f4cafdaf 100644
--- a/scripts/utils/checkProposal.ts
+++ b/scripts/utils/checkProposal.ts
@@ -1,7 +1,6 @@
-import { getAllContracts } from '../../test/integration/setup/loadContracts';
-import hre, { ethers } from 'hardhat';
-import { time } from '@openzeppelin/test-helpers';
-import { NamedContracts, namedContractsToNamedAddresses, UpgradeFuncs } from '../../types/types';
+import { getAllContracts, getAllContractAddresses } from '@test/integration/setup/loadContracts';
+import { getImpersonatedSigner, time } from '@test/helpers';
+import { NamedContracts, UpgradeFuncs } from '@custom-types/types';
 
 import * as dotenv from 'dotenv';
 
@@ -22,16 +21,26 @@ async function checkProposal() {
     throw new Error('DEPLOY_FILE or PROPOSAL_NUMBER env variable not set');
   }
 
+  // Get the upgrade setup, run and teardown scripts
+  const proposalFuncs: UpgradeFuncs = await import(`@proposals/dao/${proposalName}`);
+
   const contracts = (await getAllContracts()) as NamedContracts;
 
-  const { feiDAO } = contracts;
+  const contractAddresses = await getAllContractAddresses();
 
-  await hre.network.provider.request({
-    method: 'hardhat_impersonateAccount',
-    params: [voterAddress]
-  });
+  if (process.env.DO_SETUP) {
+    console.log('Setup');
+    await proposalFuncs.setup(
+      contractAddresses,
+      contracts as unknown as NamedContracts,
+      contracts as unknown as NamedContracts,
+      true
+    );
+  }
 
-  const voterSigner = await ethers.getSigner(voterAddress);
+  const { feiDAO } = contracts;
+
+  const voterSigner = await getImpersonatedSigner(voterAddress);
 
   console.log(`Proposal Number: ${proposalNo}`);
 
@@ -76,11 +85,6 @@ async function checkProposal() {
   await feiDAO['execute(uint256)'](proposalNo);
   console.log('Success');
 
-  // Get the upgrade setup, run and teardown scripts
-  const proposalFuncs: UpgradeFuncs = await import(`../../proposals/dao/${proposalName}`);
-
-  const contractAddresses = namedContractsToNamedAddresses(contracts);
-
   console.log('Teardown');
   await proposalFuncs.teardown(
     contractAddresses,
diff --git a/scripts/utils/exec.ts b/scripts/utils/exec.ts
index 15e70d2d5..87fabdba7 100644
--- a/scripts/utils/exec.ts
+++ b/scripts/utils/exec.ts
@@ -1,5 +1,5 @@
 import hre, { ethers } from 'hardhat';
-import { time } from '@openzeppelin/test-helpers';
+import { time } from '@test/helpers';
 
 import * as dotenv from 'dotenv';
 
diff --git a/test/unit/staking/feirari/RewardsDistributorAdmin.test.ts b/test/unit/staking/feirari/RewardsDistributorAdmin.test.ts
index 50ad2d620..3da5e4e95 100644
--- a/test/unit/staking/feirari/RewardsDistributorAdmin.test.ts
+++ b/test/unit/staking/feirari/RewardsDistributorAdmin.test.ts
@@ -1,17 +1,11 @@
-import { expectRevert, getAddresses, getCore } from '../../../helpers';
+import { expectRevert, getAddresses, getCore, ZERO_ADDRESS } from '@test/helpers';
 import { expect } from 'chai';
 import hre, { ethers } from 'hardhat';
 import { Signer, utils } from 'ethers';
-import testHelpers from '@openzeppelin/test-helpers';
-import { Core } from '../../../../types/contracts/Core';
-import { RewardsDistributorAdmin } from '../../../../types/contracts/RewardsDistributorAdmin';
-import { MockRewardsDistributor } from '../../../../types/contracts/MockRewardsDistributor';
+import { Core, RewardsDistributorAdmin, MockRewardsDistributor } from '@custom-types/contracts';
 import { keccak256 } from 'ethers/lib/utils';
 
 const toBN = ethers.BigNumber.from;
-const {
-  constants: { ZERO_ADDRESS }
-} = testHelpers;
 
 describe('RewardsDistributorAdmin', function () {
   let governorAddress: string;