Add forkSetup to test helper

balancer · Sep 12, 2023 · c0edb27 · c0edb27
1 parent c2c8030
commit c0edb27
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 40 deletions.
diff --git a/test/lib/utils/helper.ts b/test/lib/utils/helper.ts
@@ -163,6 +163,14 @@ export const setTokenBalance = async (
     });
 };
 
+/**
+ * Find ERC20 token balance storage slot (to be used on setTokenBalance)
+ *
+ * @param client client that will perform contract calls
+ * @param accountAddress Account address to probe storage slot changes
+ * @param tokenAddress Token address which we're looking for the balance slot
+ * @param isVyperMapping Whether the storage uses Vyper or Solidity mapping
+ */
 export async function findTokenBalanceSlot(
     client: Client & PublicActions & TestActions,
     accountAddress: Address,
@@ -216,3 +224,65 @@ export async function findTokenBalanceSlot(
     }
     throw new Error('Balance slot not found!');
 }
+
+/**
+ * Setup local fork with approved token balance for a given account address
+ *
+ * @param client Client that will perform transactions
+ * @param accountAddress Account address that will have token balance set and approved
+ * @param tokens Token addresses which balance will be set and approved
+ * @param slots Slot that stores token balance in memory - use npm package `slot20` to identify which slot to provide
+ * @param balances Balances in EVM amounts
+ * @param jsonRpcUrl Url with remote node to be forked locally
+ * @param blockNumber Number of the block that the fork will happen
+ * @param isVyperMapping Whether the storage uses Vyper or Solidity mapping
+ */
+export const forkSetup = async (
+    client: Client & PublicActions & TestActions & WalletActions,
+    accountAddress: Address,
+    tokens: Address[],
+    slots: number[] | undefined,
+    balances: bigint[],
+    jsonRpcUrl: string,
+    blockNumber?: bigint,
+    isVyperMapping: boolean[] = Array(tokens.length).fill(false),
+): Promise<void> => {
+    await client.reset({
+        blockNumber,
+        jsonRpcUrl,
+    });
+
+    await client.impersonateAccount({ address: accountAddress });
+
+    let _slots: number[];
+    if (slots) {
+        _slots = slots;
+    } else {
+        _slots = await Promise.all(
+            tokens.map(async (token, i) =>
+                findTokenBalanceSlot(
+                    client,
+                    accountAddress,
+                    token,
+                    isVyperMapping[i],
+                ),
+            ),
+        );
+        console.log(`slots: ${_slots}`);
+    }
+
+    for (let i = 0; i < tokens.length; i++) {
+        // Set initial account balance for each token that will be used to join pool
+        await setTokenBalance(
+            client,
+            accountAddress,
+            tokens[i],
+            _slots[i],
+            balances[i],
+            isVyperMapping[i],
+        );
+
+        // Approve appropriate allowances so that vault contract can move tokens
+        await approveToken(client, accountAddress, tokens[i]);
+    }
+};
diff --git a/test/weightedJoin.integration.test.ts b/test/weightedJoin.integration.test.ts
@@ -30,11 +30,7 @@ import { JoinParser } from '../src/entities/join/parser';
 import { Address } from '../src/types';
 import { CHAINS, ChainId, getPoolAddress } from '../src/utils';
 
-import {
-    approveToken,
-    sendTransactionGetBalances,
-    setTokenBalance,
-} from './lib/utils/helper';
+import { forkSetup, sendTransactionGetBalances } from './lib/utils/helper';
 
 const testAddress = '0x10A19e7eE7d7F8a52822f6817de8ea18204F2e4f'; // Balancer DAO Multisig
 
@@ -66,18 +62,19 @@ describe('weighted join test', () => {
     });
 
     beforeEach(async () => {
-        // reset local fork
-        await client.reset({
-            blockNumber,
-            jsonRpcUrl: process.env.ETHEREUM_RPC_URL,
-        });
-
-        // prepare test client with balance and token approvals
-        await client.impersonateAccount({ address: testAddress });
-
         // get pool state from api
         poolFromApi = await api.getPool(poolId);
 
+        await forkSetup(
+            client,
+            testAddress,
+            poolFromApi.tokens.map((t) => t.address),
+            undefined, // TODO: hardcode these values to improve test performance
+            poolFromApi.tokens.map((t) => parseUnits('100', t.decimals)),
+            process.env.ETHEREUM_RPC_URL as string,
+            blockNumber,
+        );
+
         // setup join helper
         const joinParser = new JoinParser();
         weightedJoin = joinParser.getJoin(poolFromApi.type);
@@ -100,8 +97,6 @@ describe('weighted join test', () => {
             );
             const amountIn = TokenAmount.fromHumanAmount(tokenIn, '1');
 
-            await approveToken(client, testAddress, tokenIn.address);
-
             // perform join query to get expected bpt out
             const joinInput: ExactInJoinInput = {
                 amountsIn: [amountIn],
@@ -159,8 +154,6 @@ describe('weighted join test', () => {
             );
             const amountIn = TokenAmount.fromHumanAmount(tokenIn, '1');
 
-            await approveToken(client, testAddress, tokenIn.address);
-
             // perform join query to get expected bpt out
             const joinInput: ExactInJoinInput = {
                 amountsIn: [amountIn],
@@ -229,15 +222,6 @@ describe('weighted join test', () => {
             const amountOut = TokenAmount.fromHumanAmount(tokenOut, '1');
             const tokenIn = '0x7f39c581f595b53c5cb19bd0b3f8da6c935e2ca0';
 
-            await approveToken(client, testAddress, tokenIn);
-            await setTokenBalance(
-                client,
-                testAddress,
-                tokenIn,
-                0,
-                parseUnits('100', 18),
-            );
-
             // perform join query to get expected bpt out
             const joinInput: ExactOutSingleAssetJoinInput = {
                 bptOut: amountOut,
@@ -289,19 +273,6 @@ describe('weighted join test', () => {
 
         test('proportional join', async () => {
             const amountOut = TokenAmount.fromHumanAmount(tokenOut, '1');
-            const slots = [0, 1];
-
-            for (let i = 0; i < poolFromApi.tokens.length; i++) {
-                const token = poolFromApi.tokens[i];
-                await approveToken(client, testAddress, token.address);
-                await setTokenBalance(
-                    client,
-                    testAddress,
-                    token.address,
-                    slots[i],
-                    parseUnits('100', token.decimals),
-                );
-            }
 
             // perform join query to get expected bpt out
             const joinInput: ExactOutProportionalJoinInput = {