From 2a566133ea02166655e067bde75abef942bfd342 Mon Sep 17 00:00:00 2001
From: MiksuJak <mikolaj.jakubowski@protonmail.com>
Date: Thu, 23 May 2024 12:07:19 +0200
Subject: [PATCH] Adjust integration tests

---
 test/integration/governance.test.ts |  2 +-
 test/integration/suite.ts           | 22 ++++++++++++----------
 test/integration/truefi.test.ts     |  9 +++++++++
 3 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/test/integration/governance.test.ts b/test/integration/governance.test.ts
index cc1d5f281..b26cbad81 100644
--- a/test/integration/governance.test.ts
+++ b/test/integration/governance.test.ts
@@ -50,6 +50,6 @@ describe('Governance', () => {
       'stkTRU',
       'guardian',
       'proposalCount',
-    ], '0xf6E2Da7D82ee49f76CE652bc0BeB546Cbe0Ea521')
+    ], undefined, '0xf6E2Da7D82ee49f76CE652bc0BeB546Cbe0Ea521')
   })
 })
diff --git a/test/integration/suite.ts b/test/integration/suite.ts
index 115532570..e72ce585d 100644
--- a/test/integration/suite.ts
+++ b/test/integration/suite.ts
@@ -37,22 +37,23 @@ const execGetter = <T extends Contract>(contract: T) => async (getter: Getter<T>
 export const TEST_STATE_BLOCK_NUMBER = 12010725
 
 export function upgradeSuite<T extends Contract>(blockNumber: number, Factory: ContractFactoryConstructor<T>, currentAddress: string,
-  getters: Getter<T>[], contractsOwner?: string): Promise<T>
+  getters: Getter<T>[], newGetters?: Getter<T>[] | undefined, contractsOwner?: string): Promise<T>
 export function upgradeSuite<T extends Contract>(Factory: ContractFactoryConstructor<T>, currentAddress: string,
-  getters: Getter<T>[], contractsOwner?: string): Promise<T>
+  getters: Getter<T>[], newGetters?: Getter<T>[] | undefined, contractsOwner?: string): Promise<T>
 export function upgradeSuite (...args: any[]): any {
   if (typeof args[0] === 'number') {
-    const [bn, factory, address, getters, owner] = args
-    return _upgradeSuite(factory, address, getters, owner, bn)
+    const [bn, factory, address, getters, newGetters, owner] = args
+    return _upgradeSuite(factory, address, getters, newGetters, owner, bn)
   }
-  const [factory, address, getters, owner] = args
-  return _upgradeSuite(factory, address, getters, owner)
+  const [factory, address, getters, newGetters, owner] = args
+  return _upgradeSuite(factory, address, getters, newGetters, owner)
 }
 
 async function _upgradeSuite<T extends Contract> (
   Factory: ContractFactoryConstructor<T>,
   currentAddress: string,
   getters: Getter<T>[],
+  newGetters: Getter<T>[] | undefined,
   contractsOwner: string = CONTRACTS_OWNER,
   blockNumber?: number | undefined,
 ) {
@@ -65,12 +66,13 @@ async function _upgradeSuite<T extends Contract> (
   const oldValues = await Promise.all(getters.map(execGetter(existingContract)))
   const proxy = new OwnedUpgradeabilityProxy__factory(owner).attach(currentAddress)
   await (await proxy.upgradeTo(newContract.address)).wait()
-  const newValues = await Promise.all(getters.map(execGetter(existingContract)))
+  const postUpgradeGetters = newGetters || getters
+  const newValues = await Promise.all(postUpgradeGetters.map(execGetter(existingContract)))
   for (let i = 0; i < oldValues.length; i++) {
     expect(oldValues[i], `Possible corrupted storage:
-Getter: ${getters[i]}
-Current: ${oldValues[i].toString()}
-Post upgrade: ${newValues[i].toString()} \n`).to.deep.equal(newValues[i])
+      Getter: ${getters[i]}
+      Current: ${oldValues[i].toString()}
+      Post upgrade: ${newValues[i].toString()} \n`).to.deep.equal(newValues[i])
   }
   return existingContract
 }
diff --git a/test/integration/truefi.test.ts b/test/integration/truefi.test.ts
index 48862800e..990b8f748 100644
--- a/test/integration/truefi.test.ts
+++ b/test/integration/truefi.test.ts
@@ -120,7 +120,16 @@ describe('TrueFi', () => {
   })
 
   it('LinearDistributor', async () => {
+    // @ts-ignore
     await upgradeSuite(TEST_STATE_BLOCK_NUMBER, LinearTrueDistributor__factory, '0xfB8d918428373f766B352564b70d1DcC1e3b6383', [
+      'trustToken',
+      'distributionStart',
+      'duration',
+      'totalAmount',
+      'lastDistribution',
+      'distributed',
+      'farm',
+    ], [
       'asset',
       'distributionStart',
       'duration',