@aztec/end-to-end 0.0.1-commit.2ed92850 → 0.0.1-commit.2f68f620

package/src/shared/uniswap_l1_l2.ts

@@ -11,18 +11,22 @@ import type { DeployAztecL1ContractsReturnType } from '@aztec/ethereum/deploy-az
 import { deployL1Contract } from '@aztec/ethereum/deploy-l1-contract';
 import type { ExtendedViemWalletClient } from '@aztec/ethereum/types';
 import { extractEvent } from '@aztec/ethereum/utils';
-import { CheckpointNumber, EpochNumber } from '@aztec/foundation/branded-types';
+import { EpochNumber } from '@aztec/foundation/branded-types';
 import { sha256ToField } from '@aztec/foundation/crypto/sha256';
+import { retryUntil } from '@aztec/foundation/retry';
 import { InboxAbi, UniswapPortalAbi, UniswapPortalBytecode } from '@aztec/l1-artifacts';
 import { UniswapContract } from '@aztec/noir-contracts.js/Uniswap';
 import { computeL2ToL1MessageHash } from '@aztec/stdlib/hash';
-import { computeL2ToL1MembershipWitness } from '@aztec/stdlib/messaging';
-import type { TestWallet } from '@aztec/test-wallet/server';
 
 import { jest } from '@jest/globals';
 import { type GetContractReturnType, getContract, parseEther, toFunctionSelector } from 'viem';
 
-import { type EndToEndContext, ensureAccountContractsPublished } from '../fixtures/utils.js';
+import {
+  type EndToEndContext,
+  ensureAccountContractsPublished,
+  ensureAuthRegistryPublished,
+} from '../fixtures/utils.js';
+import type { TestWallet } from '../test-wallet/test_wallet.js';
 import { CrossChainTestHarness } from './cross_chain_test_harness.js';
 
 // PSA: This tests works on forked mainnet. There is a dump of the data in `dumpedState` such that we
@@ -32,7 +36,7 @@ import { CrossChainTestHarness } from './cross_chain_test_harness.js';
 // anvil --fork-url https://mainnet.infura.io/v3/9928b52099854248b3a096be07a6b23c --fork-block-number 17514288 --chain-id 31337
 // For CI, this is configured in `run_tests.sh` and `docker-compose-images.yml`
 
-const TIMEOUT = 360_000;
+const TIMEOUT = 15 * 60 * 1000;
 
 export const uniswapL1L2TestSuite = (
   setup: () => Promise<EndToEndContext>,
@@ -84,7 +88,7 @@ export const uniswapL1L2TestSuite = (
 
       l1Client = deployL1ContractsValues.l1Client;
 
-      t.watcher?.setIsMarkingAsProven(false);
+      t.watcher.setIsMarkingAsProven(false);
 
       if (Number(await l1Client.getBlockNumber()) < expectedForkBlockNumber) {
         throw new Error('This test must be run on a fork of mainnet with the expected fork block');
@@ -98,6 +102,7 @@ export const uniswapL1L2TestSuite = (
       ownerEthAddress = EthAddress.fromString((await l1Client.getAddresses())[0]);
 
       await ensureAccountContractsPublished(wallet, [ownerAddress, sponsorAddress]);
+      await ensureAuthRegistryPublished(wallet, ownerAddress);
 
       logger.info('Deploying DAI Portal, initializing and deploying l2 contract...');
       daiCrossChainHarness = await CrossChainTestHarness.new(
@@ -130,7 +135,9 @@ export const uniswapL1L2TestSuite = (
         client: l1Client,
       });
       // deploy l2 uniswap contract and attach to portal
-      uniswapL2Contract = await UniswapContract.deploy(wallet, uniswapPortalAddress).send({ from: ownerAddress });
+      ({ contract: uniswapL2Contract } = await UniswapContract.deploy(wallet, uniswapPortalAddress).send({
+        from: ownerAddress,
+      }));
 
       const registryAddress = (await aztecNode.getNodeInfo()).l1ContractAddresses.registryAddress;
 
@@ -152,7 +159,6 @@ export const uniswapL1L2TestSuite = (
       await cleanup();
     });
 
-    // docs:start:uniswap_private
     it('should uniswap trade on L1 from L2 funds privately (swaps WETH -> DAI)', async () => {
       const wethL1BeforeBalance = await wethCrossChainHarness.getL1BalanceOf(ownerEthAddress);
 
@@ -195,7 +201,7 @@ export const uniswapL1L2TestSuite = (
       logger.info('Withdrawing weth to L1 and sending message to swap to dai');
       const [secretForDepositingSwappedDai, secretHashForDepositingSwappedDai] = await generateClaimSecret();
 
-      const l2UniswapInteractionReceipt = await uniswapL2Contract.methods
+      const { receipt: l2UniswapInteractionReceipt } = await uniswapL2Contract.methods
         .swap_private(
           wethCrossChainHarness.l2Token.address,
           wethCrossChainHarness.l2Bridge.address,
@@ -249,10 +255,13 @@ export const uniswapL1L2TestSuite = (
       // ensure that uniswap contract didn't eat the funds.
       await wethCrossChainHarness.expectPublicBalanceOnL2(uniswapL2Contract.address, 0n);
 
-      // Since the outbox is only consumable when the epoch is proven, we need to advance to the next epoch.
-      const checkpointNumber = CheckpointNumber.fromBlockNumber(l2UniswapInteractionReceipt.blockNumber!);
-      const epoch = await rollup.getEpochNumberForCheckpoint(checkpointNumber);
-      await cheatCodes.rollup.advanceToEpoch(EpochNumber(epoch + 1));
+      // Since the outbox is only consumable when the epoch is proven, we need to advance to the next epoch
+      // before we can ask the witness helper for a covering root.
+      const minedReceipt = await aztecNode.getTxReceipt(l2UniswapInteractionReceipt.txHash);
+      if (minedReceipt.epochNumber === undefined) {
+        throw new Error('L2 Uniswap interaction tx is not yet in an epoch');
+      }
+      await cheatCodes.rollup.advanceToEpoch(EpochNumber(minedReceipt.epochNumber + 1));
       await waitForProven(aztecNode, l2UniswapInteractionReceipt, { provenTimeout: 300 });
 
       // 5. Consume L2 to L1 message by calling uniswapPortal.swap_private()
@@ -260,17 +269,29 @@ export const uniswapL1L2TestSuite = (
       const daiL1BalanceOfPortalBeforeSwap = await daiCrossChainHarness.getL1BalanceOf(
         daiCrossChainHarness.tokenPortalAddress,
       );
-      const swapResult = await computeL2ToL1MembershipWitness(aztecNode, epoch, swapPrivateLeaf);
-      const withdrawResult = await computeL2ToL1MembershipWitness(aztecNode, epoch, withdrawLeaf);
+      const swapResult = await retryUntil(
+        () => aztecNode.getL2ToL1MembershipWitness(minedReceipt.txHash, swapPrivateLeaf),
+        'swap private l2 to l1 witness',
+        60,
+        1,
+      );
+      const { epochNumber: epoch, numCheckpointsInEpoch } = swapResult;
+      const withdrawResult = await retryUntil(
+        () => aztecNode.getL2ToL1MembershipWitness(minedReceipt.txHash, withdrawLeaf),
+        'withdraw l2 to l1 witness',
+        60,
+        1,
+      );
 
-      const swapPrivateL2MessageIndex = swapResult!.leafIndex;
-      const swapPrivateSiblingPath = swapResult!.siblingPath;
+      const swapPrivateL2MessageIndex = swapResult.leafIndex;
+      const swapPrivateSiblingPath = swapResult.siblingPath;
 
-      const withdrawL2MessageIndex = withdrawResult!.leafIndex;
-      const withdrawSiblingPath = withdrawResult!.siblingPath;
+      const withdrawL2MessageIndex = withdrawResult.leafIndex;
+      const withdrawSiblingPath = withdrawResult.siblingPath;
 
       const withdrawMessageMetadata = {
         _epoch: BigInt(epoch),
+        _numCheckpointsInEpoch: BigInt(withdrawResult.numCheckpointsInEpoch),
         _leafIndex: BigInt(withdrawL2MessageIndex),
         _path: withdrawSiblingPath
           .toBufferArray()
@@ -279,6 +300,7 @@ export const uniswapL1L2TestSuite = (
 
       const swapPrivateMessageMetadata = {
         _epoch: BigInt(epoch),
+        _numCheckpointsInEpoch: BigInt(numCheckpointsInEpoch),
         _leafIndex: BigInt(swapPrivateL2MessageIndex),
         _path: swapPrivateSiblingPath
           .toBufferArray()
@@ -336,10 +358,8 @@ export const uniswapL1L2TestSuite = (
       logger.info('WETH balance after swap : ', wethL2BalanceAfterSwap.toString());
       logger.info('DAI balance after swap  : ', daiL2BalanceAfterSwap.toString());
     });
-    // docs:end:uniswap_private
 
     // TODO(#7463): reenable look into this failure https://github.com/AztecProtocol/aztec-packages/actions/runs/9912612912/job/27388320150?pr=7462
-    // // docs:start:uniswap_public
     // it('should uniswap trade on L1 from L2 funds publicly (swaps WETH -> DAI)', async () => {
     //   const wethL1BeforeBalance = await wethCrossChainHarness.getL1BalanceOf(ownerEthAddress);
 
@@ -572,7 +592,6 @@ export const uniswapL1L2TestSuite = (
     //   logger.info('WETH balance after swap : ', wethL2BalanceAfterSwap.toString());
     //   logger.info('DAI balance after swap  : ', daiL2BalanceAfterSwap.toString());
     // });
-    // // docs:end:uniswap_public
 
     // Edge cases for the private flow:
     // note - tests for uniswapPortal.sol and minting asset on L2 are covered in other tests.
@@ -787,7 +806,7 @@ export const uniswapL1L2TestSuite = (
       logger.info('Withdrawing weth to L1 and sending message to swap to dai');
 
       const [, secretHashForDepositingSwappedDai] = await generateClaimSecret();
-      const withdrawReceipt = await uniswapL2Contract.methods
+      const { receipt: withdrawReceipt } = await uniswapL2Contract.methods
         .swap_private(
           wethCrossChainHarness.l2Token.address,
           wethCrossChainHarness.l2Bridge.address,
@@ -838,20 +857,41 @@ export const uniswapL1L2TestSuite = (
         chainId: new Fr(l1Client.chain.id),
       });
 
-      const epoch = await rollup.getEpochNumberForCheckpoint(
-        CheckpointNumber.fromBlockNumber(withdrawReceipt.blockNumber!),
+      // ensure that user's funds were burnt
+      await wethCrossChainHarness.expectPrivateBalanceOnL2(ownerAddress, wethL2BalanceBeforeSwap - wethAmountToBridge);
+
+      // Since the outbox is only consumable when the epoch is proven, we need to advance to the next epoch
+      // before we can ask the witness helper for a covering root.
+      const minedReceipt = await aztecNode.getTxReceipt(withdrawReceipt.txHash);
+      if (minedReceipt.epochNumber === undefined) {
+        throw new Error('Withdraw tx is not yet in an epoch');
+      }
+      await cheatCodes.rollup.advanceToEpoch(EpochNumber(minedReceipt.epochNumber + 1));
+      await waitForProven(aztecNode, withdrawReceipt, { provenTimeout: 300 });
+
+      const swapResult = await retryUntil(
+        () => aztecNode.getL2ToL1MembershipWitness(minedReceipt.txHash, swapPrivateLeaf),
+        'swap private l2 to l1 witness',
+        60,
+        1,
+      );
+      const { epochNumber: epoch, numCheckpointsInEpoch } = swapResult;
+      const withdrawResult = await retryUntil(
+        () => aztecNode.getL2ToL1MembershipWitness(minedReceipt.txHash, withdrawLeaf),
+        'withdraw l2 to l1 witness',
+        60,
+        1,
       );
-      const swapResult = await computeL2ToL1MembershipWitness(aztecNode, epoch, swapPrivateLeaf);
-      const withdrawResult = await computeL2ToL1MembershipWitness(aztecNode, epoch, withdrawLeaf);
 
-      const swapPrivateL2MessageIndex = swapResult!.leafIndex;
-      const swapPrivateSiblingPath = swapResult!.siblingPath;
+      const swapPrivateL2MessageIndex = swapResult.leafIndex;
+      const swapPrivateSiblingPath = swapResult.siblingPath;
 
-      const withdrawL2MessageIndex = withdrawResult!.leafIndex;
-      const withdrawSiblingPath = withdrawResult!.siblingPath;
+      const withdrawL2MessageIndex = withdrawResult.leafIndex;
+      const withdrawSiblingPath = withdrawResult.siblingPath;
 
       const withdrawMessageMetadata = {
         _epoch: BigInt(epoch),
+        _numCheckpointsInEpoch: BigInt(withdrawResult.numCheckpointsInEpoch),
         _leafIndex: BigInt(withdrawL2MessageIndex),
         _path: withdrawSiblingPath
           .toBufferArray()
@@ -860,19 +900,13 @@ export const uniswapL1L2TestSuite = (
 
       const swapPrivateMessageMetadata = {
         _epoch: BigInt(epoch),
+        _numCheckpointsInEpoch: BigInt(numCheckpointsInEpoch),
         _leafIndex: BigInt(swapPrivateL2MessageIndex),
         _path: swapPrivateSiblingPath
           .toBufferArray()
           .map((buf: Buffer) => `0x${buf.toString('hex')}`) as readonly `0x${string}`[],
       };
 
-      // ensure that user's funds were burnt
-      await wethCrossChainHarness.expectPrivateBalanceOnL2(ownerAddress, wethL2BalanceBeforeSwap - wethAmountToBridge);
-
-      // Since the outbox is only consumable when the epoch is proven, we need to advance to the next epoch.
-      await cheatCodes.rollup.advanceToEpoch(EpochNumber(epoch + 1));
-      await waitForProven(aztecNode, withdrawReceipt, { provenTimeout: 300 });
-
       // On L1 call swap_public!
       logger.info('call swap_public on L1');
       const swapArgs = [
@@ -916,7 +950,7 @@ export const uniswapL1L2TestSuite = (
 
       // Call swap_public on L2
       const secretHashForDepositingSwappedDai = Fr.random();
-      const withdrawReceipt = await uniswapL2Contract.methods
+      const { receipt: withdrawReceipt } = await uniswapL2Contract.methods
         .swap_public(
           ownerAddress,
           wethCrossChainHarness.l2Bridge.address,
@@ -972,20 +1006,41 @@ export const uniswapL1L2TestSuite = (
         chainId: new Fr(l1Client.chain.id),
       });
 
-      const epoch = await rollup.getEpochNumberForCheckpoint(
-        CheckpointNumber.fromBlockNumber(withdrawReceipt.blockNumber!),
+      // check weth balance of owner on L2 (we first bridged `wethAmountToBridge` into L2 and now withdrew it!)
+      await wethCrossChainHarness.expectPublicBalanceOnL2(ownerAddress, 0n);
+
+      // Since the outbox is only consumable when the epoch is proven, we need to advance to the next epoch
+      // before we can ask the witness helper for a covering root.
+      const minedReceipt = await aztecNode.getTxReceipt(withdrawReceipt.txHash);
+      if (minedReceipt.epochNumber === undefined) {
+        throw new Error('Withdraw tx is not yet in an epoch');
+      }
+      await cheatCodes.rollup.advanceToEpoch(EpochNumber(minedReceipt.epochNumber + 1));
+      await waitForProven(aztecNode, withdrawReceipt, { provenTimeout: 300 });
+
+      const swapResult = await retryUntil(
+        () => aztecNode.getL2ToL1MembershipWitness(minedReceipt.txHash, swapPublicLeaf),
+        'swap public l2 to l1 witness',
+        60,
+        1,
+      );
+      const { epochNumber: epoch, numCheckpointsInEpoch } = swapResult;
+      const withdrawResult = await retryUntil(
+        () => aztecNode.getL2ToL1MembershipWitness(minedReceipt.txHash, withdrawLeaf),
+        'withdraw l2 to l1 witness',
+        60,
+        1,
       );
-      const swapResult = await computeL2ToL1MembershipWitness(aztecNode, epoch, swapPublicLeaf);
-      const withdrawResult = await computeL2ToL1MembershipWitness(aztecNode, epoch, withdrawLeaf);
 
-      const swapPublicL2MessageIndex = swapResult!.leafIndex;
-      const swapPublicSiblingPath = swapResult!.siblingPath;
+      const swapPublicL2MessageIndex = swapResult.leafIndex;
+      const swapPublicSiblingPath = swapResult.siblingPath;
 
-      const withdrawL2MessageIndex = withdrawResult!.leafIndex;
-      const withdrawSiblingPath = withdrawResult!.siblingPath;
+      const withdrawL2MessageIndex = withdrawResult.leafIndex;
+      const withdrawSiblingPath = withdrawResult.siblingPath;
 
       const withdrawMessageMetadata = {
         _epoch: BigInt(epoch),
+        _numCheckpointsInEpoch: BigInt(withdrawResult.numCheckpointsInEpoch),
         _leafIndex: BigInt(withdrawL2MessageIndex),
         _path: withdrawSiblingPath
           .toBufferArray()
@@ -994,19 +1049,13 @@ export const uniswapL1L2TestSuite = (
 
       const swapPublicMessageMetadata = {
         _epoch: BigInt(epoch),
+        _numCheckpointsInEpoch: BigInt(numCheckpointsInEpoch),
         _leafIndex: BigInt(swapPublicL2MessageIndex),
         _path: swapPublicSiblingPath
           .toBufferArray()
           .map((buf: Buffer) => `0x${buf.toString('hex')}`) as readonly `0x${string}`[],
       };
 
-      // check weth balance of owner on L2 (we first bridged `wethAmountToBridge` into L2 and now withdrew it!)
-      await wethCrossChainHarness.expectPublicBalanceOnL2(ownerAddress, 0n);
-
-      // Since the outbox is only consumable when the epoch is proven, we need to advance to the next epoch.
-      await cheatCodes.rollup.advanceToEpoch(EpochNumber(epoch + 1));
-      await waitForProven(aztecNode, withdrawReceipt, { provenTimeout: 300 });
-
       // Call swap_private on L1
       logger.info('Execute withdraw and swap on the uniswapPortal!');
 

package/src/shared/wait_for_l1_to_l2_message.ts

@@ -0,0 +1,23 @@
+import type { Fr } from '@aztec/aztec.js/fields';
+import type { AztecNode } from '@aztec/aztec.js/node';
+import { retryUntil } from '@aztec/foundation/retry';
+
+/**
+ * Waits until the archiver has seen the L1 to L2 message and indexed it. Unlike
+ * {@link waitForL1ToL2MessageReady} from `@aztec/aztec.js/messaging`, this does not
+ * require the L2 chain to have advanced to the message's checkpoint — it only confirms
+ * the message has been picked up from L1. Use this in tests that explicitly produce L2
+ * blocks afterwards to make the message consumable.
+ */
+export function waitForL1ToL2MessageSeen(
+  node: Pick<AztecNode, 'getL1ToL2MessageCheckpoint'>,
+  l1ToL2MessageHash: Fr,
+  opts: { timeoutSeconds: number },
+) {
+  return retryUntil(
+    async () => (await node.getL1ToL2MessageCheckpoint(l1ToL2MessageHash)) !== undefined,
+    `L1 to L2 message ${l1ToL2MessageHash.toString()} seen`,
+    opts.timeoutSeconds,
+    1,
+  );
+}

package/src/simulators/lending_simulator.ts

@@ -7,6 +7,7 @@ import { SlotNumber } from '@aztec/foundation/branded-types';
 import { poseidon2Hash } from '@aztec/foundation/crypto/poseidon';
 import type { TestDateProvider } from '@aztec/foundation/timer';
 import type { LendingContract } from '@aztec/noir-contracts.js/Lending';
+import type { AztecNode, AztecNodeDebug } from '@aztec/stdlib/interfaces/client';
 
 import type { TokenSimulator } from './token_simulator.js';
 
@@ -92,30 +93,45 @@ export class LendingSimulator {
     public stableCoin: TokenSimulator,
   ) {}
 
-  async prepare() {
+  prepare() {
     this.accumulator = BASE;
-    const slot = await this.rollup.getSlotAt(BigInt(await this.cc.eth.timestamp()) + BigInt(this.ethereumSlotDuration));
-    this.time = Number(await this.rollup.getTimestampForSlot(slot));
+    this.time = 0;
   }
 
-  async progressSlots(diff: number, dateProvider?: TestDateProvider) {
+  /**
+   * Advances the simulator's accumulator and clock to match a block timestamp observed on chain.
+   * Call this BEFORE applying any accumulator-sensitive mutation (borrow/repay) so the mutation
+   * sees the same accumulator as the contract did during execution.
+   */
+  observeBlockTimestamp(ts: number) {
+    const diff = ts - this.time;
+    if (diff > 0) {
+      this.accumulator = muldivDown(this.accumulator, computeMultiplier(this.rate, BigInt(diff)), BASE);
+    }
+    this.time = ts;
+  }
+
+  async progressSlots(diff: number, _dateProvider?: TestDateProvider, node?: AztecNode & AztecNodeDebug) {
     if (diff <= 1) {
       return;
     }
 
-    const slot = await this.rollup.getSlotAt(BigInt(await this.cc.eth.timestamp()));
+    const slot = await this.rollup.getSlotAt(BigInt(await this.cc.eth.lastBlockTimestamp()));
     const targetSlot = SlotNumber(slot + diff);
     const ts = Number(await this.rollup.getTimestampForSlot(targetSlot));
-    const timeDiff = ts - this.time;
-    this.time = ts;
 
-    // Mine ethereum blocks such that the next block will be in a new slot
-    await this.cc.eth.warp(this.time - this.ethereumSlotDuration);
-    if (dateProvider) {
-      dateProvider.setTime(this.time * 1000);
+    // Queue-aware warp under AutomineSequencer: atomic warp + mineBlock that advances L2 time to the
+    // target slot. The cheat code routes through the AutomineSequencer queue when one is installed,
+    // and otherwise falls back to a manual warp + mineBlock loop.
+    if (node) {
+      await this.cc.warpL2TimeAtLeastTo(node, ts);
+    } else {
+      await this.cc.eth.warp(ts - this.ethereumSlotDuration);
     }
+
+    // Mark the latest checkpoint as proven so the rollup does not reorg pending checkpoints when
+    // time jumps far enough forward to cross an unproven epoch boundary.
     await this.cc.rollup.markAsProven(await this.rollup.getCheckpointNumber());
-    this.accumulator = muldivDown(this.accumulator, computeMultiplier(this.rate, BigInt(timeDiff)), BASE);
   }
 
   depositPrivate(from: AztecAddress, onBehalfOf: Fr, amount: bigint) {
@@ -186,14 +202,16 @@ export class LendingSimulator {
 
     expect(this.borrowed).toEqual(this.stableCoin.totalSupply - this.mintedOutside);
 
-    const asset = await this.lendingContract.methods.get_asset(0).simulate({ from: this.account.address });
+    const { result: asset } = await this.lendingContract.methods.get_asset(0).simulate({ from: this.account.address });
 
     const interestAccumulator = asset['interest_accumulator'];
     expect(interestAccumulator).toEqual(this.accumulator);
     expect(asset['last_updated_ts']).toEqual(BigInt(this.time));
 
     for (const key of [this.account.address, AztecAddress.fromField(await this.account.key())]) {
-      const privatePos = await this.lendingContract.methods.get_position(key).simulate({ from: this.account.address });
+      const { result: privatePos } = await this.lendingContract.methods
+        .get_position(key)
+        .simulate({ from: this.account.address });
       expect(new Fr(privatePos['collateral'])).toEqual(this.collateral[key.toString()] ?? Fr.ZERO);
       expect(new Fr(privatePos['static_debt'])).toEqual(this.staticDebt[key.toString()] ?? Fr.ZERO);
       expect(privatePos['debt']).toEqual(

package/src/simulators/token_simulator.ts

@@ -109,7 +109,9 @@ export class TokenSimulator {
       await Promise.all(
         chunk(calls, 5).map(batch => new BatchCall(this.defaultWallet, batch).simulate({ from: this.defaultAddress })),
       )
-    ).flat();
+    )
+      .flatMap(r => r.result)
+      .map(r => r.result);
     expect(results[0]).toEqual(this.totalSupply);
 
     // Check that all our balances match
@@ -119,39 +121,13 @@ export class TokenSimulator {
   }
 
   async checkPrivate() {
-    // Private calls
-    const defaultLookups = [];
-    const nonDefaultLookups = [];
-
     for (const address of this.accounts) {
-      if (this.lookupProvider.has(address.toString())) {
-        nonDefaultLookups.push(address);
-      } else {
-        defaultLookups.push(address);
-      }
-    }
-
-    const defaultCalls = defaultLookups.map(address => this.token.methods.balance_of_private(address));
-    const results = (
-      await Promise.all(
-        chunk(defaultCalls, 4).map(batch =>
-          new BatchCall(this.defaultWallet, batch).simulate({ from: this.defaultAddress }),
-        ),
-      )
-    ).flat();
-    for (let i = 0; i < defaultLookups.length; i++) {
-      expect(results[i]).toEqual(this.balanceOfPrivate(defaultLookups[i]));
-    }
-
-    // We are just running individual calls for the non-default lookups
-    // @todo We should also batch these
-    for (const address of nonDefaultLookups) {
       const wallet = this.lookupProvider.get(address.toString());
       const asset = wallet ? this.token.withWallet(wallet) : this.token;
 
-      const actualPrivateBalance = await asset.methods
-        .balance_of_private({ address })
-        .simulate({ from: this.defaultAddress });
+      const { result: actualPrivateBalance } = await asset.methods
+        .balance_of_private(address)
+        .simulate({ from: address });
       expect(actualPrivateBalance).toEqual(this.balanceOfPrivate(address));
     }
   }