refactor: speed up transfer test

Author: LHerskind

yarn-project/end-to-end/scripts/create_k8s_dashboard.sh

@@ -467,7 +467,7 @@ async function setupFromState(statePath: string, logger: Logger): Promise<Subsys
  * The 'restore' function is not provided, as it must be a closure within the test context to capture the results.
  */
 export const addAccounts =
-  (numberOfAccounts: number, logger: DebugLogger) =>
+  (numberOfAccounts: number, logger: DebugLogger, waitUntilProven = true) =>
   async ({ pxe }: { pxe: PXE }) => {
     // Generate account keys.
     const accountKeys: [Fr, GrumpkinScalar][] = Array.from({ length: numberOfAccounts }).map(_ => [
@@ -494,7 +494,7 @@ export const addAccounts =
 
     logger.verbose('Deploying accounts...');
     const txs = await Promise.all(accountManagers.map(account => account.deploy()));
-    await Promise.all(txs.map(tx => tx.wait({ interval: 0.1, proven: true })));
+    await Promise.all(txs.map(tx => tx.wait({ interval: 0.1, proven: waitUntilProven })));
 
     return { accountKeys };
   };
@@ -505,12 +505,16 @@ export const addAccounts =
  * @param sender - Wallet to send the deployment tx.
  * @param accountsToDeploy - Which accounts to publicly deploy.
  */
-export async function publicDeployAccounts(sender: Wallet, accountsToDeploy: (CompleteAddress | AztecAddress)[]) {
+export async function publicDeployAccounts(
+  sender: Wallet,
+  accountsToDeploy: (CompleteAddress | AztecAddress)[],
+  waitUntilProven = true,
+) {
   const accountAddressesToDeploy = accountsToDeploy.map(a => ('address' in a ? a.address : a));
   const instances = await Promise.all(accountAddressesToDeploy.map(account => sender.getContractInstance(account)));
   const batch = new BatchCall(sender, [
     (await registerContractClass(sender, SchnorrAccountContractArtifact)).request(),
     ...instances.map(instance => deployInstance(sender, instance!).request()),
   ]);
-  await batch.send().wait({ proven: true });
+  await batch.send().wait({ proven: waitUntilProven });
 }

yarn-project/end-to-end/src/fixtures/snapshot_manager.ts

@@ -3,7 +3,6 @@ import {
   type AccountWallet,
   type AccountWalletWithSecretKey,
   type AztecAddress,
-  type CompleteAddress,
   ExtendedNote,
   Fr,
   Note,
@@ -17,7 +16,7 @@ import { TokenContract } from '@aztec/noir-contracts.js';
 
 import { jest } from '@jest/globals';
 
-import { addAccounts, publicDeployAccounts } from '../fixtures/snapshot_manager.js';
+import { addAccounts } from '../fixtures/snapshot_manager.js';
 
 const { PXE_URL } = process.env;
 if (!PXE_URL) {
@@ -65,52 +64,93 @@ describe('token transfer test', () => {
   const TOKEN_SYMBOL = 'USD';
   const TOKEN_DECIMALS = 18n;
   const MINT_AMOUNT = 20n;
+
+  const WALLET_COUNT = 1;
+  const ROUNDS = 5n;
+
   let pxe: PXE;
   let wallets: AccountWalletWithSecretKey[];
-  let completeAddresses: CompleteAddress[];
+  let recipientWallet: AccountWalletWithSecretKey;
   let tokenAddress: AztecAddress;
   let tokenAtWallet0: TokenContract;
+
   beforeAll(async () => {
+    expect(ROUNDS).toBeLessThanOrEqual(MINT_AMOUNT);
+
+    // My guess is that is never proven so if we are waiting for it to prove we wait forever?
     pxe = await createCompatibleClient(PXE_URL, logger);
-    const { accountKeys } = await addAccounts(3, logger)({ pxe });
+    {
+      const { accountKeys } = await addAccounts(1, logger, false)({ pxe });
+      const accountManagers = accountKeys.map(ak => getSchnorrAccount(pxe, ak[0], ak[1], 1));
+
+      const partialAddress = accountManagers[0].getCompleteAddress().partialAddress;
+      await pxe.registerAccount(accountKeys[0][0], partialAddress);
+      recipientWallet = await accountManagers[0].getWallet();
+      logger.verbose(`Recipient Wallet address: ${recipientWallet.getAddress()} registered`);
+    }
+
+    const { accountKeys } = await addAccounts(Number(WALLET_COUNT), logger, false)({ pxe });
     const accountManagers = accountKeys.map(ak => getSchnorrAccount(pxe, ak[0], ak[1], 1));
-    wallets = await Promise.all(accountManagers.map(a => a.getWallet()));
-    completeAddresses = await pxe.getRegisteredAccounts();
-    wallets.forEach((w, i) => logger.verbose(`Wallet ${i} address: ${w.getAddress()}`));
-    await publicDeployAccounts(wallets[0], completeAddresses.slice(0, 2));
+
+    wallets = await Promise.all(
+      accountManagers.map(async (a, i) => {
+        const partialAddress = a.getCompleteAddress().partialAddress;
+        await pxe.registerAccount(accountKeys[i][0], partialAddress);
+        const wallet = await a.getWallet();
+        logger.verbose(`Wallet ${i} address: ${wallet.getAddress()} registered`);
+        return wallet;
+      }),
+    );
 
     logger.verbose(`Deploying TokenContract...`);
     const tokenContract = await TokenContract.deploy(
       wallets[0],
-      completeAddresses[0],
+      wallets[0].getAddress(),
       TOKEN_NAME,
       TOKEN_SYMBOL,
       TOKEN_DECIMALS,
     )
       .send()
-      .deployed();
+      .deployed({ timeout: 600 });
 
     tokenAddress = tokenContract.address;
     tokenAtWallet0 = await TokenContract.at(tokenAddress, wallets[0]);
 
-    logger.verbose(`Minting ${MINT_AMOUNT} publicly...`);
-    await tokenAtWallet0.methods.mint_public(completeAddresses[0].address, MINT_AMOUNT).send().wait();
-
-    logger.verbose(`Minting ${MINT_AMOUNT} privately...`);
-    const secret = Fr.random();
-    const secretHash = computeSecretHash(secret);
-    const receipt = await tokenAtWallet0.methods.mint_private(MINT_AMOUNT, secretHash).send().wait();
-
-    await addPendingShieldNoteToPXE({
-      amount: MINT_AMOUNT,
-      secretHash,
-      txHash: receipt.txHash,
-      accountAddress: completeAddresses[0].address,
-      assetAddress: tokenAddress,
-      wallet: wallets[0],
+    logger.verbose(`Minting ${MINT_AMOUNT} public assets to the ${wallets.length} wallets...`);
+
+    await Promise.all(
+      wallets.map(w => tokenAtWallet0.methods.mint_public(w.getAddress(), MINT_AMOUNT).send().wait({ timeout: 600 })),
+    );
+
+    logger.verbose(`Minting ${MINT_AMOUNT} private assets to the ${wallets.length} wallets...`);
+    const secrets: Fr[] = wallets.map(() => Fr.random());
+
+    const txs = await Promise.all(
+      wallets.map((w, i) =>
+        tokenAtWallet0.methods.mint_private(MINT_AMOUNT, computeSecretHash(secrets[i])).send().wait({ timeout: 600 }),
+      ),
+    );
+
+    wallets.forEach(async (wallet, i) => {
+      await addPendingShieldNoteToPXE({
+        amount: MINT_AMOUNT,
+        secretHash: computeSecretHash(secrets[i]),
+        txHash: txs[i].txHash,
+        accountAddress: wallet.getAddress(),
+        assetAddress: tokenAddress,
+        wallet: wallet,
+      });
     });
-    const txClaim = tokenAtWallet0.methods.redeem_shield(completeAddresses[0].address, MINT_AMOUNT, secret).send();
-    await txClaim.wait({ debug: true });
+
+    await Promise.all(
+      wallets.map(async (w, i) =>
+        (await TokenContract.at(tokenAddress, w)).methods
+          .redeem_shield(w.getAddress(), MINT_AMOUNT, secrets[i])
+          .send()
+          .wait({ timeout: 600 }),
+      ),
+    );
+
     logger.verbose(`Minting complete.`);
   });
 
@@ -119,34 +159,54 @@ describe('token transfer test', () => {
     expect(name).toBe(TOKEN_NAME);
   });
 
-  it('can transfer 1 publicly', async () => {
+  it('can transfer 1 token privately and publicly', async () => {
+    const recipient = recipientWallet.getAddress();
     const transferAmount = 1n;
-    const numTransfers = MINT_AMOUNT / transferAmount;
-    const initialBalance = await tokenAtWallet0.methods.balance_of_public(completeAddresses[0].address).simulate();
-    expect(initialBalance).toBeGreaterThanOrEqual(transferAmount);
-    for (let i = 1n; i <= numTransfers; i++) {
-      await tokenAtWallet0.methods
-        .transfer_public(completeAddresses[0].address, completeAddresses[1].address, transferAmount, 0)
-        .send()
-        .wait();
-    }
-    const finalBalance0 = await tokenAtWallet0.methods.balance_of_public(completeAddresses[0].address).simulate();
-    expect(finalBalance0).toBe(0n);
-    const finalBalance1 = await tokenAtWallet0.methods.balance_of_public(completeAddresses[1].address).simulate();
-    expect(finalBalance1).toBe(MINT_AMOUNT);
-  });
 
-  it('can transfer 1 privately', async () => {
-    const transferAmount = 1n;
-    const numTransfers = MINT_AMOUNT / transferAmount;
-    const initialBalance = await tokenAtWallet0.methods.balance_of_private(completeAddresses[0].address).simulate();
-    expect(initialBalance).toBeGreaterThanOrEqual(transferAmount);
-    for (let i = 1n; i <= numTransfers; i++) {
-      await tokenAtWallet0.methods.transfer(completeAddresses[1].address, transferAmount).send().wait();
+    wallets.forEach(async w => {
+      expect(MINT_AMOUNT).toBe(
+        await (await TokenContract.at(tokenAddress, w)).methods.balance_of_private(w.getAddress()).simulate(),
+      );
+      expect(MINT_AMOUNT).toBe(await tokenAtWallet0.methods.balance_of_public(w.getAddress()).simulate());
+    });
+
+    expect(0n).toBe(
+      await (await TokenContract.at(tokenAddress, recipientWallet)).methods.balance_of_private(recipient).simulate(),
+    );
+    expect(0n).toBe(await tokenAtWallet0.methods.balance_of_public(recipient).simulate());
+
+    // For each round, make both private and public transfers
+    for (let i = 1n; i <= ROUNDS; i++) {
+      const txs = await Promise.all([
+        ...wallets.map(async w =>
+          (await TokenContract.at(tokenAddress, w)).methods.transfer(recipient, transferAmount),
+        ),
+        ...wallets.map(async w =>
+          (
+            await TokenContract.at(tokenAddress, w)
+          ).methods.transfer_public(w.getAddress(), recipient, transferAmount, 0),
+        ),
+      ]);
+
+      txs.forEach(async t => await t.prove());
+
+      await Promise.all(txs.map(t => t.send().wait({ timeout: 600 })));
     }
-    const finalBalance0 = await tokenAtWallet0.methods.balance_of_private(completeAddresses[0].address).simulate();
-    expect(finalBalance0).toBe(0n);
-    const finalBalance1 = await tokenAtWallet0.methods.balance_of_private(completeAddresses[1].address).simulate();
-    expect(finalBalance1).toBe(MINT_AMOUNT);
+
+    wallets.forEach(async w => {
+      expect(MINT_AMOUNT - ROUNDS * transferAmount).toBe(
+        await (await TokenContract.at(tokenAddress, w)).methods.balance_of_private(w.getAddress()).simulate(),
+      );
+      expect(MINT_AMOUNT - ROUNDS * transferAmount).toBe(
+        await tokenAtWallet0.methods.balance_of_public(w.getAddress()).simulate(),
+      );
+    });
+
+    expect(ROUNDS * transferAmount * BigInt(wallets.length)).toBe(
+      await (await TokenContract.at(tokenAddress, recipientWallet)).methods.balance_of_private(recipient).simulate(),
+    );
+    expect(ROUNDS * transferAmount * BigInt(wallets.length)).toBe(
+      await tokenAtWallet0.methods.balance_of_public(recipient).simulate(),
+    );
   });
 });