@aztec/end-to-end 0.0.0-test.1 → 0.0.1-commit.017a351

package/src/shared/uniswap_l1_l2.ts

@@ -1,30 +1,32 @@
-import {
-  type AccountWallet,
-  AztecAddress,
-  type AztecNode,
-  type CheatCodes,
-  EthAddress,
-  Fr,
-  type Logger,
-  type PXE,
-  computeAuthWitMessageHash,
-  generateClaimSecret,
-} from '@aztec/aztec.js';
-import {
-  type DeployL1ContractsReturnType,
-  type ViemPublicClient,
-  type ViemWalletClient,
-  deployL1Contract,
-  extractEvent,
-} from '@aztec/ethereum';
-import { sha256ToField } from '@aztec/foundation/crypto';
-import { InboxAbi, RollupAbi, UniswapPortalAbi, UniswapPortalBytecode } from '@aztec/l1-artifacts';
+import { AztecAddress, EthAddress } from '@aztec/aztec.js/addresses';
+import { computeAuthWitMessageHash } from '@aztec/aztec.js/authorization';
+import { waitForProven } from '@aztec/aztec.js/contracts';
+import { generateClaimSecret } from '@aztec/aztec.js/ethereum';
+import { Fr } from '@aztec/aztec.js/fields';
+import type { Logger } from '@aztec/aztec.js/log';
+import type { AztecNode } from '@aztec/aztec.js/node';
+import { CheatCodes } from '@aztec/aztec/testing';
+import { RollupContract } from '@aztec/ethereum/contracts';
+import type { DeployAztecL1ContractsReturnType } from '@aztec/ethereum/deploy-aztec-l1-contracts';
+import { deployL1Contract } from '@aztec/ethereum/deploy-l1-contract';
+import type { ExtendedViemWalletClient } from '@aztec/ethereum/types';
+import { extractEvent } from '@aztec/ethereum/utils';
+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 { jest } from '@jest/globals';
 import { type GetContractReturnType, getContract, parseEther, toFunctionSelector } from 'viem';
 
-import { ensureAccountsPubliclyDeployed } 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
@@ -34,38 +36,13 @@ 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`
 
-// docs:start:uniswap_l1_l2_test_setup_const
-const TIMEOUT = 360_000;
-
-/** Objects to be returned by the uniswap setup function */
-export type UniswapSetupContext = {
-  /** Aztec Node instance */
-  aztecNode: AztecNode;
-  /** The Private eXecution Environment (PXE). */
-  pxe: PXE;
-  /** Logger instance named as the current test. */
-  logger: Logger;
-  /** Viem Public client instance. */
-  publicClient: ViemPublicClient;
-  /** Viem Wallet Client instance. */
-  walletClient: ViemWalletClient;
-  /** The owner wallet. */
-  ownerWallet: AccountWallet;
-  /** The sponsor wallet. */
-  sponsorWallet: AccountWallet;
-  /**  */
-  deployL1ContractsValues: DeployL1ContractsReturnType;
-  /** Cheat codes instance. */
-  cheatCodes: CheatCodes;
-};
-// docs:end:uniswap_l1_l2_test_setup_const
+const TIMEOUT = 15 * 60 * 1000;
 
 export const uniswapL1L2TestSuite = (
-  setup: () => Promise<UniswapSetupContext>,
+  setup: () => Promise<EndToEndContext>,
   cleanup: () => Promise<void>,
   expectedForkBlockNumber = 17514288,
 ) => {
-  // docs:start:uniswap_l1_l2_test_beforeAll
   describe('uniswap_trade_on_l1_from_l2', () => {
     jest.setTimeout(TIMEOUT);
 
@@ -73,25 +50,22 @@ export const uniswapL1L2TestSuite = (
     const DAI_ADDRESS: EthAddress = EthAddress.fromString('0x6B175474E89094C44Da98b954EedeAC495271d0F');
 
     let aztecNode: AztecNode;
-    let pxe: PXE;
     let logger: Logger;
 
-    let walletClient: ViemWalletClient;
-    let publicClient: ViemPublicClient;
+    let l1Client: ExtendedViemWalletClient;
 
-    let ownerWallet: AccountWallet;
+    let wallet: TestWallet;
     let ownerAddress: AztecAddress;
     let ownerEthAddress: EthAddress;
     // does transactions on behalf of owner on Aztec:
-    let sponsorWallet: AccountWallet;
-    // let sponsorAddress: AztecAddress;
+    let sponsorAddress: AztecAddress;
 
     let daiCrossChainHarness: CrossChainTestHarness;
     let wethCrossChainHarness: CrossChainTestHarness;
 
-    let deployL1ContractsValues: DeployL1ContractsReturnType;
-    let rollup: GetContractReturnType<typeof RollupAbi, ViemWalletClient>;
-    let uniswapPortal: GetContractReturnType<typeof UniswapPortalAbi, ViemWalletClient>;
+    let deployL1ContractsValues: DeployAztecL1ContractsReturnType;
+    let rollup: RollupContract;
+    let uniswapPortal: GetContractReturnType<typeof UniswapPortalAbi, ExtendedViemWalletClient>;
     let uniswapPortalAddress: EthAddress;
     let uniswapL2Contract: UniswapContract;
 
@@ -100,43 +74,42 @@ export const uniswapL1L2TestSuite = (
     const minimumOutputAmount = 0n;
 
     let cheatCodes: CheatCodes;
-
+    let version: number;
     beforeAll(async () => {
+      const t = await setup();
       ({
         aztecNode,
-        pxe,
         logger,
-        publicClient,
-        walletClient,
-        ownerWallet,
-        sponsorWallet,
         deployL1ContractsValues,
         cheatCodes,
-      } = await setup());
+        wallet,
+        accounts: [ownerAddress, sponsorAddress],
+      } = t);
+
+      l1Client = deployL1ContractsValues.l1Client;
 
-      if (Number(await publicClient.getBlockNumber()) < expectedForkBlockNumber) {
+      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');
       }
 
-      rollup = getContract({
-        address: deployL1ContractsValues.l1ContractAddresses.rollupAddress.toString(),
-        abi: RollupAbi,
-        client: deployL1ContractsValues.walletClient,
-      });
-
-      ownerAddress = ownerWallet.getAddress();
-      // sponsorAddress = sponsorWallet.getAddress();
-      ownerEthAddress = EthAddress.fromString((await walletClient.getAddresses())[0]);
+      rollup = new RollupContract(
+        deployL1ContractsValues.l1Client,
+        deployL1ContractsValues.l1ContractAddresses.rollupAddress,
+      );
+      version = Number(await rollup.getVersion());
+      ownerEthAddress = EthAddress.fromString((await l1Client.getAddresses())[0]);
 
-      await ensureAccountsPubliclyDeployed(ownerWallet, [ownerWallet, sponsorWallet]);
+      await ensureAccountContractsPublished(wallet, [ownerAddress, sponsorAddress]);
+      await ensureAuthRegistryPublished(wallet, ownerAddress);
 
       logger.info('Deploying DAI Portal, initializing and deploying l2 contract...');
       daiCrossChainHarness = await CrossChainTestHarness.new(
         aztecNode,
-        pxe,
-        publicClient,
-        walletClient,
-        ownerWallet,
+        deployL1ContractsValues.l1Client,
+        wallet,
+        ownerAddress,
         logger,
         DAI_ADDRESS,
       );
@@ -144,31 +117,29 @@ export const uniswapL1L2TestSuite = (
       logger.info('Deploying WETH Portal, initializing and deploying l2 contract...');
       wethCrossChainHarness = await CrossChainTestHarness.new(
         aztecNode,
-        pxe,
-        publicClient,
-        walletClient,
-        ownerWallet,
+        l1Client,
+        wallet,
+        ownerAddress,
         logger,
         WETH9_ADDRESS,
       );
 
       logger.info('Deploy Uniswap portal on L1 and L2...');
-      uniswapPortalAddress = await deployL1Contract(
-        walletClient,
-        publicClient,
-        UniswapPortalAbi,
-        UniswapPortalBytecode,
-      ).then(({ address }) => address);
+      uniswapPortalAddress = await deployL1Contract(l1Client, UniswapPortalAbi, UniswapPortalBytecode).then(
+        ({ address }) => address,
+      );
 
       uniswapPortal = getContract({
         address: uniswapPortalAddress.toString(),
         abi: UniswapPortalAbi,
-        client: walletClient,
+        client: l1Client,
       });
       // deploy l2 uniswap contract and attach to portal
-      uniswapL2Contract = await UniswapContract.deploy(ownerWallet, uniswapPortalAddress).send().deployed();
+      ({ contract: uniswapL2Contract } = await UniswapContract.deploy(wallet, uniswapPortalAddress).send({
+        from: ownerAddress,
+      }));
 
-      const registryAddress = (await pxe.getNodeInfo()).l1ContractAddresses.registryAddress;
+      const registryAddress = (await aztecNode.getNodeInfo()).l1ContractAddresses.registryAddress;
 
       await uniswapPortal.write.initialize(
         [registryAddress.toString(), uniswapL2Contract.address.toString()],
@@ -177,19 +148,17 @@ export const uniswapL1L2TestSuite = (
 
       // Give me some WETH so I can deposit to L2 and do the swap...
       logger.info('Getting some weth');
-      const hash = await walletClient.sendTransaction({ to: WETH9_ADDRESS.toString(), value: parseEther('1000') });
-      await publicClient.waitForTransactionReceipt({ hash });
+      const hash = await l1Client.sendTransaction({ to: WETH9_ADDRESS.toString(), value: parseEther('1000') });
+      await l1Client.waitForTransactionReceipt({ hash });
 
       const wethBalance = await wethCrossChainHarness.getL1BalanceOf(ownerEthAddress);
       expect(wethBalance).toBe(parseEther('1000'));
     });
-    // docs:end:uniswap_l1_l2_test_beforeAll
 
     afterAll(async () => {
       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);
 
@@ -218,7 +187,7 @@ export const uniswapL1L2TestSuite = (
       // 3. Owner gives uniswap approval to transfer the funds to public to self on its behalf
       logger.info('Approving uniswap to transfer funds to public to self on my behalf');
       const nonceForWETHTransferToPublicApproval = new Fr(1n);
-      await ownerWallet.createAuthWit({
+      const transferToPublicAuhtwit = await wallet.createAuthWit(ownerAddress, {
         caller: uniswapL2Contract.address,
         action: wethCrossChainHarness.l2Token.methods.transfer_to_public(
           ownerAddress,
@@ -232,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,
@@ -244,8 +213,7 @@ export const uniswapL1L2TestSuite = (
           secretHashForDepositingSwappedDai,
           ownerEthAddress,
         )
-        .send()
-        .wait();
+        .send({ from: ownerAddress, authWitnesses: [transferToPublicAuhtwit] });
 
       const swapPrivateFunction = 'swap_private(address,uint256,uint24,address,uint256,bytes32,address)';
       const swapPrivateContent = sha256ToField([
@@ -259,13 +227,13 @@ export const uniswapL1L2TestSuite = (
         ownerEthAddress.toBuffer32(),
       ]);
 
-      const swapPrivateLeaf = sha256ToField([
-        uniswapL2Contract.address,
-        new Fr(1), // aztec version
-        EthAddress.fromString(uniswapPortal.address).toBuffer32(),
-        new Fr(publicClient.chain.id), // chain id
-        swapPrivateContent,
-      ]);
+      const swapPrivateLeaf = computeL2ToL1MessageHash({
+        l2Sender: uniswapL2Contract.address,
+        l1Recipient: EthAddress.fromString(uniswapPortal.address),
+        content: swapPrivateContent,
+        rollupVersion: new Fr(version),
+        chainId: new Fr(l1Client.chain.id),
+      });
 
       const withdrawContent = sha256ToField([
         Buffer.from(toFunctionSelector('withdraw(address,uint256,address)').substring(2), 'hex'),
@@ -274,39 +242,56 @@ export const uniswapL1L2TestSuite = (
         uniswapPortalAddress.toBuffer32(),
       ]);
 
-      const withdrawLeaf = sha256ToField([
-        wethCrossChainHarness.l2Bridge.address,
-        new Fr(1), // aztec version
-        wethCrossChainHarness.tokenPortalAddress.toBuffer32(),
-        new Fr(publicClient.chain.id), // chain id
-        withdrawContent,
-      ]);
+      const withdrawLeaf = computeL2ToL1MessageHash({
+        l2Sender: wethCrossChainHarness.l2Bridge.address,
+        l1Recipient: wethCrossChainHarness.tokenPortalAddress,
+        content: withdrawContent,
+        rollupVersion: new Fr(version),
+        chainId: new Fr(l1Client.chain.id),
+      });
 
       // ensure that user's funds were burnt
       await wethCrossChainHarness.expectPrivateBalanceOnL2(ownerAddress, wethL2BalanceBeforeSwap - wethAmountToBridge);
       // ensure that uniswap contract didn't eat the funds.
       await wethCrossChainHarness.expectPublicBalanceOnL2(uniswapL2Contract.address, 0n);
 
-      // Since the outbox is only consumable when the block is proven, we need to set the block to be proven
-      await cheatCodes.rollup.markAsProven(await rollup.read.getPendingBlockNumber());
+      // 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()
       logger.info('Execute withdraw and swap on the uniswapPortal!');
       const daiL1BalanceOfPortalBeforeSwap = await daiCrossChainHarness.getL1BalanceOf(
         daiCrossChainHarness.tokenPortalAddress,
       );
-
-      const [swapPrivateL2MessageIndex, swapPrivateSiblingPath] = await aztecNode.getL2ToL1MessageMembershipWitness(
-        l2UniswapInteractionReceipt.blockNumber!,
-        swapPrivateLeaf,
+      const swapResult = await retryUntil(
+        () => aztecNode.getL2ToL1MembershipWitness(minedReceipt.txHash, swapPrivateLeaf),
+        'swap private l2 to l1 witness',
+        60,
+        1,
       );
-      const [withdrawL2MessageIndex, withdrawSiblingPath] = await aztecNode.getL2ToL1MessageMembershipWitness(
-        l2UniswapInteractionReceipt.blockNumber!,
-        withdrawLeaf,
+      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 withdrawL2MessageIndex = withdrawResult.leafIndex;
+      const withdrawSiblingPath = withdrawResult.siblingPath;
+
       const withdrawMessageMetadata = {
-        _l2BlockNumber: BigInt(l2UniswapInteractionReceipt.blockNumber!),
+        _epoch: BigInt(epoch),
+        _numCheckpointsInEpoch: BigInt(withdrawResult.numCheckpointsInEpoch),
         _leafIndex: BigInt(withdrawL2MessageIndex),
         _path: withdrawSiblingPath
           .toBufferArray()
@@ -314,7 +299,8 @@ export const uniswapL1L2TestSuite = (
       };
 
       const swapPrivateMessageMetadata = {
-        _l2BlockNumber: BigInt(l2UniswapInteractionReceipt.blockNumber!),
+        _epoch: BigInt(epoch),
+        _numCheckpointsInEpoch: BigInt(numCheckpointsInEpoch),
         _leafIndex: BigInt(swapPrivateL2MessageIndex),
         _path: swapPrivateSiblingPath
           .toBufferArray()
@@ -333,7 +319,7 @@ export const uniswapL1L2TestSuite = (
       ] as const;
 
       // this should also insert a message into the inbox.
-      const txReceipt = await daiCrossChainHarness.publicClient.waitForTransactionReceipt({
+      const txReceipt = await daiCrossChainHarness.l1Client.waitForTransactionReceipt({
         hash: await uniswapPortal.write.swapPrivate(swapArgs),
       });
 
@@ -372,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);
 
@@ -461,10 +445,10 @@ export const uniswapL1L2TestSuite = (
     //       ownerEthAddress,
     //       nonceForSwap,
     //     );
-    //   await ownerWallet.setPublicAuthWit({ caller: sponsorAddress, action }, true).send().wait();
+    //   await ownerWallet.setPublicAuthWit({ caller: sponsorAddress, action }, true).send();
 
     //   // 4.2 Call swap_public from user2 on behalf of owner
-    //   const uniswapL2Interaction = await action.send().wait();
+    //   const uniswapL2Interaction = await action.send();
 
     //   const swapPublicContent = sha256ToField([
     //     Buffer.from(
@@ -525,7 +509,7 @@ export const uniswapL1L2TestSuite = (
     //   );
 
     //   const withdrawMessageMetadata = {
-    //     _l2BlockNumber: BigInt(uniswapL2Interaction.blockNumber!),
+    //     _epoch: epoch,
     //     _leafIndex: BigInt(withdrawL2MessageIndex),
     //     _path: withdrawSiblingPath
     //       .toBufferArray()
@@ -533,7 +517,7 @@ export const uniswapL1L2TestSuite = (
     //   };
 
     //   const swapPrivateMessageMetadata = {
-    //     _l2BlockNumber: BigInt(uniswapL2Interaction.blockNumber!),
+    //     _epoch: epoch,
     //     _leafIndex: BigInt(swapPrivateL2MessageIndex),
     //     _path: swapPrivateSiblingPath
     //       .toBufferArray()
@@ -608,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.
@@ -620,14 +603,16 @@ export const uniswapL1L2TestSuite = (
       const expectedMessageHash = await computeAuthWitMessageHash(
         {
           caller: uniswapL2Contract.address,
-          action: wethCrossChainHarness.l2Token.methods.transfer_to_public(
-            ownerAddress,
-            uniswapL2Contract.address,
-            wethAmountToBridge,
-            nonceForWETHTransferToPublicApproval,
-          ),
+          call: await wethCrossChainHarness.l2Token.methods
+            .transfer_to_public(
+              ownerAddress,
+              uniswapL2Contract.address,
+              wethAmountToBridge,
+              nonceForWETHTransferToPublicApproval,
+            )
+            .getFunctionCall(),
         },
-        { chainId: ownerWallet.getChainId(), version: ownerWallet.getVersion() },
+        await wallet.getChainInfo(),
       );
 
       await expect(
@@ -643,7 +628,7 @@ export const uniswapL1L2TestSuite = (
             Fr.random(),
             ownerEthAddress,
           )
-          .prove(),
+          .simulate({ from: ownerAddress }),
       ).rejects.toThrow(`Unknown auth witness for message hash ${expectedMessageHash.toString()}`);
     });
 
@@ -655,7 +640,7 @@ export const uniswapL1L2TestSuite = (
       // 2. owner gives uniswap approval to transfer the funds to public:
       logger.info('Approving uniswap to transfer funds to public to self on my behalf');
       const nonceForWETHTransferToPublicApproval = new Fr(3n);
-      await ownerWallet.createAuthWit({
+      const transferToPublicAuthwith = await wallet.createAuthWit(ownerAddress, {
         caller: uniswapL2Contract.address,
         action: wethCrossChainHarness.l2Token.methods.transfer_to_public(
           ownerAddress,
@@ -680,7 +665,7 @@ export const uniswapL1L2TestSuite = (
             Fr.random(),
             ownerEthAddress,
           )
-          .prove(),
+          .send({ from: ownerAddress, authWitnesses: [transferToPublicAuthwith] }),
       ).rejects.toThrow('Assertion failed: input_asset address is not the same as seen in the bridge contract');
     });
 
@@ -692,7 +677,8 @@ export const uniswapL1L2TestSuite = (
 
       // 2. Give approval to uniswap to transfer funds to itself
       const nonceForWETHTransferApproval = new Fr(2n);
-      const validateActionInteraction = await ownerWallet.setPublicAuthWit(
+      const validateActionInteraction = await wallet.setPublicAuthWit(
+        ownerAddress,
         {
           caller: uniswapL2Contract.address,
           action: wethCrossChainHarness.l2Token.methods.transfer_in_public(
@@ -704,7 +690,7 @@ export const uniswapL1L2TestSuite = (
         },
         true,
       );
-      await validateActionInteraction.send().wait();
+      await validateActionInteraction.send();
 
       // No approval to call `swap` but should work even without it:
       const [_, secretHashForDepositingSwappedDai] = await generateClaimSecret();
@@ -723,8 +709,7 @@ export const uniswapL1L2TestSuite = (
           ownerEthAddress,
           Fr.ZERO, // nonce for swap -> doesn't matter
         )
-        .send()
-        .wait();
+        .send({ from: ownerAddress });
       // check weth balance of owner on L2 (we first bridged `wethAmountToBridge` into L2 and now withdrew it!)
       await wethCrossChainHarness.expectPublicBalanceOnL2(ownerAddress, 0n);
     });
@@ -736,32 +721,35 @@ export const uniswapL1L2TestSuite = (
       const nonceForWETHTransferApproval = new Fr(3n);
       const nonceForSwap = new Fr(3n);
       const secretHashForDepositingSwappedDai = new Fr(4n);
-      const action = uniswapL2Contract
-        .withWallet(sponsorWallet)
-        .methods.swap_public(
-          ownerAddress,
-          wethCrossChainHarness.l2Bridge.address,
-          wethAmountToBridge,
-          daiCrossChainHarness.l2Bridge.address,
-          nonceForWETHTransferApproval,
-          uniswapFeeTier,
-          minimumOutputAmount,
-          ownerAddress,
-          secretHashForDepositingSwappedDai,
-          ownerEthAddress,
-          nonceForSwap,
-        );
-      const validateActionInteraction = await ownerWallet.setPublicAuthWit({ caller: approvedUser, action }, true);
-      await validateActionInteraction.send().wait();
+      const action = uniswapL2Contract.methods.swap_public(
+        ownerAddress,
+        wethCrossChainHarness.l2Bridge.address,
+        wethAmountToBridge,
+        daiCrossChainHarness.l2Bridge.address,
+        nonceForWETHTransferApproval,
+        uniswapFeeTier,
+        minimumOutputAmount,
+        ownerAddress,
+        secretHashForDepositingSwappedDai,
+        ownerEthAddress,
+        nonceForSwap,
+      );
+      const validateActionInteraction = await wallet.setPublicAuthWit(
+        ownerAddress,
+        { caller: approvedUser, action },
+        true,
+      );
+      await validateActionInteraction.send();
 
-      await expect(action.simulate()).rejects.toThrow(/unauthorized/);
+      await expect(action.simulate({ from: sponsorAddress })).rejects.toThrow(/unauthorized/);
     });
 
     it("uniswap can't pull funds without transfer approval", async () => {
       // swap should fail since no transfer approval to uniswap:
       const nonceForWETHTransferApproval = new Fr(4n);
 
-      const validateActionInteraction = await ownerWallet.setPublicAuthWit(
+      const validateActionInteraction = await wallet.setPublicAuthWit(
+        ownerAddress,
         {
           caller: uniswapL2Contract.address,
           action: wethCrossChainHarness.l2Token.methods.transfer_in_public(
@@ -773,7 +761,7 @@ export const uniswapL1L2TestSuite = (
         },
         true,
       );
-      await validateActionInteraction.send().wait();
+      await validateActionInteraction.send();
 
       await expect(
         uniswapL2Contract.methods
@@ -790,7 +778,7 @@ export const uniswapL1L2TestSuite = (
             ownerEthAddress,
             Fr.ZERO,
           )
-          .simulate(),
+          .simulate({ from: ownerAddress }),
       ).rejects.toThrow(/unauthorized/);
     });
 
@@ -803,7 +791,7 @@ export const uniswapL1L2TestSuite = (
       // Owner gives uniswap approval to transfer the funds to public to self on its behalf
       logger.info('Approving uniswap to transfer the funds to public to self on my behalf');
       const nonceForWETHTransferToPublicApproval = new Fr(4n);
-      await ownerWallet.createAuthWit({
+      const transferToPublicAuhtwit = await wallet.createAuthWit(ownerAddress, {
         caller: uniswapL2Contract.address,
         action: wethCrossChainHarness.l2Token.methods.transfer_to_public(
           ownerAddress,
@@ -818,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,
@@ -830,8 +818,7 @@ export const uniswapL1L2TestSuite = (
           secretHashForDepositingSwappedDai,
           ownerEthAddress,
         )
-        .send()
-        .wait();
+        .send({ from: ownerAddress, authWitnesses: [transferToPublicAuhtwit] });
 
       const swapPrivateContent = sha256ToField([
         Buffer.from(
@@ -847,13 +834,13 @@ export const uniswapL1L2TestSuite = (
         ownerEthAddress.toBuffer32(),
       ]);
 
-      const swapPrivateLeaf = sha256ToField([
-        uniswapL2Contract.address,
-        new Fr(1), // aztec version
-        EthAddress.fromString(uniswapPortal.address).toBuffer32(),
-        new Fr(publicClient.chain.id), // chain id
-        swapPrivateContent,
-      ]);
+      const swapPrivateLeaf = computeL2ToL1MessageHash({
+        l2Sender: uniswapL2Contract.address,
+        l1Recipient: EthAddress.fromString(uniswapPortal.address),
+        content: swapPrivateContent,
+        rollupVersion: new Fr(version),
+        chainId: new Fr(l1Client.chain.id),
+      });
 
       const withdrawContent = sha256ToField([
         Buffer.from(toFunctionSelector('withdraw(address,uint256,address)').substring(2), 'hex'),
@@ -862,25 +849,49 @@ export const uniswapL1L2TestSuite = (
         uniswapPortalAddress.toBuffer32(),
       ]);
 
-      const withdrawLeaf = sha256ToField([
-        wethCrossChainHarness.l2Bridge.address,
-        new Fr(1), // aztec version
-        wethCrossChainHarness.tokenPortalAddress.toBuffer32(),
-        new Fr(publicClient.chain.id), // chain id
-        withdrawContent,
-      ]);
+      const withdrawLeaf = computeL2ToL1MessageHash({
+        l2Sender: wethCrossChainHarness.l2Bridge.address,
+        l1Recipient: wethCrossChainHarness.tokenPortalAddress,
+        content: withdrawContent,
+        rollupVersion: new Fr(version),
+        chainId: new Fr(l1Client.chain.id),
+      });
+
+      // ensure that user's funds were burnt
+      await wethCrossChainHarness.expectPrivateBalanceOnL2(ownerAddress, wethL2BalanceBeforeSwap - wethAmountToBridge);
 
-      const [swapPrivateL2MessageIndex, swapPrivateSiblingPath] = await aztecNode.getL2ToL1MessageMembershipWitness(
-        withdrawReceipt.blockNumber!,
-        swapPrivateLeaf,
+      // 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 [withdrawL2MessageIndex, withdrawSiblingPath] = await aztecNode.getL2ToL1MessageMembershipWitness(
-        withdrawReceipt.blockNumber!,
-        withdrawLeaf,
+      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 withdrawL2MessageIndex = withdrawResult.leafIndex;
+      const withdrawSiblingPath = withdrawResult.siblingPath;
+
       const withdrawMessageMetadata = {
-        _l2BlockNumber: BigInt(withdrawReceipt.blockNumber!),
+        _epoch: BigInt(epoch),
+        _numCheckpointsInEpoch: BigInt(withdrawResult.numCheckpointsInEpoch),
         _leafIndex: BigInt(withdrawL2MessageIndex),
         _path: withdrawSiblingPath
           .toBufferArray()
@@ -888,19 +899,14 @@ export const uniswapL1L2TestSuite = (
       };
 
       const swapPrivateMessageMetadata = {
-        _l2BlockNumber: BigInt(withdrawReceipt.blockNumber!),
+        _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 block is proven, we need to set the block to be proven
-      await cheatCodes.rollup.markAsProven(await rollup.read.getPendingBlockNumber());
-
       // On L1 call swap_public!
       logger.info('call swap_public on L1');
       const swapArgs = [
@@ -927,7 +933,8 @@ export const uniswapL1L2TestSuite = (
 
       // Owner gives uniswap approval to transfer funds on its behalf
       const nonceForWETHTransferApproval = new Fr(5n);
-      const validateActionInteraction = await ownerWallet.setPublicAuthWit(
+      const validateActionInteraction = await wallet.setPublicAuthWit(
+        ownerAddress,
         {
           caller: uniswapL2Contract.address,
           action: wethCrossChainHarness.l2Token.methods.transfer_in_public(
@@ -939,11 +946,11 @@ export const uniswapL1L2TestSuite = (
         },
         true,
       );
-      await validateActionInteraction.send().wait();
+      await validateActionInteraction.send();
 
       // 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,
@@ -957,8 +964,7 @@ export const uniswapL1L2TestSuite = (
           ownerEthAddress,
           Fr.ZERO,
         )
-        .send()
-        .wait();
+        .send({ from: ownerAddress });
 
       const swapPublicContent = sha256ToField([
         Buffer.from(
@@ -977,13 +983,13 @@ export const uniswapL1L2TestSuite = (
         ownerEthAddress.toBuffer32(),
       ]);
 
-      const swapPublicLeaf = sha256ToField([
-        uniswapL2Contract.address,
-        new Fr(1), // aztec version
-        EthAddress.fromString(uniswapPortal.address).toBuffer32(),
-        new Fr(publicClient.chain.id), // chain id
-        swapPublicContent,
-      ]);
+      const swapPublicLeaf = computeL2ToL1MessageHash({
+        l2Sender: uniswapL2Contract.address,
+        l1Recipient: EthAddress.fromString(uniswapPortal.address),
+        content: swapPublicContent,
+        rollupVersion: new Fr(version),
+        chainId: new Fr(l1Client.chain.id),
+      });
 
       const withdrawContent = sha256ToField([
         Buffer.from(toFunctionSelector('withdraw(address,uint256,address)').substring(2), 'hex'),
@@ -992,25 +998,49 @@ export const uniswapL1L2TestSuite = (
         uniswapPortalAddress.toBuffer32(),
       ]);
 
-      const withdrawLeaf = sha256ToField([
-        wethCrossChainHarness.l2Bridge.address,
-        new Fr(1), // aztec version
-        wethCrossChainHarness.tokenPortalAddress.toBuffer32(),
-        new Fr(publicClient.chain.id), // chain id
-        withdrawContent,
-      ]);
+      const withdrawLeaf = computeL2ToL1MessageHash({
+        l2Sender: wethCrossChainHarness.l2Bridge.address,
+        l1Recipient: wethCrossChainHarness.tokenPortalAddress,
+        content: withdrawContent,
+        rollupVersion: new Fr(version),
+        chainId: new Fr(l1Client.chain.id),
+      });
 
-      const [swapPublicL2MessageIndex, swapPublicSiblingPath] = await aztecNode.getL2ToL1MessageMembershipWitness(
-        withdrawReceipt.blockNumber!,
-        swapPublicLeaf,
+      // 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 [withdrawL2MessageIndex, withdrawSiblingPath] = await aztecNode.getL2ToL1MessageMembershipWitness(
-        withdrawReceipt.blockNumber!,
-        withdrawLeaf,
+      const { epochNumber: epoch, numCheckpointsInEpoch } = swapResult;
+      const withdrawResult = await retryUntil(
+        () => aztecNode.getL2ToL1MembershipWitness(minedReceipt.txHash, withdrawLeaf),
+        'withdraw l2 to l1 witness',
+        60,
+        1,
       );
 
+      const swapPublicL2MessageIndex = swapResult.leafIndex;
+      const swapPublicSiblingPath = swapResult.siblingPath;
+
+      const withdrawL2MessageIndex = withdrawResult.leafIndex;
+      const withdrawSiblingPath = withdrawResult.siblingPath;
+
       const withdrawMessageMetadata = {
-        _l2BlockNumber: BigInt(withdrawReceipt.blockNumber!),
+        _epoch: BigInt(epoch),
+        _numCheckpointsInEpoch: BigInt(withdrawResult.numCheckpointsInEpoch),
         _leafIndex: BigInt(withdrawL2MessageIndex),
         _path: withdrawSiblingPath
           .toBufferArray()
@@ -1018,19 +1048,14 @@ export const uniswapL1L2TestSuite = (
       };
 
       const swapPublicMessageMetadata = {
-        _l2BlockNumber: BigInt(withdrawReceipt.blockNumber!),
+        _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 block is proven, we need to set the block to be proven
-      await cheatCodes.rollup.markAsProven(await rollup.read.getPendingBlockNumber());
-
       // Call swap_private on L1
       logger.info('Execute withdraw and swap on the uniswapPortal!');