From 381d32d80d6b98ffabd2dd724694636f66ff2801 Mon Sep 17 00:00:00 2001
From: Alexandr Tarelkin <alexandrtarelkin92@gmail.com>
Date: Sun, 5 Feb 2023 20:19:21 +0300
Subject: [PATCH] fix test setup

---
 test/helpers/factories.js | 23 +++++++++++++++--------
 test/helpers/protocol.js  | 23 +++++++++++++++++++----
 2 files changed, 34 insertions(+), 12 deletions(-)

diff --git a/test/helpers/factories.js b/test/helpers/factories.js
index 325653562..4ef7a306c 100644
--- a/test/helpers/factories.js
+++ b/test/helpers/factories.js
@@ -181,8 +181,8 @@ async function hashConsensusTimeTravellableFactory({ legacyOracle, voting, repor
   return consensus
 }
 
-async function accountingOracleFactory({ voting, pool, consensusContract, legacyOracle }) {
-  const base = await AccountingOracle.new(pool.address, SECONDS_PER_SLOT, GENESIS_TIME)
+async function accountingOracleFactory({ voting, pool, lidoLocator, consensusContract, legacyOracle }) {
+  const base = await AccountingOracle.new(lidoLocator.address, SECONDS_PER_SLOT, GENESIS_TIME)
   const proxy = await OssifiableProxy.new(base.address, voting.address, '0x')
   const oracle = await AccountingOracle.at(proxy.address)
 
@@ -257,13 +257,14 @@ async function elRewardsVaultFactory({ pool, treasury }) {
   return await LidoExecutionLayerRewardsVault.new(pool.address, treasury.address)
 }
 
-async function withdrawalQueueFactory({ appManager, wsteth }) {
+async function withdrawalQueueFactory({ appManager, oracle, wsteth }) {
   const withdrawalQueue = (await withdrawals.deploy(appManager.address, wsteth.address)).queue
 
   await withdrawalQueue.initialize(appManager.address, appManager.address, appManager.address, appManager.address)
 
   const BUNKER_MODE_REPORT_ROLE = await withdrawalQueue.BUNKER_MODE_REPORT_ROLE()
   await withdrawalQueue.grantRole(BUNKER_MODE_REPORT_ROLE, appManager.address, { from: appManager.address })
+  await withdrawalQueue.grantRole(BUNKER_MODE_REPORT_ROLE, oracle.address, { from: appManager.address })
 
   return withdrawalQueue
 }
@@ -301,19 +302,24 @@ async function selfOwnedStETHBurnerFactory({ appManager, treasury, pool, voting
   return burner
 }
 
-async function lidoLocatorMockFactory(protocol) {
+async function lidoLocatorFactory(protocol) {
+  const base = await lidoLocatorMockImplFactory(protocol)
+  return await OssifiableProxy.new(base.address, protocol.appManager.address, '0x')
+}
+
+async function lidoLocatorMockImplFactory(protocol) {
   return LidoLocatorMock.new({
     lido: protocol.pool.address,
     depositSecurityModule: protocol.depositSecurityModule.address,
     elRewardsVault: protocol.elRewardsVault.address,
-    accountingOracle: protocol.oracle.address,
+    accountingOracle: protocol.oracle ? protocol.oracle.address : ZERO_ADDRESS,
     legacyOracle: protocol.legacyOracle.address,
     safetyNetsRegistry: ZERO_ADDRESS,
     selfOwnedStEthBurner: protocol.selfOwnedStETHBurner.address,
     validatorExitBus: ZERO_ADDRESS,
     stakingRouter: protocol.stakingRouter.address,
     treasury: protocol.treasury.address,
-    withdrawalQueue: protocol.withdrawalQueue.address,
+    withdrawalQueue: protocol.withdrawalQueue ? protocol.withdrawalQueue.address : ZERO_ADDRESS,
     withdrawalVault: protocol.withdrawalVault.address,
     postTokenRebaseReceiver: protocol.legacyOracle.address
   })
@@ -347,12 +353,13 @@ module.exports = {
   withdrawalCredentialsFactory,
   stakingModulesFactory,
   guardiansFactory,
-  lidoLocatorMockFactory,
+  lidoLocatorMockImplFactory,
   selfOwnedStETHBurnerFactory,
   postSetup,
   legacyOracleFactory,
   legacyOracleMockFactory,
   hashConsensusFactory,
   hashConsensusTimeTravellableFactory,
-  reportProcessorFactory
+  reportProcessorFactory,
+  lidoLocatorFactory
 }
diff --git a/test/helpers/protocol.js b/test/helpers/protocol.js
index 7058b683b..1b32eb1be 100644
--- a/test/helpers/protocol.js
+++ b/test/helpers/protocol.js
@@ -2,6 +2,8 @@ const { ethers } = require('hardhat')
 
 const { newDao } = require('./dao')
 
+const OssifiableProxy = artifacts.require('OssifiableProxy')
+
 const factories = require('./factories')
 const DEAFAULT_FACTORIES = {
   appManagerFactory: factories.appManagerFactory,
@@ -23,9 +25,10 @@ const DEAFAULT_FACTORIES = {
   withdrawalCredentialsFactory: factories.withdrawalCredentialsFactory,
   stakingModulesFactory: factories.stakingModulesFactory,
   guardiansFactory: factories.guardiansFactory,
-  lidoLocatorFactory: factories.lidoLocatorMockFactory,
+  lidoLocatorMockImplFactory: factories.lidoLocatorMockImplFactory,
   selfOwnedStETHBurnerFactory: factories.selfOwnedStETHBurnerFactory,
-  postSetup: factories.postSetup
+  postSetup: factories.postSetup,
+  lidoLocatorFactory: factories.lidoLocatorFactory
 }
 
 const getFactory = (config, factoryName) => {
@@ -54,7 +57,6 @@ async function deployProtocol(config = {}) {
 
   protocol.reportProcessor = await getFactory(config, 'reportProcessorFactory')(protocol)
   protocol.consensusContract = await getFactory(config, 'hashConsensusFactory')(protocol)
-  protocol.oracle = await getFactory(config, 'accountingOracleFactory')(protocol)
 
   protocol.depositContract = await getFactory(config, 'depositContractFactory')(protocol)
 
@@ -65,11 +67,18 @@ async function deployProtocol(config = {}) {
 
   protocol.depositSecurityModule = await getFactory(config, 'depositSecurityModuleFactory')(protocol)
   protocol.elRewardsVault = await getFactory(config, 'elRewardsVaultFactory')(protocol)
-  protocol.withdrawalQueue = await getFactory(config, 'withdrawalQueueFactory')(protocol)
   protocol.withdrawalVault = await getFactory(config, 'withdrawalVaultFactory')(protocol)
   protocol.eip712StETH = await getFactory(config, 'eip712StETHFactory')(protocol)
   protocol.selfOwnedStETHBurner = await getFactory(config, 'selfOwnedStETHBurnerFactory')(protocol)
+
   protocol.lidoLocator = await getFactory(config, 'lidoLocatorFactory')(protocol)
+  protocol.oracle = await getFactory(config, 'accountingOracleFactory')(protocol)
+  protocol.withdrawalQueue = await getFactory(config, 'withdrawalQueueFactory')(protocol)
+  await upgradeOssifiableProxy(
+    protocol.lidoLocator.address,
+    (await getFactory(config, 'lidoLocatorMockImplFactory')(protocol)).address,
+    protocol.appManager.address
+  )
 
   await getFactory(config, 'postSetup')(protocol)
 
@@ -93,6 +102,12 @@ async function addStakingModules(stakingModulesFactory, protocol) {
   return stakingModules.map(({ module }) => module)
 }
 
+async function upgradeOssifiableProxy(proxyAddress, newImplemantation, proxyOwner) {
+  const proxy = await OssifiableProxy.at(proxyAddress)
+
+  await proxy.proxy__upgradeTo(newImplemantation, { from: proxyOwner })
+}
+
 module.exports = {
   deployProtocol
 }