chore(e2e): fix mbps proposed chain test

yarn-project/end-to-end/src/e2e_account_contracts.test.ts

@@ -11,16 +11,17 @@ 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';
 
 export class TestWalletInternals extends TestWallet {
   static override async create(node: AztecNode): Promise<TestWalletInternals> {
     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) {

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()) });

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,

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<PXEConfig>,
     options: PXECreationOptions = { loggers: {} },
   ): Promise<TestWallet> {
+    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<AccountManager> {
@@ -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<string | symbol, unknown>)[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;
+  }
+}