refactor: un-exposing PXE from e2e tests (#17163)

e2e tests should interact only with a wallet (and maybe a node) because
PXE is becoming only a lib used by a wallet. This PR is a step in that
direction.

# Followup work
1. Drop PXE JSON RPC server,
2. have Wallet instantiate PXE instead of having it passed in as an arg
to a constructor,
3. revisit the PXE methods exposed on the TestWallet.

Author: benesjan

yarn-project/accounts/src/testing/index.ts

@@ -18,11 +18,11 @@ import {
 } from './configuration.js';
 
 export {
-  type InitialAccountData,
   INITIAL_TEST_ACCOUNT_SALTS,
   INITIAL_TEST_ENCRYPTION_KEYS,
   INITIAL_TEST_SECRET_KEYS,
   INITIAL_TEST_SIGNING_KEYS,
+  type InitialAccountData,
 } from './configuration.js';
 
 /**
@@ -45,10 +45,11 @@ export function getInitialTestAccountsData(): Promise<InitialAccountData[]> {
 
 /**
  * Queries a PXE for it's registered accounts.
- * @param pxe - PXE instance.
+ * @param testWalletOrPxe - Test wallet or pxe instance to use to get the registered accounts.
  * @returns A set of key data for each of the initial accounts.
  */
-export async function getDeployedTestAccounts(pxe: PXE): Promise<InitialAccountData[]> {
+export async function getDeployedTestAccounts(testWalletOrPxe: { getPxe(): PXE } | PXE): Promise<InitialAccountData[]> {
+  const pxe = 'getPxe' in testWalletOrPxe ? testWalletOrPxe.getPxe() : testWalletOrPxe;
   const registeredAccounts = await pxe.getRegisteredAccounts();
   const testAccounts = await getInitialTestAccountsData();
   return testAccounts.filter(t => registeredAccounts.some(r => r.address.equals(t.address)));

yarn-project/aztec/src/cli/aztec_start_action.ts

@@ -63,8 +63,8 @@ export async function aztecStart(options: any, userLog: LogFn, debugLogger: Logg
       const { startBlobSink } = await import('./cmds/start_blob_sink.js');
       await startBlobSink(options, signalHandlers, userLog);
     } else if (options.pxe) {
-      const { startPXE } = await import('./cmds/start_pxe.js');
-      ({ config } = await startPXE(options, signalHandlers, services, userLog));
+      const { startPXEServiceGetWallet } = await import('./cmds/start_pxe.js');
+      ({ config } = await startPXEServiceGetWallet(options, services, userLog));
     } else if (options.archiver) {
       const { startArchiver } = await import('./cmds/start_archiver.js');
       ({ config } = await startArchiver(options, signalHandlers, services));

yarn-project/aztec/src/cli/cmds/start_bot.ts

@@ -1,13 +1,14 @@
 import { type BotConfig, BotRunner, botConfigMappings, getBotRunnerApiHandler } from '@aztec/bot';
 import type { NamespacedApiHandlers } from '@aztec/foundation/json-rpc/server';
 import type { LogFn } from '@aztec/foundation/log';
-import { type AztecNode, type PXE, createAztecNodeClient } from '@aztec/stdlib/interfaces/client';
+import { type AztecNode, type AztecNodeAdmin, createAztecNodeClient } from '@aztec/stdlib/interfaces/client';
 import type { TelemetryClient } from '@aztec/telemetry-client';
 import {
   getConfigEnvVars as getTelemetryClientConfig,
   initTelemetryClient,
   makeTracedFetch,
 } from '@aztec/telemetry-client';
+import { TestWallet } from '@aztec/test-wallet';
 
 import { extractRelevantOptions } from '../util.js';
 import { getVersions } from '../versioning.js';
@@ -32,28 +33,28 @@ export async function startBot(
     throw new Error('The bot requires access to a Node');
   }
 
-  const node = createAztecNodeClient(config.nodeUrl, getVersions(), fetch);
+  const aztecNode = createAztecNodeClient(config.nodeUrl, getVersions(), fetch);
 
-  // Start a PXE client that is used by the bot if required
-  let pxe: PXE | undefined;
-  if (options.pxe) {
-    const { addPXE } = await import('./start_pxe.js');
-    ({ pxe } = await addPXE(options, signalHandlers, services, userLog, { node }));
-  }
+  // Start a PXE client and get a wallet that is used by the bot if required
+  const { startPXEServiceGetWallet } = await import('./start_pxe.js');
+  const { wallet } = await startPXEServiceGetWallet(options, services, userLog, { node: aztecNode });
 
   const telemetry = initTelemetryClient(getTelemetryClientConfig());
-  await addBot(options, signalHandlers, services, { pxe, telemetry, node });
+  await addBot(options, signalHandlers, services, wallet, aztecNode, telemetry, undefined);
 }
 
 export function addBot(
   options: any,
   signalHandlers: (() => Promise<void>)[],
   services: NamespacedApiHandlers,
-  deps: { pxe?: PXE; node?: AztecNode; telemetry: TelemetryClient },
+  wallet: TestWallet,
+  aztecNode: AztecNode,
+  telemetry: TelemetryClient,
+  aztecNodeAdmin?: AztecNodeAdmin,
 ) {
   const config = extractRelevantOptions<BotConfig>(options, botConfigMappings, 'bot');
 
-  const botRunner = new BotRunner(config, deps);
+  const botRunner = new BotRunner(config, wallet, aztecNode, telemetry, aztecNodeAdmin);
   if (!config.noStart) {
     void botRunner.start(); // Do not block since bot setup takes time
   }

yarn-project/aztec/src/cli/cmds/start_node.ts

@@ -7,7 +7,7 @@ import { getPublicClient } from '@aztec/ethereum';
 import { SecretValue } from '@aztec/foundation/config';
 import type { NamespacedApiHandlers } from '@aztec/foundation/json-rpc/server';
 import type { LogFn } from '@aztec/foundation/log';
-import { AztecNodeAdminApiSchema, AztecNodeApiSchema, type PXE } from '@aztec/stdlib/interfaces/client';
+import { AztecNodeAdminApiSchema, AztecNodeApiSchema } from '@aztec/stdlib/interfaces/client';
 import { P2PApiSchema } from '@aztec/stdlib/interfaces/server';
 import {
   type TelemetryClientConfig,
@@ -128,17 +128,13 @@ export async function startNode(
   // Add node stop function to signal handlers
   signalHandlers.push(node.stop.bind(node));
 
-  // Add a PXE client that connects to this node if requested
-  let pxe: PXE | undefined;
-  if (options.pxe) {
-    const { addPXE } = await import('./start_pxe.js');
-    ({ pxe } = await addPXE(options, signalHandlers, services, userLog, { node }));
-  }
-
   // Add a txs bot if requested
   if (options.bot) {
     const { addBot } = await import('./start_bot.js');
-    await addBot(options, signalHandlers, services, { pxe, node, telemetry });
+    const { startPXEServiceGetWallet } = await import('./start_pxe.js');
+    const { wallet } = await startPXEServiceGetWallet(options, services, userLog, { node });
+
+    await addBot(options, signalHandlers, services, wallet, node, telemetry, undefined);
   }
 
   if (nodeConfig.autoUpdate !== 'disabled' && nodeConfig.autoUpdateUrl) {

yarn-project/aztec/src/cli/cmds/start_pxe.ts

@@ -1,36 +1,21 @@
 import type { NamespacedApiHandlers } from '@aztec/foundation/json-rpc/server';
 import type { LogFn } from '@aztec/foundation/log';
-import {
-  type CliPXEOptions,
-  type PXEService,
-  type PXEServiceConfig,
-  allPxeConfigMappings,
-  createPXEService,
-} from '@aztec/pxe/server';
+import { type CliPXEOptions, type PXEServiceConfig, allPxeConfigMappings, createPXEService } from '@aztec/pxe/server';
 import { type AztecNode, PXESchema, createAztecNodeClient } from '@aztec/stdlib/interfaces/client';
 import { makeTracedFetch } from '@aztec/telemetry-client';
+import { TestWallet } from '@aztec/test-wallet';
 
 import { extractRelevantOptions } from '../util.js';
 import { getVersions } from '../versioning.js';
 
-export type { PXEServiceConfig, CliPXEOptions };
+export type { CliPXEOptions, PXEServiceConfig };
 
-export async function startPXE(
+export async function startPXEServiceGetWallet(
   options: any,
-  signalHandlers: (() => Promise<void>)[],
-  services: NamespacedApiHandlers,
-  userLog: LogFn,
-): Promise<{ pxe: PXEService; config: PXEServiceConfig & CliPXEOptions }> {
-  return await addPXE(options, signalHandlers, services, userLog, {});
-}
-
-export async function addPXE(
-  options: any,
-  _signalHandlers: (() => Promise<void>)[],
   services: NamespacedApiHandlers,
   userLog: LogFn,
   deps: { node?: AztecNode } = {},
-): Promise<{ pxe: PXEService; config: PXEServiceConfig & CliPXEOptions }> {
+): Promise<{ wallet: TestWallet; config: PXEServiceConfig & CliPXEOptions }> {
   const pxeConfig = extractRelevantOptions<PXEServiceConfig & CliPXEOptions>(options, allPxeConfigMappings, 'pxe');
   const nodeUrl = pxeConfig.nodeUrl;
 
@@ -42,8 +27,10 @@ export async function addPXE(
   const node = deps.node ?? createAztecNodeClient(nodeUrl!, getVersions(pxeConfig), makeTracedFetch([1, 2, 3], true));
   const pxe = await createPXEService(node, pxeConfig as PXEServiceConfig);
 
+  const wallet = new TestWallet(pxe, node);
+
   // Add PXE to services list
   services.pxe = [pxe, PXESchema];
 
-  return { pxe, config: pxeConfig };
+  return { wallet, config: pxeConfig };
 }

yarn-project/aztec/src/examples/token.ts

@@ -22,7 +22,7 @@ const TRANSFER_AMOUNT = 33n;
 async function main() {
   logger.info('Running token contract test on HTTP interface.');
 
-  const accounts = await getDeployedTestAccounts(pxe);
+  const accounts = await getDeployedTestAccounts(wallet);
   const [alice, bob] = await Promise.all(
     accounts.map(async acc => {
       const accountManager = await wallet.createSchnorrAccount(acc.secret, acc.salt);

yarn-project/aztec/src/sandbox/banana_fpc.ts

@@ -6,7 +6,7 @@ import { FPCContract } from '@aztec/noir-contracts.js/FPC';
 import { TokenContract } from '@aztec/noir-contracts.js/Token';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
 import { type ContractInstanceWithAddress, getContractInstanceFromInstantiationParams } from '@aztec/stdlib/contract';
-import type { PXE } from '@aztec/stdlib/interfaces/client';
+import type { TestWallet } from '@aztec/test-wallet';
 
 const BANANA_COIN_SALT = new Fr(0);
 const bananaCoinArgs = {
@@ -62,20 +62,20 @@ export async function setupBananaFPC(initialAccounts: InitialAccountData[], wall
   log(`FPC: ${fpc.address}`);
 }
 
-export async function getDeployedBananaCoinAddress(pxe: PXE) {
+export async function getDeployedBananaCoinAddress(wallet: TestWallet) {
   const initialAccounts = await getInitialTestAccountsData();
   const bananaCoin = await getBananaCoinAddress(initialAccounts);
-  const contracts = await pxe.getContracts();
+  const contracts = await wallet.getContracts();
   if (!contracts.find(c => c.equals(bananaCoin))) {
     throw new Error('BananaCoin not deployed.');
   }
   return bananaCoin;
 }
 
-export async function getDeployedBananaFPCAddress(pxe: PXE) {
+export async function getDeployedBananaFPCAddress(wallet: TestWallet) {
   const initialAccounts = await getInitialTestAccountsData();
   const fpc = await getBananaFPCInstance(initialAccounts);
-  const contracts = await pxe.getContracts();
+  const contracts = await wallet.getContracts();
   if (!contracts.find(c => c.equals(fpc.address))) {
     throw new Error('BananaFPC not deployed.');
   }

yarn-project/aztec/src/sandbox/sponsored_fpc.ts

@@ -1,11 +1,7 @@
-import {
-  type ContractInstanceWithAddress,
-  Fr,
-  type PXE,
-  getContractInstanceFromInstantiationParams,
-} from '@aztec/aztec.js';
+import { type ContractInstanceWithAddress, Fr, getContractInstanceFromInstantiationParams } from '@aztec/aztec.js';
 import { SPONSORED_FPC_SALT } from '@aztec/constants';
 import { SponsoredFPCContract } from '@aztec/noir-contracts.js/SponsoredFPC';
+import type { TestWallet } from '@aztec/test-wallet';
 
 async function getSponsoredFPCInstance(): Promise<ContractInstanceWithAddress> {
   return await getContractInstanceFromInstantiationParams(SponsoredFPCContract.artifact, {
@@ -17,9 +13,9 @@ export async function getSponsoredFPCAddress() {
   return (await getSponsoredFPCInstance()).address;
 }
 
-export async function getDeployedSponsoredFPCAddress(pxe: PXE) {
+export async function getDeployedSponsoredFPCAddress(wallet: TestWallet) {
   const fpc = await getSponsoredFPCAddress();
-  const contracts = await pxe.getContracts();
+  const contracts = await wallet.getContracts();
   if (!contracts.find(c => c.equals(fpc))) {
     throw new Error('SponsoredFPC not deployed.');
   }

yarn-project/aztec/src/testing/aztec_cheat_codes.ts

@@ -2,18 +2,18 @@ import { Fr } from '@aztec/foundation/fields';
 import { createLogger } from '@aztec/foundation/log';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
 import { deriveStorageSlotInMap } from '@aztec/stdlib/hash';
-import type { AztecNode, PXE } from '@aztec/stdlib/interfaces/client';
-import type { Note } from '@aztec/stdlib/note';
+import type { AztecNode } from '@aztec/stdlib/interfaces/client';
+import type { Note, NotesFilter, UniqueNote } from '@aztec/stdlib/note';
 
 /**
  * A class that provides utility functions for interacting with the aztec chain.
  */
 export class AztecCheatCodes {
   constructor(
     /**
-     * The PXE Service to use for interacting with the chain
+     * The test wallet or pxe to use for getting notes
      */
-    public pxe: PXE,
+    public testWalletOrPxe: { getNotes(filter: NotesFilter): Promise<UniqueNote[]> },
     /**
      * The Aztec Node to use for interacting with the chain
      */
@@ -71,7 +71,7 @@ export class AztecCheatCodes {
    * @returns The notes stored at the given slot
    */
   public async loadPrivate(recipient: AztecAddress, contract: AztecAddress, slot: Fr | bigint): Promise<Note[]> {
-    const extendedNotes = await this.pxe.getNotes({
+    const extendedNotes = await this.testWalletOrPxe.getNotes({
       recipient,
       contractAddress: contract,
       storageSlot: new Fr(slot),

yarn-project/aztec/src/testing/cheat_codes.ts

@@ -1,7 +1,8 @@
 import { retryUntil } from '@aztec/aztec.js';
 import { EthCheatCodes, RollupCheatCodes } from '@aztec/ethereum/test';
 import type { SequencerClient } from '@aztec/sequencer-client';
-import type { AztecNode, PXE } from '@aztec/stdlib/interfaces/client';
+import type { AztecNode } from '@aztec/stdlib/interfaces/client';
+import type { NotesFilter, UniqueNote } from '@aztec/stdlib/note';
 
 import { AztecCheatCodes } from './aztec_cheat_codes.js';
 
@@ -18,9 +19,13 @@ export class CheatCodes {
     public rollup: RollupCheatCodes,
   ) {}
 
-  static async create(rpcUrls: string[], pxe: PXE, node: AztecNode): Promise<CheatCodes> {
+  static async create(
+    rpcUrls: string[],
+    testWalletOrPxe: { getNotes(filter: NotesFilter): Promise<UniqueNote[]> },
+    node: AztecNode,
+  ): Promise<CheatCodes> {
     const ethCheatCodes = new EthCheatCodes(rpcUrls);
-    const aztecCheatCodes = new AztecCheatCodes(pxe, node);
+    const aztecCheatCodes = new AztecCheatCodes(testWalletOrPxe, node);
     const rollupCheatCodes = new RollupCheatCodes(
       ethCheatCodes,
       await node.getNodeInfo().then(n => n.l1ContractAddresses),