@aztec/end-to-end 0.0.0-test.1 → 0.0.1-commit.1142ef1

package/src/shared/uniswap_l1_l2.ts

@@ -1,30 +1,28 @@
-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 { CheckpointNumber, EpochNumber } from '@aztec/foundation/branded-types';
+import { sha256ToField } from '@aztec/foundation/crypto/sha256';
+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 { ensureAccountsPubliclyDeployed } from '../fixtures/utils.js';
+import { type EndToEndContext, ensureAccountContractsPublished } from '../fixtures/utils.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 +32,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
-
 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 +46,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 +70,41 @@ 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;
+
+      t.watcher?.setIsMarkingAsProven(false);
 
-      if (Number(await publicClient.getBlockNumber()) < expectedForkBlockNumber) {
+      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]);
 
       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 +112,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();
+      uniswapL2Contract = await UniswapContract.deploy(wallet, uniswapPortalAddress)
+        .send({ from: ownerAddress })
+        .deployed();
 
-      const registryAddress = (await pxe.getNodeInfo()).l1ContractAddresses.registryAddress;
+      const registryAddress = (await aztecNode.getNodeInfo()).l1ContractAddresses.registryAddress;
 
       await uniswapPortal.write.initialize(
         [registryAddress.toString(), uniswapL2Contract.address.toString()],
@@ -177,13 +143,12 @@ 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();
@@ -218,7 +183,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,
@@ -244,7 +209,7 @@ export const uniswapL1L2TestSuite = (
           secretHashForDepositingSwappedDai,
           ownerEthAddress,
         )
-        .send()
+        .send({ from: ownerAddress, authWitnesses: [transferToPublicAuhtwit] })
         .wait();
 
       const swapPrivateFunction = 'swap_private(address,uint256,uint24,address,uint256,bytes32,address)';
@@ -259,13 +224,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 +239,41 @@ 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.
+      const checkpointNumber = CheckpointNumber.fromBlockNumber(l2UniswapInteractionReceipt.blockNumber!);
+      const epoch = await rollup.getEpochNumberForCheckpoint(checkpointNumber);
+      await cheatCodes.rollup.advanceToEpoch(EpochNumber(epoch + 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 swapResult = await computeL2ToL1MembershipWitness(aztecNode, epoch, swapPrivateLeaf);
+      const withdrawResult = await computeL2ToL1MembershipWitness(aztecNode, epoch, withdrawLeaf);
 
-      const [swapPrivateL2MessageIndex, swapPrivateSiblingPath] = await aztecNode.getL2ToL1MessageMembershipWitness(
-        l2UniswapInteractionReceipt.blockNumber!,
-        swapPrivateLeaf,
-      );
-      const [withdrawL2MessageIndex, withdrawSiblingPath] = await aztecNode.getL2ToL1MessageMembershipWitness(
-        l2UniswapInteractionReceipt.blockNumber!,
-        withdrawLeaf,
-      );
+      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),
         _leafIndex: BigInt(withdrawL2MessageIndex),
         _path: withdrawSiblingPath
           .toBufferArray()
@@ -314,7 +281,7 @@ export const uniswapL1L2TestSuite = (
       };
 
       const swapPrivateMessageMetadata = {
-        _l2BlockNumber: BigInt(l2UniswapInteractionReceipt.blockNumber!),
+        _epoch: BigInt(epoch),
         _leafIndex: BigInt(swapPrivateL2MessageIndex),
         _path: swapPrivateSiblingPath
           .toBufferArray()
@@ -333,7 +300,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),
       });
 
@@ -525,7 +492,7 @@ export const uniswapL1L2TestSuite = (
     //   );
 
     //   const withdrawMessageMetadata = {
-    //     _l2BlockNumber: BigInt(uniswapL2Interaction.blockNumber!),
+    //     _epoch: epoch,
     //     _leafIndex: BigInt(withdrawL2MessageIndex),
     //     _path: withdrawSiblingPath
     //       .toBufferArray()
@@ -533,7 +500,7 @@ export const uniswapL1L2TestSuite = (
     //   };
 
     //   const swapPrivateMessageMetadata = {
-    //     _l2BlockNumber: BigInt(uniswapL2Interaction.blockNumber!),
+    //     _epoch: epoch,
     //     _leafIndex: BigInt(swapPrivateL2MessageIndex),
     //     _path: swapPrivateSiblingPath
     //       .toBufferArray()
@@ -620,14 +587,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 +612,7 @@ export const uniswapL1L2TestSuite = (
             Fr.random(),
             ownerEthAddress,
           )
-          .prove(),
+          .simulate({ from: ownerAddress }),
       ).rejects.toThrow(`Unknown auth witness for message hash ${expectedMessageHash.toString()}`);
     });
 
@@ -655,7 +624,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 +649,8 @@ export const uniswapL1L2TestSuite = (
             Fr.random(),
             ownerEthAddress,
           )
-          .prove(),
+          .send({ from: ownerAddress, authWitnesses: [transferToPublicAuthwith] })
+          .wait(),
       ).rejects.toThrow('Assertion failed: input_asset address is not the same as seen in the bridge contract');
     });
 
@@ -692,7 +662,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(
@@ -723,7 +694,7 @@ export const uniswapL1L2TestSuite = (
           ownerEthAddress,
           Fr.ZERO, // nonce for swap -> doesn't matter
         )
-        .send()
+        .send({ from: ownerAddress })
         .wait();
       // check weth balance of owner on L2 (we first bridged `wethAmountToBridge` into L2 and now withdrew it!)
       await wethCrossChainHarness.expectPublicBalanceOnL2(ownerAddress, 0n);
@@ -736,32 +707,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);
+      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().wait();
 
-      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(
@@ -790,7 +764,7 @@ export const uniswapL1L2TestSuite = (
             ownerEthAddress,
             Fr.ZERO,
           )
-          .simulate(),
+          .simulate({ from: ownerAddress }),
       ).rejects.toThrow(/unauthorized/);
     });
 
@@ -803,7 +777,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,
@@ -830,7 +804,7 @@ export const uniswapL1L2TestSuite = (
           secretHashForDepositingSwappedDai,
           ownerEthAddress,
         )
-        .send()
+        .send({ from: ownerAddress, authWitnesses: [transferToPublicAuhtwit] })
         .wait();
 
       const swapPrivateContent = sha256ToField([
@@ -847,13 +821,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 +836,28 @@ 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 [swapPrivateL2MessageIndex, swapPrivateSiblingPath] = await aztecNode.getL2ToL1MessageMembershipWitness(
-        withdrawReceipt.blockNumber!,
-        swapPrivateLeaf,
-      );
-      const [withdrawL2MessageIndex, withdrawSiblingPath] = await aztecNode.getL2ToL1MessageMembershipWitness(
-        withdrawReceipt.blockNumber!,
-        withdrawLeaf,
+      const epoch = await rollup.getEpochNumberForCheckpoint(
+        CheckpointNumber.fromBlockNumber(withdrawReceipt.blockNumber!),
       );
+      const swapResult = await computeL2ToL1MembershipWitness(aztecNode, epoch, swapPrivateLeaf);
+      const withdrawResult = await computeL2ToL1MembershipWitness(aztecNode, epoch, withdrawLeaf);
+
+      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),
         _leafIndex: BigInt(withdrawL2MessageIndex),
         _path: withdrawSiblingPath
           .toBufferArray()
@@ -888,7 +865,7 @@ export const uniswapL1L2TestSuite = (
       };
 
       const swapPrivateMessageMetadata = {
-        _l2BlockNumber: BigInt(withdrawReceipt.blockNumber!),
+        _epoch: BigInt(epoch),
         _leafIndex: BigInt(swapPrivateL2MessageIndex),
         _path: swapPrivateSiblingPath
           .toBufferArray()
@@ -898,8 +875,9 @@ export const uniswapL1L2TestSuite = (
       // 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());
+      // 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');
@@ -927,7 +905,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(
@@ -957,7 +936,7 @@ export const uniswapL1L2TestSuite = (
           ownerEthAddress,
           Fr.ZERO,
         )
-        .send()
+        .send({ from: ownerAddress })
         .wait();
 
       const swapPublicContent = sha256ToField([
@@ -977,13 +956,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 +971,28 @@ 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,
-      );
-      const [withdrawL2MessageIndex, withdrawSiblingPath] = await aztecNode.getL2ToL1MessageMembershipWitness(
-        withdrawReceipt.blockNumber!,
-        withdrawLeaf,
+      const epoch = await rollup.getEpochNumberForCheckpoint(
+        CheckpointNumber.fromBlockNumber(withdrawReceipt.blockNumber!),
       );
+      const swapResult = await computeL2ToL1MembershipWitness(aztecNode, epoch, swapPublicLeaf);
+      const withdrawResult = await computeL2ToL1MembershipWitness(aztecNode, epoch, withdrawLeaf);
+
+      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),
         _leafIndex: BigInt(withdrawL2MessageIndex),
         _path: withdrawSiblingPath
           .toBufferArray()
@@ -1018,7 +1000,7 @@ export const uniswapL1L2TestSuite = (
       };
 
       const swapPublicMessageMetadata = {
-        _l2BlockNumber: BigInt(withdrawReceipt.blockNumber!),
+        _epoch: BigInt(epoch),
         _leafIndex: BigInt(swapPublicL2MessageIndex),
         _path: swapPublicSiblingPath
           .toBufferArray()
@@ -1028,8 +1010,9 @@ export const uniswapL1L2TestSuite = (
       // 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());
+      // 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!');