@aztec/end-to-end 4.0.0-nightly.20250907 → 4.0.0-nightly.20260108

package/src/e2e_p2p/p2p_network.ts

@@ -1,23 +1,24 @@
-import { getSchnorrWalletWithSecretKey } from '@aztec/accounts/schnorr';
 import type { InitialAccountData } from '@aztec/accounts/testing';
 import type { AztecNodeConfig, AztecNodeService } from '@aztec/aztec-node';
-import { type AccountWalletWithSecretKey, AztecAddress, EthAddress, Fr } from '@aztec/aztec.js';
+import { AztecAddress, EthAddress } from '@aztec/aztec.js/addresses';
+import { Fr } from '@aztec/aztec.js/fields';
+import { getL1ContractsConfigEnvVars } from '@aztec/ethereum/config';
 import {
   type EmpireSlashingProposerContract,
-  type ExtendedViemWalletClient,
   GSEContract,
-  MultiAdderArtifact,
-  type Operator,
   RollupContract,
   type TallySlashingProposerContract,
-  type ViemClient,
-  createL1TxUtilsFromViemWallet,
-  deployL1Contract,
-  getL1ContractsConfigEnvVars,
-} from '@aztec/ethereum';
+} from '@aztec/ethereum/contracts';
+import type { Operator } from '@aztec/ethereum/deploy-aztec-l1-contracts';
+import { deployL1Contract } from '@aztec/ethereum/deploy-l1-contract';
+import { MultiAdderArtifact } from '@aztec/ethereum/l1-artifacts';
+import { createL1TxUtilsFromViemWallet } from '@aztec/ethereum/l1-tx-utils';
 import { ChainMonitor } from '@aztec/ethereum/test';
+import type { ExtendedViemWalletClient, ViemClient } from '@aztec/ethereum/types';
+import { EpochNumber } from '@aztec/foundation/branded-types';
 import { SecretValue } from '@aztec/foundation/config';
 import { type Logger, createLogger } from '@aztec/foundation/log';
+import { retryUntil } from '@aztec/foundation/retry';
 import { RollupAbi, SlasherAbi, TestERC20Abi } from '@aztec/l1-artifacts';
 import { SpamContract } from '@aztec/noir-test-contracts.js/Spam';
 import type { BootstrapNode } from '@aztec/p2p/bootstrap';
@@ -26,6 +27,7 @@ import { tryStop } from '@aztec/stdlib/interfaces/server';
 import { SlashFactoryContract } from '@aztec/stdlib/l1-contracts';
 import type { PublicDataTreeLeaf } from '@aztec/stdlib/trees';
 import { ZkPassportProofParams } from '@aztec/stdlib/zkpassport';
+import type { TestWallet } from '@aztec/test-wallet/server';
 import { getGenesisValues } from '@aztec/world-state/testing';
 
 import getPort from 'get-port';
@@ -43,7 +45,7 @@ import {
   createSnapshotManager,
   deployAccounts,
 } from '../fixtures/snapshot_manager.js';
-import { getPrivateKeyFromIndex, getSponsoredFPCAddress } from '../fixtures/utils.js';
+import { type SetupOptions, getPrivateKeyFromIndex, getSponsoredFPCAddress } from '../fixtures/utils.js';
 import { getEndToEndTestTelemetryClient } from '../fixtures/with_telemetry_utils.js';
 
 // Use a fixed bootstrap node private key so that we can re-use the same snapshot and the nodes can find each other
@@ -75,23 +77,22 @@ export class P2PNetworkTest {
   public prefilledPublicData: PublicDataTreeLeaf[] = [];
 
   // The re-execution test needs a wallet and a spam contract
-  public wallet?: AccountWalletWithSecretKey;
+  public wallet?: TestWallet;
   public defaultAccountAddress?: AztecAddress;
   public spamContract?: SpamContract;
 
   public bootstrapNode?: BootstrapNode;
 
   constructor(
-    testName: string,
+    public readonly testName: string,
     public bootstrapNodeEnr: string,
     public bootNodePort: number,
     public numberOfValidators: number,
-    initialValidatorConfig: AztecNodeConfig,
+    initialValidatorConfig: SetupOptions,
     public numberOfNodes = 0,
     // If set enable metrics collection
     private metricsPort?: number,
     startProverNode?: boolean,
-    mockZkPassportVerifier?: boolean,
   ) {
     this.logger = createLogger(`e2e:e2e_p2p:${testName}`);
 
@@ -120,7 +121,6 @@ export class P2PNetworkTest {
           initialValidatorConfig.slashingRoundSizeInEpochs ?? l1ContractsConfig.slashingRoundSizeInEpochs,
         slasherFlavor: initialValidatorConfig.slasherFlavor ?? 'tally',
         aztecTargetCommitteeSize: numberOfValidators,
-        salt: 420,
         metricsPort: metricsPort,
         numberOfInitialFundedAccounts: 2,
         startProverNode,
@@ -139,7 +139,6 @@ export class P2PNetworkTest {
         aztecTargetCommitteeSize: numberOfValidators,
         initialValidators: [],
         zkPassportArgs: {
-          mockZkPassportVerifier,
           zkPassportDomain: zkPassportParams.domain,
           zkPassportScope: zkPassportParams.scope,
         },
@@ -155,16 +154,14 @@ export class P2PNetworkTest {
     metricsPort,
     initialConfig,
     startProverNode,
-    mockZkPassportVerifier,
   }: {
     testName: string;
     numberOfNodes: number;
     numberOfValidators: number;
     basePort?: number;
     metricsPort?: number;
-    initialConfig?: Partial<AztecNodeConfig>;
+    initialConfig?: SetupOptions;
     startProverNode?: boolean;
-    mockZkPassportVerifier?: boolean;
   }) {
     const port = basePort || (await getPort());
 
@@ -185,7 +182,6 @@ export class P2PNetworkTest {
       numberOfNodes,
       metricsPort,
       startProverNode,
-      mockZkPassportVerifier,
     );
   }
 
@@ -198,7 +194,7 @@ export class P2PNetworkTest {
 
   async addBootstrapNode() {
     await this.snapshotManager.snapshot('add-bootstrap-node', async ({ aztecNodeConfig }) => {
-      const telemetry = getEndToEndTestTelemetryClient(this.metricsPort);
+      const telemetry = await getEndToEndTestTelemetryClient(this.metricsPort);
       this.bootstrapNode = await createBootstrapNodeFromPrivateKey(
         BOOTSTRAP_NODE_PRIVATE_KEY,
         this.bootNodePort,
@@ -230,88 +226,85 @@ export class P2PNetworkTest {
 
   async applyBaseSnapshots() {
     await this.addBootstrapNode();
-    await this.snapshotManager.snapshot(
-      'add-validators',
-      async ({ deployL1ContractsValues, dateProvider, cheatCodes }) => {
-        const rollup = getContract({
-          address: deployL1ContractsValues.l1ContractAddresses.rollupAddress.toString(),
-          abi: RollupAbi,
-          client: deployL1ContractsValues.l1Client,
-        });
+    await this.snapshotManager.snapshot('add-validators', async ({ deployL1ContractsValues, cheatCodes }) => {
+      const rollup = getContract({
+        address: deployL1ContractsValues.l1ContractAddresses.rollupAddress.toString(),
+        abi: RollupAbi,
+        client: deployL1ContractsValues.l1Client,
+      });
 
-        this.logger.info(`Adding ${this.numberOfValidators} validators`);
+      this.logger.info(`Adding ${this.numberOfValidators} validators`);
 
-        const stakingAsset = getContract({
-          address: deployL1ContractsValues.l1ContractAddresses.stakingAssetAddress.toString(),
-          abi: TestERC20Abi,
-          client: deployL1ContractsValues.l1Client,
-        });
+      const stakingAsset = getContract({
+        address: deployL1ContractsValues.l1ContractAddresses.stakingAssetAddress.toString(),
+        abi: TestERC20Abi,
+        client: deployL1ContractsValues.l1Client,
+      });
 
-        const { address: multiAdderAddress } = await deployL1Contract(
-          deployL1ContractsValues.l1Client,
-          MultiAdderArtifact.contractAbi,
-          MultiAdderArtifact.contractBytecode,
-          [rollup.address, deployL1ContractsValues.l1Client.account.address],
-        );
+      const { address: multiAdderAddress } = await deployL1Contract(
+        deployL1ContractsValues.l1Client,
+        MultiAdderArtifact.contractAbi,
+        MultiAdderArtifact.contractBytecode,
+        [rollup.address, deployL1ContractsValues.l1Client.account.address],
+      );
 
-        const multiAdder = getContract({
-          address: multiAdderAddress.toString(),
-          abi: MultiAdderArtifact.contractAbi,
-          client: deployL1ContractsValues.l1Client,
-        });
+      const multiAdder = getContract({
+        address: multiAdderAddress.toString(),
+        abi: MultiAdderArtifact.contractAbi,
+        client: deployL1ContractsValues.l1Client,
+      });
 
-        const stakeNeeded = l1ContractsConfig.activationThreshold * BigInt(this.numberOfValidators);
-        await Promise.all(
-          [await stakingAsset.write.mint([multiAdder.address, stakeNeeded], {} as any)].map(txHash =>
-            deployL1ContractsValues.l1Client.waitForTransactionReceipt({ hash: txHash }),
-          ),
-        );
+      const stakeNeeded = (await rollup.read.getActivationThreshold()) * BigInt(this.numberOfValidators);
+      await Promise.all(
+        [await stakingAsset.write.mint([multiAdder.address, stakeNeeded], {} as any)].map(txHash =>
+          deployL1ContractsValues.l1Client.waitForTransactionReceipt({ hash: txHash }),
+        ),
+      );
 
-        const { validators } = this.getValidators();
-        this.validators = validators;
+      const { validators } = this.getValidators();
+      this.validators = validators;
 
-        const gseAddress = deployL1ContractsValues.l1ContractAddresses.gseAddress!;
-        if (!gseAddress) {
-          throw new Error('GSE contract not deployed');
-        }
+      const gseAddress = deployL1ContractsValues.l1ContractAddresses.gseAddress!;
+      if (!gseAddress) {
+        throw new Error('GSE contract not deployed');
+      }
 
-        const gseContract = new GSEContract(deployL1ContractsValues.l1Client, gseAddress.toString());
+      const gseContract = new GSEContract(deployL1ContractsValues.l1Client, gseAddress.toString());
 
-        const makeValidatorTuples = async (validator: Operator) => {
-          const registrationTuple = await gseContract.makeRegistrationTuple(validator.bn254SecretKey.getValue());
-          return {
-            attester: validator.attester.toString() as `0x${string}`,
-            withdrawer: validator.withdrawer.toString() as `0x${string}`,
-            ...registrationTuple,
-          };
+      const makeValidatorTuples = async (validator: Operator) => {
+        const registrationTuple = await gseContract.makeRegistrationTuple(validator.bn254SecretKey.getValue());
+        return {
+          attester: validator.attester.toString() as `0x${string}`,
+          withdrawer: validator.withdrawer.toString() as `0x${string}`,
+          ...registrationTuple,
         };
-        const validatorTuples = await Promise.all(validators.map(makeValidatorTuples));
+      };
+      const validatorTuples = await Promise.all(validators.map(makeValidatorTuples));
 
-        await deployL1ContractsValues.l1Client.waitForTransactionReceipt({
-          hash: await multiAdder.write.addValidators([validatorTuples]),
-        });
-
-        const timestamp = await cheatCodes.rollup.advanceToEpoch(2n, { updateDateProvider: dateProvider });
+      await deployL1ContractsValues.l1Client.waitForTransactionReceipt({
+        hash: await multiAdder.write.addValidators([validatorTuples]),
+      });
 
-        // Send and await a tx to make sure we mine a block for the warp to correctly progress.
-        await this._sendDummyTx(deployL1ContractsValues.l1Client);
+      await cheatCodes.rollup.advanceToEpoch(
+        EpochNumber.fromBigInt(
+          BigInt(await cheatCodes.rollup.getEpoch()) + (await rollup.read.getLagInEpochsForValidatorSet()) + 1n,
+        ),
+      );
 
-        // Set the system time in the node, only after we have warped the time and waited for a block
-        // Time is only set in the NEXT block
-        dateProvider.setTime(Number(timestamp) * 1000);
-      },
-    );
+      // Send and await a tx to make sure we mine a block for the warp to correctly progress.
+      await this._sendDummyTx(deployL1ContractsValues.l1Client);
+    });
   }
 
   async setupAccount() {
     await this.snapshotManager.snapshot(
       'setup-account',
-      deployAccounts(1, this.logger, false),
-      async ({ deployedAccounts }, { pxe }) => {
+      deployAccounts(1, this.logger),
+      ({ deployedAccounts }, { wallet }) => {
         this.deployedAccounts = deployedAccounts;
-        const [account] = deployedAccounts;
-        this.wallet = await getSchnorrWalletWithSecretKey(pxe, account.secret, account.signingKey, account.salt);
-        this.defaultAccountAddress = this.wallet.getAddress();
+        [{ address: this.defaultAccountAddress }] = deployedAccounts;
+        this.wallet = wallet;
+        return Promise.resolve();
       },
     );
   }
@@ -329,11 +322,12 @@ export class P2PNetworkTest {
           .deployed();
         return { contractAddress: spamContract.address };
       },
-      async ({ contractAddress }) => {
+      ({ contractAddress }) => {
         if (!this.wallet) {
           throw new Error('Call snapshot t.setupAccount before deploying account contract');
         }
-        this.spamContract = await SpamContract.at(contractAddress, this.wallet);
+        this.spamContract = SpamContract.at(contractAddress, this.wallet);
+        return Promise.resolve();
       },
     );
   }
@@ -393,6 +387,48 @@ export class P2PNetworkTest {
     this.logger.info('Nodes stopped');
   }
 
+  /**
+   * Wait for P2P mesh to be fully formed across all nodes.
+   * This ensures that all nodes are connected to each other before proceeding,
+   * preventing race conditions where validators propose blocks before the network is ready.
+   *
+   * @param nodes - Array of nodes to check for P2P connectivity
+   * @param expectedNodeCount - Expected number of nodes in the network (defaults to nodes.length)
+   * @param timeoutSeconds - Maximum time to wait for connections (default: 30 seconds)
+   * @param checkIntervalSeconds - How often to check connectivity (default: 0.1 seconds)
+   */
+  async waitForP2PMeshConnectivity(
+    nodes: AztecNodeService[],
+    expectedNodeCount?: number,
+    timeoutSeconds = 30,
+    checkIntervalSeconds = 0.1,
+  ) {
+    const nodeCount = expectedNodeCount ?? nodes.length;
+    const minPeerCount = nodeCount - 1;
+
+    this.logger.warn(
+      `Waiting for all ${nodeCount} nodes to connect to P2P mesh (at least ${minPeerCount} peers each)...`,
+    );
+
+    await Promise.all(
+      nodes.map(async (node, index) => {
+        const p2p = node.getP2P();
+        await retryUntil(
+          async () => {
+            const peers = await p2p.getPeers();
+            // Each node should be connected to at least N-1 other nodes
+            return peers.length >= minPeerCount ? true : undefined;
+          },
+          `Node ${index} to connect to at least ${minPeerCount} peers`,
+          timeoutSeconds,
+          checkIntervalSeconds,
+        );
+      }),
+    );
+
+    this.logger.warn('All nodes connected to P2P mesh');
+  }
+
   async teardown() {
     await this.monitor.stop();
     await tryStop(this.bootstrapNode, this.logger);
@@ -406,7 +442,7 @@ export class P2PNetworkTest {
     slashFactory: SlashFactoryContract;
   }> {
     if (!this.ctx.deployL1ContractsValues) {
-      throw new Error('DeployL1ContractsValues not set');
+      throw new Error('DeployAztecL1ContractsValues not set');
     }
 
     const rollup = new RollupContract(
@@ -415,7 +451,7 @@ export class P2PNetworkTest {
     );
 
     const slasherContract = getContract({
-      address: getAddress(await rollup.getSlasher()),
+      address: getAddress((await rollup.getSlasherAddress()).toString()),
       abi: SlasherAbi,
       client: this.ctx.deployL1ContractsValues.l1Client,
     });

package/src/e2e_p2p/shared.ts

@@ -1,28 +1,29 @@
-import { getSchnorrAccount } from '@aztec/accounts/schnorr';
 import type { InitialAccountData } from '@aztec/accounts/testing';
 import type { AztecNodeService } from '@aztec/aztec-node';
-import {
-  AztecAddress,
-  Fr,
-  type Logger,
-  ProvenTx,
-  type SentTx,
-  TxStatus,
-  getContractInstanceFromInstantiationParams,
-  retryUntil,
-} from '@aztec/aztec.js';
+import { AztecAddress } from '@aztec/aztec.js/addresses';
+import { type SentTx, getContractInstanceFromInstantiationParams } from '@aztec/aztec.js/contracts';
+import { Fr } from '@aztec/aztec.js/fields';
+import type { Logger } from '@aztec/aztec.js/log';
+import { Tx, TxStatus } from '@aztec/aztec.js/tx';
 import type { RollupCheatCodes } from '@aztec/aztec/testing';
-import type { EmpireSlashingProposerContract, RollupContract, TallySlashingProposerContract } from '@aztec/ethereum';
+import type {
+  EmpireSlashingProposerContract,
+  RollupContract,
+  TallySlashingProposerContract,
+} from '@aztec/ethereum/contracts';
+import { EpochNumber } from '@aztec/foundation/branded-types';
 import { timesAsync, unique } from '@aztec/foundation/collection';
-import type { TestDateProvider } from '@aztec/foundation/timer';
+import { EthAddress } from '@aztec/foundation/eth-address';
+import { retryUntil } from '@aztec/foundation/retry';
+import { pluralize } from '@aztec/foundation/string';
 import type { SpamContract } from '@aztec/noir-test-contracts.js/Spam';
 import { TestContract, TestContractArtifact } from '@aztec/noir-test-contracts.js/Test';
-import { PXEService, createPXEService, getPXEServiceConfig as getRpcConfig } from '@aztec/pxe/server';
+import { getPXEConfig, getPXEConfig as getRpcConfig } from '@aztec/pxe/server';
 import { getRoundForOffense } from '@aztec/slasher';
 import type { AztecNodeAdmin } from '@aztec/stdlib/interfaces/client';
 import type { SlashFactoryContract } from '@aztec/stdlib/l1-contracts';
+import { TestWallet, proveInteraction } from '@aztec/test-wallet/server';
 
-import type { NodeContext } from '../fixtures/setup_p2p_test.js';
 import { submitTxsTo } from '../shared/submit-transactions.js';
 
 // submits a set of transactions to the provided Private eXecution Environment (PXE)
@@ -55,56 +56,46 @@ export const submitComplexTxsTo = async (
   return txs;
 };
 
-// creates an instance of the PXE and submit a given number of transactions to it.
-export const createPXEServiceAndSubmitTransactions = async (
+// creates a wallet and submit a given number of transactions through it.
+export const submitTransactions = async (
   logger: Logger,
   node: AztecNodeService,
   numTxs: number,
   fundedAccount: InitialAccountData,
-): Promise<NodeContext> => {
+): Promise<SentTx[]> => {
   const rpcConfig = getRpcConfig();
   rpcConfig.proverEnabled = false;
-  const pxeService = await createPXEService(node, rpcConfig, { useLogSuffix: true });
-
-  const account = await getSchnorrAccount(
-    pxeService,
-    fundedAccount.secret,
-    fundedAccount.signingKey,
-    fundedAccount.salt,
-  );
-  await account.register();
-  const wallet = await account.getWallet();
-
-  const txs = await submitTxsTo(pxeService, numTxs, wallet, logger);
-  return { txs, pxeService, node };
+  const wallet = await TestWallet.create(node, { ...getPXEConfig(), proverEnabled: false }, { useLogSuffix: true });
+  const fundedAccountManager = await wallet.createSchnorrAccount(fundedAccount.secret, fundedAccount.salt);
+  return submitTxsTo(wallet, fundedAccountManager.address, numTxs, logger);
 };
 
-export async function createPXEServiceAndPrepareTransactions(
+export async function prepareTransactions(
   logger: Logger,
   node: AztecNodeService,
   numTxs: number,
   fundedAccount: InitialAccountData,
-): Promise<{ pxeService: PXEService; txs: ProvenTx[]; node: AztecNodeService }> {
+): Promise<Tx[]> {
   const rpcConfig = getRpcConfig();
   rpcConfig.proverEnabled = false;
-  const pxe = await createPXEService(node, rpcConfig, { useLogSuffix: true });
 
-  const account = await getSchnorrAccount(pxe, fundedAccount.secret, fundedAccount.signingKey, fundedAccount.salt);
-  await account.register();
-  const wallet = await account.getWallet();
+  const wallet = await TestWallet.create(node, { ...getPXEConfig(), proverEnabled: false }, { useLogSuffix: true });
+  const fundedAccountManager = await wallet.createSchnorrAccount(fundedAccount.secret, fundedAccount.salt);
 
-  const testContractInstance = await getContractInstanceFromInstantiationParams(TestContractArtifact, {});
-  await wallet.registerContract({ instance: testContractInstance, artifact: TestContractArtifact });
-  const contract = await TestContract.at(testContractInstance.address, wallet);
+  const testContractInstance = await getContractInstanceFromInstantiationParams(TestContractArtifact, {
+    salt: Fr.random(),
+  });
+  await wallet.registerContract(testContractInstance, TestContractArtifact);
+  const contract = TestContract.at(testContractInstance.address, wallet);
 
-  const txs = await timesAsync(numTxs, async () => {
-    const tx = await contract.methods.emit_nullifier(Fr.random()).prove({ from: account.getAddress() });
+  return timesAsync(numTxs, async () => {
+    const tx = await proveInteraction(wallet, contract.methods.emit_nullifier(Fr.random()), {
+      from: fundedAccountManager.address,
+    });
     const txHash = tx.getTxHash();
     logger.info(`Tx prepared with hash ${txHash}`);
     return tx;
   });
-
-  return { txs, pxeService: pxe, node };
 }
 
 export function awaitProposalExecution(
@@ -147,7 +138,7 @@ export async function awaitCommitteeExists({
   logger: Logger;
 }): Promise<readonly `0x${string}`[]> {
   logger.info(`Waiting for committee to be set`);
-  let committee: readonly `0x${string}`[] | undefined;
+  let committee: EthAddress[] | undefined;
   await retryUntil(
     async () => {
       committee = await rollup.getCurrentEpochCommittee();
@@ -156,7 +147,7 @@ export async function awaitCommitteeExists({
     'non-empty committee',
     60,
   );
-  return committee!;
+  return committee!.map(c => c.toString() as `0x${string}`);
 }
 
 export async function awaitOffenseDetected({
@@ -164,26 +155,31 @@ export async function awaitOffenseDetected({
   nodeAdmin,
   slashingRoundSize,
   epochDuration,
+  waitUntilOffenseCount,
+  timeoutSeconds = 120,
 }: {
   nodeAdmin: AztecNodeAdmin;
   logger: Logger;
   slashingRoundSize: number;
   epochDuration: number;
+  waitUntilOffenseCount?: number;
+  timeoutSeconds?: number;
 }) {
-  logger.info(`Waiting for an offense to be detected`);
+  const targetOffenseCount = waitUntilOffenseCount ?? 1;
+  logger.warn(`Waiting for ${pluralize('offense', targetOffenseCount)} to be detected`);
   const offenses = await retryUntil(
     async () => {
       const offenses = await nodeAdmin.getSlashOffenses('all');
-      if (offenses.length > 0) {
+      if (offenses.length >= targetOffenseCount) {
         return offenses;
       }
     },
     'non-empty offenses',
-    60,
+    timeoutSeconds,
   );
   logger.info(
     `Hit ${offenses.length} offenses on rounds ${unique(offenses.map(o => getRoundForOffense(o, { slashingRoundSize, epochDuration })))}`,
-    offenses,
+    { offenses },
   );
   return offenses;
 }
@@ -200,8 +196,9 @@ export async function awaitCommitteeKicked({
   slashingProposer,
   slashingRoundSize,
   aztecSlotDuration,
+  aztecEpochDuration,
   logger,
-  dateProvider,
+  offenseEpoch,
 }: {
   rollup: RollupContract;
   cheatCodes: RollupCheatCodes;
@@ -210,20 +207,22 @@ export async function awaitCommitteeKicked({
   slashingProposer: EmpireSlashingProposerContract | TallySlashingProposerContract | undefined;
   slashingRoundSize: number;
   aztecSlotDuration: number;
-  dateProvider: TestDateProvider;
+  aztecEpochDuration: number;
   logger: Logger;
+  offenseEpoch: number;
 }) {
   if (!slashingProposer) {
     throw new Error('No slashing proposer configured. Cannot test slashing.');
   }
 
-  logger.info(`Advancing epochs so we start slashing`);
   await cheatCodes.debugRollup();
-  await cheatCodes.advanceToNextEpoch({ updateDateProvider: dateProvider });
-  await cheatCodes.advanceToNextEpoch({ updateDateProvider: dateProvider });
 
-  // Await for the slash payload to be created if empire (no payload is created on tally until execution time)
   if (slashingProposer.type === 'empire') {
+    // Await for the slash payload to be created if empire (no payload is created on tally until execution time)
+    const targetEpoch = EpochNumber((await cheatCodes.getEpoch()) + (await rollup.getLagInEpochsForValidatorSet()) + 1);
+    logger.info(`Advancing to epoch ${targetEpoch} so we start slashing`);
+    await cheatCodes.advanceToEpoch(targetEpoch);
+
     const slashPayloadEvents = await retryUntil(
       async () => {
         const events = await slashFactory.getSlashPayloadCreatedEvents();
@@ -238,6 +237,15 @@ export async function awaitCommitteeKicked({
     expect(unique(slashPayloadEvents[0].slashes.map(slash => slash.validator.toString()))).toHaveLength(
       committee.length,
     );
+  } else {
+    // Use the slash offset to ensure we are in the right epoch for tally
+    const slashOffsetInRounds = await slashingProposer.getSlashOffsetInRounds();
+    const slashingRoundSizeInEpochs = slashingRoundSize / aztecEpochDuration;
+    const slashingOffsetInEpochs = Number(slashOffsetInRounds) * slashingRoundSizeInEpochs;
+    const firstEpochInOffenseRound = offenseEpoch - (offenseEpoch % slashingRoundSizeInEpochs);
+    const targetEpoch = firstEpochInOffenseRound + slashingOffsetInEpochs;
+    logger.info(`Advancing to epoch ${targetEpoch} so we start slashing`);
+    await cheatCodes.advanceToEpoch(EpochNumber(targetEpoch), { offset: -aztecSlotDuration / 2 });
   }
 
   const attestersPre = await rollup.getAttesters();
@@ -248,7 +256,7 @@ export async function awaitCommitteeKicked({
     expect(attesterInfo.status).toEqual(1); // Validating
   }
 
-  const timeout = slashingRoundSize * 2 * aztecSlotDuration;
+  const timeout = slashingRoundSize * 2 * aztecSlotDuration + 30;
   logger.info(`Waiting for slash to be executed (timeout ${timeout}s)`);
   await awaitProposalExecution(slashingProposer, timeout, logger);
 
@@ -265,10 +273,11 @@ export async function awaitCommitteeKicked({
     expect(attesterInfo.status).toEqual(2); // Living
   }
 
-  logger.info(`Advancing two epochs to check current committee`);
+  logger.info(`Advancing to check current committee`);
   await cheatCodes.debugRollup();
-  await cheatCodes.advanceToNextEpoch({ updateDateProvider: dateProvider });
-  await cheatCodes.advanceToNextEpoch({ updateDateProvider: dateProvider });
+  await cheatCodes.advanceToEpoch(
+    EpochNumber((await cheatCodes.getEpoch()) + (await rollup.getLagInEpochsForValidatorSet()) + 1),
+  );
   await cheatCodes.debugRollup();
 
   const committeeNextEpoch = await rollup.getCurrentEpochCommittee();