diff --git a/yarn-project/end-to-end/src/e2e_account_contracts.test.ts b/yarn-project/end-to-end/src/e2e_account_contracts.test.ts index 3e5c071cf6c2..e07d5251b0d5 100644 --- a/yarn-project/end-to-end/src/e2e_account_contracts.test.ts +++ b/yarn-project/end-to-end/src/e2e_account_contracts.test.ts @@ -11,7 +11,7 @@ import { randomBytes } from '@aztec/foundation/crypto/random'; import { ChildContract } from '@aztec/noir-test-contracts.js/Child'; import { createPXE, getPXEConfig } from '@aztec/pxe/server'; import { deriveSigningKey } from '@aztec/stdlib/keys'; -import { TestWallet } from '@aztec/test-wallet/server'; +import { AztecNodeProxy, TestWallet } from '@aztec/test-wallet/server'; import { setup } from './fixtures/utils.js'; @@ -19,8 +19,9 @@ export class TestWalletInternals extends TestWallet { static override async create(node: AztecNode): Promise { const pxeConfig = getPXEConfig(); pxeConfig.proverEnabled = false; - const pxe = await createPXE(node, pxeConfig); - return new TestWalletInternals(pxe, node); + const nodeRef = new AztecNodeProxy(node); + const pxe = await createPXE(nodeRef, pxeConfig); + return new TestWalletInternals(pxe, nodeRef); } replaceAccountAt(account: Account, address: AztecAddress) { diff --git a/yarn-project/end-to-end/src/e2e_epochs/epochs_mbps.parallel.test.ts b/yarn-project/end-to-end/src/e2e_epochs/epochs_mbps.parallel.test.ts index e9375b5742b5..22b3a068191b 100644 --- a/yarn-project/end-to-end/src/e2e_epochs/epochs_mbps.parallel.test.ts +++ b/yarn-project/end-to-end/src/e2e_epochs/epochs_mbps.parallel.test.ts @@ -110,10 +110,9 @@ describe('e2e_epochs/epochs_mbps', () => { ({ context, logger, rollup } = test); wallet = context.wallet; - archiver = (context.aztecNode as AztecNodeService).getBlockSource() as Archiver; from = context.accounts[0]; - // Deploy cross-chain contract if needed (before stopping the initial sequencer). + // Deploy cross-chain contract if needed (before stopping the initial node). // Unlike emit_nullifier (which has #[noinitcheck]), cross-chain methods require a deployed contract. if (deployCrossChainContract) { logger.warn(`Deploying cross-chain test contract before stopping initial sequencer`); @@ -132,6 +131,12 @@ describe('e2e_epochs/epochs_mbps', () => { ); logger.warn(`Started ${NODE_COUNT} validator nodes.`, { validators: validators.map(v => v.attester.toString()) }); + // Point the wallet at a validator node. The initial node-0 has all validator keys in its config, + // so it rejects block proposals from validators thinking they come from itself. By redirecting + // the wallet to a validator node, the PXE correctly tracks proposed blocks. + wallet.updateNode(nodes[0]); + archiver = nodes[0].getBlockSource() as Archiver; + // Register contract for sending txs. contract = await test.registerTestContract(wallet); logger.warn(`Test setup completed.`, { validators: validators.map(v => v.attester.toString()) }); diff --git a/yarn-project/test-wallet/src/server.ts b/yarn-project/test-wallet/src/server.ts index 0fb54e4ebc04..8be2174b09ad 100644 --- a/yarn-project/test-wallet/src/server.ts +++ b/yarn-project/test-wallet/src/server.ts @@ -1,4 +1,4 @@ -export { TestWallet } from './wallet/server.js'; +export { TestWallet, AztecNodeProxy } from './wallet/server.js'; export { type AccountData } from './wallet/test_wallet.js'; export { deployFundedSchnorrAccounts, diff --git a/yarn-project/test-wallet/src/wallet/server.ts b/yarn-project/test-wallet/src/wallet/server.ts index 9f3095f3a9e4..5efbf1aed28c 100644 --- a/yarn-project/test-wallet/src/wallet/server.ts +++ b/yarn-project/test-wallet/src/wallet/server.ts @@ -7,7 +7,7 @@ import { getContractInstanceFromInstantiationParams } from '@aztec/aztec.js/cont import { Fq, Fr } from '@aztec/aztec.js/fields'; import type { AztecNode } from '@aztec/aztec.js/node'; import { AccountManager } from '@aztec/aztec.js/wallet'; -import { type PXEConfig, type PXECreationOptions, createPXE, getPXEConfig } from '@aztec/pxe/server'; +import { PXE, type PXEConfig, type PXECreationOptions, createPXE, getPXEConfig } from '@aztec/pxe/server'; import { deriveSigningKey } from '@aztec/stdlib/keys'; import { BaseTestWallet } from './test_wallet.js'; @@ -18,17 +18,33 @@ import { BaseTestWallet } from './test_wallet.js'; * from the `pxe/server` package. */ export class TestWallet extends BaseTestWallet { + constructor( + pxe: PXE, + private readonly nodeRef: AztecNodeProxy, + ) { + super(pxe, nodeRef); + } + static async create( node: AztecNode, overridePXEConfig?: Partial, options: PXECreationOptions = { loggers: {} }, ): Promise { + const nodeRef = new AztecNodeProxy(node); const pxeConfig = Object.assign(getPXEConfig(), { proverEnabled: overridePXEConfig?.proverEnabled ?? false, ...overridePXEConfig, }); - const pxe = await createPXE(node, pxeConfig, options); - return new TestWallet(pxe, node); + const pxe = await createPXE(nodeRef, pxeConfig, options); + return new TestWallet(pxe, nodeRef); + } + + /** + * Updates the underlying node that this wallet and its PXE communicate with. + * @param node - The new AztecNode to forward all calls to. + */ + updateNode(node: AztecNode): void { + this.nodeRef.updateTargetNode(node); } createSchnorrAccount(secret: Fr, salt: Fr, signingKey?: Fq): Promise { @@ -89,3 +105,40 @@ export class TestWallet extends BaseTestWallet { }; } } + +/** + * Extends AztecNode via declaration merging so instances can be used wherever AztecNode is expected. + * The actual method forwarding is handled by a Proxy in the class constructor. + */ +// eslint-disable-next-line @typescript-eslint/no-unsafe-declaration-merging +export interface AztecNodeProxy extends AztecNode {} + +/** + * Mutable wrapper around an AztecNode that forwards all calls to the current target. + * Allows swapping the underlying node at runtime via updateTargetNode, which is useful + * for tests that need to redirect a wallet from one node to another without recreating it. + */ +// eslint-disable-next-line @typescript-eslint/no-unsafe-declaration-merging +export class AztecNodeProxy { + constructor(private target: AztecNode) { + return new Proxy(this, { + get: (obj, prop, receiver) => { + // Own properties and methods (updateTargetNode, target) are served directly. + if (Reflect.has(obj, prop)) { + return Reflect.get(obj, prop, receiver); + } + // Everything else is forwarded to the current target node. + const val = (obj.target as unknown as Record)[prop]; + return typeof val === 'function' ? val.bind(obj.target) : val; + }, + }); + } + + /** + * Updates the underlying node that this reference points to. + * @param node - The new node to forward calls to. + */ + updateTargetNode(node: AztecNode): void { + this.target = node; + } +}