@aztec/end-to-end 0.0.1-commit.a072138 → 0.0.1-commit.a4600f49

package/src/shared/timing_env.mjs

@@ -0,0 +1,363 @@
+import { appendFileSync, writeFileSync } from 'node:fs';
+import { basename } from 'node:path';
+
+import CustomEnvironment from '../../../foundation/src/jest/env.mjs';
+
+// Per-test e2e timing environment. Gated entirely on the TEST_TIMING_FILE env var: when unset, this
+// behaves exactly like the base CustomEnvironment (it only delegates). When set, it records, per test
+// worker process, the time spent in jest before/after hooks and the test body, and folds in the named
+// spans recorded via the `testSpan()` wrapper (fixtures/timing.ts) through a collector shared on
+// `this.global`.
+//
+// Output is one JSONL file per worker process (the env runs once per worker). Each line carries a
+// `type` discriminator and is one of:
+//   - `type: 'test'` (`name` set): beforeEach hooks, the it() body, afterEach hooks for one test; or
+//   - `type: 'suite'` (`name: null`): the suite-scoped beforeAll/afterAll hooks for the whole file.
+//
+// Each line also gains a `spans` map: per `category:label` tag the test/suite touched, the collector
+// records `{ count, totalMs, busyMs, maxMs }`. `busyMs` is the duration of the union of the spans'
+// `[start, end)` intervals on the shared `performance.now()` clock, so concurrent spans (Promise.all
+// fan-out) do not inflate it the way `totalMs` does. The legacy `setupFnMs`/`teardownFnMs` fields are
+// kept for back-compat, derived from the `setup:env`/`teardown:env` spans.
+//
+// When `TEST_TIMING_SPANS=1`, the collector additionally retains every raw span occurrence and emits
+// one `type:"span"` line per occurrence (owner, name, ms) for deep dives, at the cost of a larger file.
+export default class TimingEnvironment extends CustomEnvironment {
+  constructor(config, context) {
+    super(config, context);
+
+    this.timingFile = process.env.TEST_TIMING_FILE;
+    if (!this.timingFile) {
+      return;
+    }
+
+    this.suite = basename(context?.testPath ?? config?.testPath ?? 'unknown').replace(/\.test\.[cm]?[jt]s$/, '');
+    this.meta = {
+      commit: process.env.COMMIT_HASH ?? null,
+      branch: process.env.TARGET_BRANCH ?? process.env.REF_NAME ?? null,
+      runId: process.env.RUN_ID ?? null,
+    };
+
+    // Opt-in per-occurrence emission for deep dives (off by default keeps the JSONL one line per test).
+    this.emitSpanLines = process.env.TEST_TIMING_SPANS === '1';
+
+    // Shared collector. The instrumented helpers (running in the sandbox realm) read
+    // `globalThis.__e2eTimings`; this env (host realm) reads `this.global.__e2eTimings` — they are the
+    // same object. `current` is the full name of the test currently running (null during
+    // beforeAll/afterAll), used to attribute each span to a test or to the suite-scoped line. `spans`
+    // accumulates every recorded span until flush. The `start`/`end` timestamps come from the sandbox
+    // realm's `performance.now()`; Node's perf clock is process-wide (one monotonic origin across all
+    // realms), so merging the intervals here for `busyMs` is valid.
+    this.collector = { current: null, spans: [] };
+    this.global.__e2eTimings = this.collector;
+
+    // Records to flush as JSONL on teardown. One per test plus one suite-scoped record.
+    this.records = [];
+    this.recordsByName = new Map();
+    // The suite-scoped record (beforeAll/afterAll); bodyMs/totalMs are computed at flush time.
+    this.suiteRecord = { name: null, status: 'passed', beforeHooksMs: 0, afterHooksMs: 0 };
+    // Start times keyed by hook type (beforeAll/beforeEach/afterAll/afterEach).
+    this.hookStarts = {};
+    this.testStarts = new Map();
+    // Guards against double-flushing (we flush on both the teardown event and the teardown method).
+    this.flushed = false;
+    // Raw per-occurrence spans retained only when TEST_TIMING_SPANS=1, emitted as `type:"span"` lines.
+    this.rawSpans = [];
+  }
+
+  // Flush on the teardown method too: the jest lifecycle always calls this, whereas the 'teardown'
+  // event is not reliably delivered when a test run is interrupted. Whichever fires first wins.
+  async teardown() {
+    this.finalizeAndFlush();
+    await super.teardown();
+  }
+
+  async handleTestEvent(event, state) {
+    // Run the base env first so its unhandledRejection patching for after-hooks stays intact.
+    await super.handleTestEvent(event, state);
+
+    if (!this.timingFile) {
+      return;
+    }
+
+    const now = Date.now();
+    switch (event.name) {
+      case 'hook_start': {
+        this.hookStarts[event.hook.type] = now;
+        break;
+      }
+      case 'hook_success':
+      case 'hook_failure': {
+        const start = this.hookStarts[event.hook.type];
+        if (start === undefined) {
+          break;
+        }
+        delete this.hookStarts[event.hook.type];
+        const ms = now - start;
+        if (event.hook.type === 'beforeEach') {
+          this.addToCurrent('beforeHooksMs', ms);
+        } else if (event.hook.type === 'afterEach') {
+          this.addToCurrent('afterHooksMs', ms);
+        } else if (event.hook.type === 'beforeAll') {
+          this.suiteRecord.beforeHooksMs += ms;
+        } else if (event.hook.type === 'afterAll') {
+          this.suiteRecord.afterHooksMs += ms;
+        }
+        break;
+      }
+      case 'test_start': {
+        const name = this.fullTestName(event.test);
+        const record = {
+          type: 'test',
+          name,
+          status: 'passed',
+          setupFnMs: 0,
+          beforeHooksMs: 0,
+          bodyMs: 0,
+          teardownFnMs: 0,
+          afterHooksMs: 0,
+          totalMs: 0,
+          startedAt: new Date(now).toISOString(),
+        };
+        this.records.push(record);
+        this.recordsByName.set(name, record);
+        this.testStarts.set(name, now);
+        this.collector.current = name;
+        break;
+      }
+      case 'test_fn_start': {
+        this.testStarts.set(`fn:${this.fullTestName(event.test)}`, now);
+        break;
+      }
+      case 'test_fn_success':
+      case 'test_fn_failure': {
+        const name = this.fullTestName(event.test);
+        const start = this.testStarts.get(`fn:${name}`);
+        const record = this.recordsByName.get(name);
+        if (record && start !== undefined) {
+          record.bodyMs = now - start;
+        }
+        if (record && event.name === 'test_fn_failure') {
+          record.status = 'failed';
+        }
+        break;
+      }
+      case 'test_done': {
+        const name = this.fullTestName(event.test);
+        const record = this.recordsByName.get(name);
+        const start = this.testStarts.get(name);
+        if (record && start !== undefined) {
+          record.totalMs = now - start;
+        }
+        if (record && event.test?.errors?.length) {
+          record.status = 'failed';
+        }
+        this.collector.current = null;
+        break;
+      }
+      case 'teardown': {
+        this.finalizeAndFlush();
+        break;
+      }
+      default:
+        break;
+    }
+  }
+
+  /** The record for the test currently running, if any. */
+  currentRecord() {
+    return this.collector.current ? this.recordsByName.get(this.collector.current) : undefined;
+  }
+
+  /** Adds `ms` to a hook bucket on the current per-test record, or the suite record during beforeAll/afterAll. */
+  addToCurrent(field, ms) {
+    const record = this.currentRecord();
+    if (record) {
+      record[field] += ms;
+    }
+  }
+
+  /** Builds the full test name by walking up the describe blocks, matching jest's currentTestName. */
+  fullTestName(test) {
+    if (!test) {
+      return null;
+    }
+    const parts = [];
+    let node = test;
+    while (node && node.name && node.name !== 'ROOT_DESCRIBE_BLOCK') {
+      parts.unshift(node.name);
+      node = node.parent;
+    }
+    return parts.join(' ');
+  }
+
+  /** Folds the collected spans into the per-test / suite records via {@link foldSpansInto}, then drains the collector. */
+  foldSpans() {
+    const rawSpans = foldSpansInto(this.collector.spans, {
+      recordsByName: this.recordsByName,
+      suiteRecord: this.suiteRecord,
+      emitSpanLines: this.emitSpanLines,
+    });
+    this.rawSpans.push(...rawSpans);
+    this.collector.spans = [];
+  }
+
+  /** Folds spans into records, finalizes the suite-scoped record, and writes all JSONL. */
+  finalizeAndFlush() {
+    if (this.flushed || !this.timingFile) {
+      return;
+    }
+    this.flushed = true;
+    this.foldSpans();
+
+    const lines = [];
+    for (const record of this.records) {
+      lines.push(this.toLine({ ...record }));
+    }
+    lines.push(
+      this.toLine({
+        type: 'suite',
+        name: null,
+        status: 'passed',
+        setupFnMs: this.suiteRecord.setupFnMs ?? 0,
+        beforeHooksMs: this.suiteRecord.beforeHooksMs,
+        teardownFnMs: this.suiteRecord.teardownFnMs ?? 0,
+        afterHooksMs: this.suiteRecord.afterHooksMs,
+        totalMs: this.suiteRecord.beforeHooksMs + this.suiteRecord.afterHooksMs,
+        spans: this.suiteRecord.spans,
+      }),
+    );
+    for (const raw of this.rawSpans) {
+      lines.push(this.toLine({ type: 'span', name: raw.owner, span: raw.name, ms: raw.ms }));
+    }
+
+    const payload = lines.join('\n') + '\n';
+    try {
+      // One file per worker process, but a worker may run several suites: append so we keep them all.
+      appendFileSync(this.timingFile, payload);
+    } catch {
+      try {
+        writeFileSync(this.timingFile, payload);
+      } catch {
+        // Timing is best-effort; never fail a test because we couldn't write the file.
+      }
+    }
+  }
+
+  /** Flattens metadata onto a record and serializes to a single JSON line, rounding all Ms fields. */
+  toLine(record) {
+    const obj = { suite: this.suite, ...record, ...this.meta };
+    for (const key of Object.keys(obj)) {
+      if (key.endsWith('Ms') && typeof obj[key] === 'number') {
+        obj[key] = Math.round(obj[key]);
+      }
+    }
+    // Round the nested span aggregates' Ms fields too (totalMs/busyMs/maxMs per tag).
+    if (obj.spans && typeof obj.spans === 'object') {
+      for (const agg of Object.values(obj.spans)) {
+        for (const key of Object.keys(agg)) {
+          if (key.endsWith('Ms') && typeof agg[key] === 'number') {
+            agg[key] = Math.round(agg[key]);
+          }
+        }
+      }
+    }
+    return JSON.stringify(obj);
+  }
+}
+
+/**
+ * Aggregates a list of `{ start, end }` spans of one tag into `{ count, totalMs, busyMs, maxMs }`.
+ * `totalMs` is the naive sum of per-occurrence durations (correct for serial repeats); `busyMs` is
+ * the duration of the union of the `[start, end)` intervals (concurrency-correct — N concurrent
+ * spans collapse to their wall-clock span instead of summing to ~N×); `maxMs` is the longest single
+ * occurrence. All durations are on the shared process-wide `performance.now()` clock. Does not mutate
+ * the input array.
+ */
+export function aggregateSpans(spans) {
+  let totalMs = 0;
+  let maxMs = 0;
+  for (const span of spans) {
+    const ms = span.end - span.start;
+    totalMs += ms;
+    if (ms > maxMs) {
+      maxMs = ms;
+    }
+  }
+  // busyMs: sort by start and merge overlapping intervals, summing the merged lengths.
+  const sorted = [...spans].sort((a, b) => a.start - b.start);
+  let busyMs = 0;
+  let mergeStart = null;
+  let mergeEnd = null;
+  for (const span of sorted) {
+    if (mergeStart === null) {
+      mergeStart = span.start;
+      mergeEnd = span.end;
+    } else if (span.start <= mergeEnd) {
+      if (span.end > mergeEnd) {
+        mergeEnd = span.end;
+      }
+    } else {
+      busyMs += mergeEnd - mergeStart;
+      mergeStart = span.start;
+      mergeEnd = span.end;
+    }
+  }
+  if (mergeStart !== null) {
+    busyMs += mergeEnd - mergeStart;
+  }
+  return { count: spans.length, totalMs, busyMs, maxMs };
+}
+
+/**
+ * Groups `spans` by owner, then by tag, and attaches a `spans` map of aggregates to the matching
+ * per-test record (looked up in `recordsByName`) or, for `null`-owner spans, to `suiteRecord`. Also
+ * derives the back-compat `setupFnMs` / `teardownFnMs` fields from the `setup:env:<mode>` /
+ * `teardown:env` aggregates (additive over any pre-existing value). Spans whose owner matches no
+ * record (e.g. the pinned `other:mempool-feeder` owner) are dropped. Returns the raw per-occurrence
+ * list (`{ owner, name, ms }`), which is empty unless `emitSpanLines` is set.
+ */
+export function foldSpansInto(spans, { recordsByName, suiteRecord, emitSpanLines = false }) {
+  const rawSpans = [];
+  // Group by owner (null owner → suite record), then by tag.
+  const byOwner = new Map();
+  for (const span of spans) {
+    const owner = span.owner ?? null;
+    let byName = byOwner.get(owner);
+    if (!byName) {
+      byName = new Map();
+      byOwner.set(owner, byName);
+    }
+    let list = byName.get(span.name);
+    if (!list) {
+      list = [];
+      byName.set(span.name, list);
+    }
+    list.push(span);
+    if (emitSpanLines) {
+      rawSpans.push({ owner, name: span.name, ms: span.end - span.start });
+    }
+  }
+
+  for (const [owner, byName] of byOwner) {
+    const record = owner === null ? suiteRecord : recordsByName.get(owner);
+    if (!record) {
+      continue;
+    }
+    const aggregates = {};
+    for (const [name, list] of byName) {
+      aggregates[name] = aggregateSpans(list);
+    }
+    record.spans = aggregates;
+    // Back-compat: the legacy fn-span buckets are derived from the top-level env setup/teardown tags.
+    // The top-level setup is tagged `setup:env:<proverMode>` (the mode is a closed set, distinct from
+    // the `:anvil`/`:l1-deploy`/... sub-phase tags), so sum those three exact tags.
+    const setupFnMs =
+      (aggregates['setup:env:none']?.totalMs ?? 0) +
+      (aggregates['setup:env:fake']?.totalMs ?? 0) +
+      (aggregates['setup:env:real']?.totalMs ?? 0);
+    record.setupFnMs = setupFnMs + (record.setupFnMs ?? 0);
+    record.teardownFnMs = (aggregates['teardown:env']?.totalMs ?? 0) + (record.teardownFnMs ?? 0);
+  }
+  return rawSpans;
+}

package/src/shared/uniswap_l1_l2.ts

@@ -11,18 +11,18 @@ 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, 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 +32,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,8 +84,6 @@ export const uniswapL1L2TestSuite = (
 
       l1Client = deployL1ContractsValues.l1Client;
 
-      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');
       }
@@ -97,7 +95,7 @@ export const uniswapL1L2TestSuite = (
       version = Number(await rollup.getVersion());
       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 +128,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 +152,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 +194,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 +248,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 +262,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 +293,7 @@ export const uniswapL1L2TestSuite = (
 
       const swapPrivateMessageMetadata = {
         _epoch: BigInt(epoch),
+        _numCheckpointsInEpoch: BigInt(numCheckpointsInEpoch),
         _leafIndex: BigInt(swapPrivateL2MessageIndex),
         _path: swapPrivateSiblingPath
           .toBufferArray()
@@ -336,10 +351,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 +585,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 +799,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 +850,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 +893,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 +943,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 +999,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 +1042,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!');