@aztec/ethereum 0.0.0-test.1 → 0.0.1-commit.001888fc

package/src/test/rollup_cheat_codes.ts

@@ -0,0 +1,330 @@
+import { OutboxContract, RollupContract } from '@aztec/ethereum/contracts';
+import type { L1ContractAddresses } from '@aztec/ethereum/l1-contract-addresses';
+import type { ViemPublicClient } from '@aztec/ethereum/types';
+import { CheckpointNumber, EpochNumber, SlotNumber } from '@aztec/foundation/branded-types';
+import { EthAddress } from '@aztec/foundation/eth-address';
+import { createLogger } from '@aztec/foundation/log';
+import type { DateProvider } from '@aztec/foundation/timer';
+import { RollupAbi } from '@aztec/l1-artifacts/RollupAbi';
+
+import {
+  type GetContractReturnType,
+  type Hex,
+  createPublicClient,
+  fallback,
+  getContract,
+  hexToBigInt,
+  http,
+} from 'viem';
+
+import { EthCheatCodes } from './eth_cheat_codes.js';
+
+/** Cheat codes for the L1 rollup contract. */
+export class RollupCheatCodes {
+  private client: ViemPublicClient;
+  private rollup: GetContractReturnType<typeof RollupAbi, ViemPublicClient>;
+
+  private logger = createLogger('aztecjs:cheat_codes');
+
+  constructor(
+    private ethCheatCodes: EthCheatCodes,
+    addresses: Pick<L1ContractAddresses, 'rollupAddress'>,
+  ) {
+    this.client = createPublicClient({
+      chain: ethCheatCodes.chain,
+      transport: fallback(ethCheatCodes.rpcUrls.map(url => http(url, { batch: false }))),
+    });
+    this.rollup = getContract({
+      abi: RollupAbi,
+      address: addresses.rollupAddress.toString(),
+      client: this.client,
+    });
+  }
+
+  static create(
+    rpcUrls: string[],
+    addresses: Pick<L1ContractAddresses, 'rollupAddress'>,
+    dateProvider: DateProvider,
+  ): RollupCheatCodes {
+    const ethCheatCodes = new EthCheatCodes(rpcUrls, dateProvider);
+    return new RollupCheatCodes(ethCheatCodes, addresses);
+  }
+
+  /** Returns the current slot */
+  public async getSlot(): Promise<SlotNumber> {
+    const ts = BigInt((await this.client.getBlock()).timestamp);
+    return SlotNumber.fromBigInt(await this.rollup.read.getSlotAt([ts]));
+  }
+
+  /** Returns the number of seconds until the start of the given slot based on L1 block timestamp. */
+  public async getSecondsUntilSlot(slot: SlotNumber): Promise<number> {
+    const [currentTimestamp, targetTimestamp] = await Promise.all([
+      this.client.getBlock().then(b => BigInt(b.timestamp)),
+      this.rollup.read.getTimestampForSlot([BigInt(slot)]),
+    ]);
+    return Math.max(0, Number(targetTimestamp - currentTimestamp));
+  }
+
+  /** Returns the current epoch */
+  public async getEpoch(): Promise<EpochNumber> {
+    const slotNumber = await this.getSlot();
+    return EpochNumber.fromBigInt(await this.rollup.read.getEpochAtSlot([BigInt(slotNumber)]));
+  }
+
+  /**
+   * Returns the pending and proven chain tips
+   * @returns The pending and proven chain tips
+   */
+  public async getTips(): Promise<{
+    /** The pending chain tip */ pending: CheckpointNumber;
+    /** The proven chain tip */ proven: CheckpointNumber;
+  }> {
+    const { pending, proven } = await this.rollup.read.getTips();
+    return {
+      pending: CheckpointNumber.fromBigInt(pending),
+      proven: CheckpointNumber.fromBigInt(proven),
+    };
+  }
+
+  /**
+   * Logs the current state of the rollup contract.
+   */
+  public async debugRollup() {
+    const rollup = new RollupContract(this.client, this.rollup.address);
+    const pendingNum = await rollup.getCheckpointNumber();
+    const provenNum = await rollup.getProvenCheckpointNumber();
+    const validators = await rollup.getAttesters();
+    const committee = await rollup.getCurrentEpochCommittee();
+    const archive = await rollup.archive();
+    const slot = await this.getSlot();
+    const epochNum = await rollup.getEpochNumberForSlotNumber(slot);
+
+    this.logger.info(`Pending checkpoint num: ${pendingNum}`);
+    this.logger.info(`Proven checkpoint num: ${provenNum}`);
+    this.logger.info(`Validators: ${validators.map(v => v.toString()).join(', ')}`);
+    this.logger.info(`Committee: ${committee?.map(v => v.toString()).join(', ')}`);
+    this.logger.info(`Archive: ${archive}`);
+    this.logger.info(`Epoch num: ${epochNum}`);
+    this.logger.info(`Slot: ${slot}`);
+  }
+
+  /** Fetches the epoch and slot duration config from the rollup contract */
+  public async getConfig(): Promise<{
+    /** Epoch duration */ epochDuration: bigint;
+    /** Slot duration */ slotDuration: number;
+  }> {
+    const [epochDuration, slotDuration] = await Promise.all([
+      this.rollup.read.getEpochDuration(),
+      this.rollup.read.getSlotDuration(),
+    ]);
+    return { epochDuration, slotDuration: Number(slotDuration) };
+  }
+
+  /**
+   * Advances time to the beginning of the given epoch
+   * @param epoch - The epoch to advance to
+   * @param opts - Options
+   */
+  public async advanceToEpoch(
+    epoch: EpochNumber,
+    opts: {
+      /** Offset in seconds */
+      offset?: number;
+    } = {},
+  ) {
+    const { epochDuration: slotsInEpoch } = await this.getConfig();
+    const slotNumber = SlotNumber(Number(epoch) * Number(slotsInEpoch));
+    const timestamp = (await this.rollup.read.getTimestampForSlot([BigInt(slotNumber)])) + BigInt(opts.offset ?? 0);
+    try {
+      await this.ethCheatCodes.warp(Number(timestamp), { ...opts, silent: true, resetBlockInterval: true });
+      this.logger.warn(`Warped to epoch ${epoch}`, { offset: opts.offset, timestamp });
+    } catch (err) {
+      this.logger.warn(`Warp to epoch ${epoch} failed: ${err}`);
+    }
+    return timestamp;
+  }
+
+  /** Warps time in L1 until the next epoch */
+  public async advanceToNextEpoch() {
+    const slot = await this.getSlot();
+    const { epochDuration, slotDuration } = await this.getConfig();
+    const slotsUntilNextEpoch = epochDuration - (BigInt(slot) % epochDuration) + 1n;
+    const timeToNextEpoch = slotsUntilNextEpoch * BigInt(slotDuration);
+    const l1Timestamp = BigInt((await this.client.getBlock()).timestamp);
+    await this.ethCheatCodes.warp(Number(l1Timestamp + timeToNextEpoch), {
+      silent: true,
+      resetBlockInterval: true,
+    });
+    this.logger.warn(`Advanced to next epoch`);
+  }
+
+  /** Warps time in L1 until the beginning of the next slot. */
+  public async advanceToNextSlot() {
+    const currentSlot = await this.getSlot();
+    const nextSlot = SlotNumber(currentSlot + 1);
+    const timestamp = await this.rollup.read.getTimestampForSlot([BigInt(nextSlot)]);
+    await this.ethCheatCodes.warp(Number(timestamp), { silent: true, resetBlockInterval: true });
+    this.logger.warn(`Advanced to slot ${nextSlot}`);
+    return [timestamp, nextSlot];
+  }
+
+  /**
+   * Warps time in L1 equivalent to however many slots.
+   * @param howMany - The number of slots to advance.
+   */
+  public async advanceSlots(howMany: number) {
+    const l1Timestamp = (await this.client.getBlock()).timestamp;
+    const slotDuration = Number(await this.rollup.read.getSlotDuration());
+    const timeToWarp = BigInt(howMany) * BigInt(slotDuration);
+    await this.ethCheatCodes.warp(l1Timestamp + timeToWarp, { silent: true, resetBlockInterval: true });
+    const [slot, epoch] = await Promise.all([this.getSlot(), this.getEpoch()]);
+    this.logger.warn(`Advanced ${howMany} slots up to slot ${slot} in epoch ${epoch}`);
+  }
+
+  /**
+   * Marks the specified checkpoint (or latest if none) as proven
+   * @param maybeCheckpointNumber - The checkpoint number to mark as proven (defaults to latest pending)
+   */
+  public markAsProven(maybeCheckpointNumber?: CheckpointNumber) {
+    return this.ethCheatCodes.execWithPausedAnvil(async () => {
+      const tipsBefore = await this.getTips();
+      const { pending, proven } = tipsBefore;
+
+      let checkpointNumber = maybeCheckpointNumber;
+      if (checkpointNumber === undefined || checkpointNumber > pending) {
+        checkpointNumber = pending;
+      }
+      if (checkpointNumber <= proven) {
+        this.logger.debug(`Checkpoint ${checkpointNumber} is already proven`);
+        return;
+      }
+
+      // @note @LHerskind this is heavily dependent on the storage layout and size of values
+      // The rollupStore is a struct and if the size of elements or the struct changes, this can break
+      const provenCheckpointNumberSlot = hexToBigInt(RollupContract.stfStorageSlot);
+
+      // Need to pack it as a single 32 byte word
+      const newValue = (BigInt(tipsBefore.pending) << 128n) | BigInt(checkpointNumber);
+      await this.ethCheatCodes.store(EthAddress.fromString(this.rollup.address), provenCheckpointNumberSlot, newValue);
+
+      const tipsAfter = await this.getTips();
+      if (tipsAfter.pending < tipsAfter.proven) {
+        throw new Error('Overwrote pending tip to a checkpoint in the past');
+      }
+
+      this.logger.info(
+        `Proven tip moved: ${tipsBefore.proven} -> ${tipsAfter.proven}. Pending tip: ${tipsAfter.pending}.`,
+      );
+    });
+  }
+
+  /**
+   * Overrides the inProgress field of the Inbox contract state
+   * @param howMuch - How many checkpoints to move it forward
+   */
+  public advanceInboxInProgress(howMuch: number | bigint): Promise<bigint> {
+    return this.ethCheatCodes.execWithPausedAnvil(async () => {
+      // Storage slot 2 contains the InboxState struct
+      const inboxStateSlot = 2n;
+
+      // Get inbox and its current state values
+      const inboxAddress = await this.rollup.read.getInbox();
+      const currentStateValue = await this.ethCheatCodes.load(EthAddress.fromString(inboxAddress), inboxStateSlot);
+
+      // Extract current values from the packed storage slot
+      // Storage layout: rollingHash (128 bits) | totalMessagesInserted (64 bits) | inProgress (64 bits)
+      const currentRollingHash = currentStateValue & ((1n << 128n) - 1n);
+      const currentTotalMessages = (currentStateValue >> 128n) & ((1n << 64n) - 1n);
+      const currentInProgress = currentStateValue >> 192n;
+      const newInProgress = currentInProgress + BigInt(howMuch);
+
+      // Pack new values: rollingHash (low 128 bits) | totalMessages (middle 64 bits) | inProgress (high 64 bits)
+      const newValue = (BigInt(newInProgress) << 192n) | (currentTotalMessages << 128n) | currentRollingHash;
+
+      await this.ethCheatCodes.store(EthAddress.fromString(inboxAddress), inboxStateSlot, newValue, {
+        silent: true,
+      });
+
+      this.logger.warn(`Inbox inProgress advanced from ${currentInProgress} to ${newInProgress}`, {
+        inbox: inboxAddress,
+        oldValue: '0x' + currentStateValue.toString(16),
+        newValue: '0x' + newValue.toString(16),
+        rollingHash: currentRollingHash,
+        totalMessages: currentTotalMessages,
+        oldInProgress: currentInProgress,
+        newInProgress,
+      });
+
+      return newInProgress;
+    });
+  }
+
+  public insertOutbox(epoch: EpochNumber, outHash: bigint) {
+    return this.ethCheatCodes.execWithPausedAnvil(async () => {
+      const outboxAddress = await this.rollup.read.getOutbox();
+      const epochRootSlot = OutboxContract.getEpochRootStorageSlot(epoch);
+      await this.ethCheatCodes.store(EthAddress.fromString(outboxAddress), epochRootSlot, outHash);
+      this.logger.warn(`Advanced outbox to epoch ${epoch} with out hash ${outHash}`);
+    });
+  }
+
+  /**
+   * Executes an action impersonated as the owner of the Rollup contract.
+   * @param action - The action to execute
+   */
+  public async asOwner(
+    action: (owner: Hex, rollup: GetContractReturnType<typeof RollupAbi, ViemPublicClient>) => Promise<void>,
+  ) {
+    const owner = await this.rollup.read.owner();
+    await this.ethCheatCodes.startImpersonating(owner);
+    await action(owner, this.rollup);
+    await this.ethCheatCodes.stopImpersonating(owner);
+  }
+
+  /**
+   * Sets up the epoch.
+   */
+  public async setupEpoch() {
+    // Doesn't need to be done as owner, but the functionality is here...
+    await this.asOwner(async (account, rollup) => {
+      const hash = await rollup.write.setupEpoch({ account });
+      await this.client.waitForTransactionReceipt({ hash });
+      this.logger.warn(`Setup epoch`);
+    });
+  }
+
+  /** Directly calls the L1 gas fee oracle. */
+  public async updateL1GasFeeOracle() {
+    await this.asOwner(async (account, rollup) => {
+      const hash = await rollup.write.updateL1GasFeeOracle({ account, chain: this.client.chain });
+      await this.client.waitForTransactionReceipt({ hash });
+      this.logger.warn(`Updated L1 gas fee oracle`);
+    });
+  }
+
+  /**
+   * Bumps proving cost per mana.
+   * @param bumper - Callback to calculate the new proving cost per mana based on current value.
+   */
+  public async bumpProvingCostPerMana(bumper: (before: bigint) => bigint) {
+    const currentCost = await this.rollup.read.getProvingCostPerManaInEth();
+    const newCost = bumper(currentCost);
+    await this.setProvingCostPerMana(newCost);
+  }
+
+  /**
+   * Directly updates proving cost per mana.
+   * @param ethValue - The new proving cost per mana in ETH
+   */
+  public async setProvingCostPerMana(ethValue: bigint) {
+    await this.asOwner(async (account, rollup) => {
+      const hash = await rollup.write.setProvingCostPerMana([ethValue], {
+        account,
+        chain: this.client.chain,
+        gasLimit: 1000000n,
+      });
+      await this.client.waitForTransactionReceipt({ hash });
+      this.logger.warn(`Updated proving cost per mana to ${ethValue}`);
+    });
+  }
+}

package/src/test/start_anvil.ts

@@ -1,52 +1,196 @@
+import { createLogger } from '@aztec/foundation/log';
 import { makeBackoff, retry } from '@aztec/foundation/retry';
 import { fileURLToPath } from '@aztec/foundation/url';
 
-import { type Anvil, createAnvil } from '@viem/anvil';
+import { type ChildProcess, spawn } from 'child_process';
 import { dirname, resolve } from 'path';
 
+/** Minimal interface matching the @viem/anvil Anvil shape used by callers. */
+export interface Anvil {
+  readonly port: number;
+  readonly host: string;
+  readonly status: 'listening' | 'idle';
+  stop(): Promise<void>;
+}
+
 /**
  * Ensures there's a running Anvil instance and returns the RPC URL.
  */
 export async function startAnvil(
   opts: {
+    port?: number;
     l1BlockTime?: number;
+    log?: boolean;
+    captureMethodCalls?: boolean;
+    accounts?: number;
+    chainId?: number;
+    /** The hardfork to use (e.g. 'cancun', 'latest'). */
+    hardfork?: string;
+    /**
+     * Number of slots per epoch used by anvil to compute the 'finalized' and 'safe' block tags.
+     * Anvil reports `finalized = latest - slotsInAnEpoch * 2`.
+     * Defaults to 1 so the finalized block advances immediately, making tests that check
+     * L1-finality-based logic work without needing hundreds of mined blocks.
+     */
+    slotsInAnEpoch?: number;
   } = {},
-): Promise<{ anvil: Anvil; rpcUrl: string; stop: () => Promise<void> }> {
+): Promise<{ anvil: Anvil; methodCalls?: string[]; rpcUrl: string; stop: () => Promise<void> }> {
   const anvilBinary = resolve(dirname(fileURLToPath(import.meta.url)), '../../', 'scripts/anvil_kill_wrapper.sh');
+  const logger = opts.log ? createLogger('ethereum:anvil') : undefined;
+  const methodCalls = opts.captureMethodCalls ? ([] as string[]) : undefined;
 
-  let port: number | undefined;
+  let detectedPort: number | undefined;
 
-  // Start anvil.
-  // We go via a wrapper script to ensure if the parent dies, anvil dies.
   const anvil = await retry(
     async () => {
-      const anvil = createAnvil({
-        anvilBinary,
-        port: 0,
-        blockTime: opts.l1BlockTime,
-        stopTimeout: 1000,
+      const port = opts.port ?? (process.env.ANVIL_PORT ? parseInt(process.env.ANVIL_PORT) : 8545);
+      const args: string[] = [
+        '--host',
+        '127.0.0.1',
+        '--port',
+        String(port),
+        '--accounts',
+        String(opts.accounts ?? 20),
+        '--gas-limit',
+        String(45_000_000),
+        '--chain-id',
+        String(opts.chainId ?? 31337),
+      ];
+      if (opts.l1BlockTime !== undefined) {
+        args.push('--block-time', String(opts.l1BlockTime));
+      }
+      if (opts.hardfork !== undefined) {
+        args.push('--hardfork', opts.hardfork);
+      }
+      args.push('--slots-in-an-epoch', String(opts.slotsInAnEpoch ?? 1));
+
+      const child = spawn(anvilBinary, args, {
+        stdio: ['ignore', 'pipe', 'pipe'],
+        env: { ...process.env, RAYON_NUM_THREADS: '1' },
       });
 
-      // Listen to the anvil output to get the port.
-      const removeHandler = anvil.on('message', (message: string) => {
-        if (port === undefined && message.includes('Listening on')) {
-          port = parseInt(message.match(/Listening on ([^:]+):(\d+)/)![2]);
-        }
+      // Wait for "Listening on" or an early exit.
+      await new Promise<void>((resolve, reject) => {
+        let stderr = '';
+
+        const onStdout = (data: Buffer) => {
+          const text = data.toString();
+          logger?.debug(text.trim());
+          methodCalls?.push(...(text.match(/eth_[^\s]+/g) || []));
+
+          if (detectedPort === undefined && text.includes('Listening on')) {
+            const match = text.match(/Listening on ([^:]+):(\d+)/);
+            if (match) {
+              detectedPort = parseInt(match[2]);
+            }
+          }
+          if (detectedPort !== undefined) {
+            child.stdout?.removeListener('data', onStdout);
+            child.stderr?.removeListener('data', onStderr);
+            child.removeListener('close', onClose);
+            resolve();
+          }
+        };
+
+        const onStderr = (data: Buffer) => {
+          stderr += data.toString();
+          logger?.debug(data.toString().trim());
+        };
+
+        const onClose = (code: number | null) => {
+          child.stdout?.removeListener('data', onStdout);
+          child.stderr?.removeListener('data', onStderr);
+          reject(new Error(`Anvil exited with code ${code} before listening. stderr: ${stderr}`));
+        };
+
+        child.stdout?.on('data', onStdout);
+        child.stderr?.on('data', onStderr);
+        child.once('close', onClose);
       });
-      await anvil.start();
-      removeHandler();
 
-      return anvil;
+      // Continue piping for logging / method-call capture after startup.
+      if (logger || opts.captureMethodCalls) {
+        child.stdout?.on('data', (data: Buffer) => {
+          const text = data.toString();
+          logger?.debug(text.trim());
+          methodCalls?.push(...(text.match(/eth_[^\s]+/g) || []));
+        });
+        child.stderr?.on('data', (data: Buffer) => {
+          logger?.debug(data.toString().trim());
+        });
+      } else {
+        // Consume streams so the child process doesn't block on full pipe buffers.
+        child.stdout?.resume();
+        child.stderr?.resume();
+      }
+
+      return child;
     },
     'Start anvil',
     makeBackoff([5, 5, 5]),
   );
 
-  if (!port) {
+  if (!detectedPort) {
     throw new Error('Failed to start anvil');
   }
 
-  // Monkeypatch the anvil instance to include the actually assigned port
-  Object.defineProperty(anvil, 'port', { value: port, writable: false });
-  return { anvil, stop: () => anvil.stop(), rpcUrl: `http://127.0.0.1:${port}` };
+  const port = detectedPort;
+  let status: 'listening' | 'idle' = 'listening';
+
+  anvil.once('close', () => {
+    status = 'idle';
+  });
+
+  const stop = async () => {
+    if (status === 'idle') {
+      return;
+    }
+    await killChild(anvil);
+  };
+
+  const anvilObj: Anvil = {
+    port,
+    host: '127.0.0.1',
+    get status() {
+      return status;
+    },
+    stop,
+  };
+
+  return { anvil: anvilObj, methodCalls, stop, rpcUrl: `http://127.0.0.1:${port}` };
+}
+
+/** Send SIGTERM, wait up to 5 s, then SIGKILL. All timers are always cleared. */
+function killChild(child: ChildProcess): Promise<void> {
+  return new Promise<void>(resolve => {
+    if (child.exitCode !== null || child.killed) {
+      child.stdout?.destroy();
+      child.stderr?.destroy();
+      resolve();
+      return;
+    }
+
+    let killTimer: NodeJS.Timeout | undefined;
+
+    const onClose = () => {
+      if (killTimer !== undefined) {
+        clearTimeout(killTimer);
+      }
+      // Destroy stdio streams so their PipeWrap handles don't keep the event loop alive.
+      child.stdout?.destroy();
+      child.stderr?.destroy();
+      resolve();
+    };
+
+    child.once('close', onClose);
+    child.kill('SIGTERM');
+
+    killTimer = setTimeout(() => {
+      killTimer = undefined;
+      child.kill('SIGKILL');
+    }, 5000);
+
+    // Ensure the timer does not prevent Node from exiting.
+    killTimer.unref();
+  });
 }

package/src/test/upgrade_utils.ts

@@ -1,36 +1,37 @@
 import type { Logger } from '@aztec/foundation/log';
-import { TestERC20Abi as FeeJuiceAbi } from '@aztec/l1-artifacts';
+import { DateProvider } from '@aztec/foundation/timer';
 import { GovernanceAbi } from '@aztec/l1-artifacts/GovernanceAbi';
+import { TestERC20Abi as StakingAssetAbi } from '@aztec/l1-artifacts/TestERC20Abi';
 
 import { type GetContractReturnType, type PrivateKeyAccount, getContract } from 'viem';
 
-import { EthCheatCodes } from '../eth_cheat_codes.js';
+import { extractProposalIdFromLogs } from '../contracts/governance.js';
 import type { L1ContractAddresses } from '../l1_contract_addresses.js';
-import type { L1Clients } from '../types.js';
+import { createL1TxUtils } from '../l1_tx_utils/index.js';
+import type { ExtendedViemWalletClient, ViemPublicClient } from '../types.js';
+import { EthCheatCodes } from './eth_cheat_codes.js';
 
 export async function executeGovernanceProposal(
   proposalId: bigint,
-  governance: GetContractReturnType<typeof GovernanceAbi, L1Clients['publicClient']>,
+  governance: GetContractReturnType<typeof GovernanceAbi, ViemPublicClient>,
   voteAmount: bigint,
   privateKey: PrivateKeyAccount,
-  publicClient: L1Clients['publicClient'],
-  walletClient: L1Clients['walletClient'],
+  l1Client: ExtendedViemWalletClient,
   rpcUrls: string[],
   logger: Logger,
 ) {
   const proposal = await governance.read.getProposal([proposalId]);
 
+  const l1TxUtils = createL1TxUtils(l1Client);
+
   const waitL1Block = async () => {
-    await publicClient.waitForTransactionReceipt({
-      hash: await walletClient.sendTransaction({
-        to: privateKey.address,
-        value: 1n,
-        account: privateKey,
-      }),
+    await l1TxUtils.sendAndMonitorTransaction({
+      to: l1Client.account.address,
+      value: 1n,
     });
   };
 
-  const cheatCodes = new EthCheatCodes(rpcUrls, logger);
+  const cheatCodes = new EthCheatCodes(rpcUrls, new DateProvider(), logger);
 
   const timeToActive = proposal.creation + proposal.config.votingDelay;
   logger.info(`Warping to ${timeToActive + 1n}`);
@@ -40,7 +41,7 @@ export async function executeGovernanceProposal(
 
   logger.info(`Voting`);
   const voteTx = await governance.write.vote([proposalId, voteAmount, true], { account: privateKey });
-  await publicClient.waitForTransactionReceipt({ hash: voteTx });
+  await l1Client.waitForTransactionReceipt({ hash: voteTx });
   logger.info(`Voted`);
 
   const timeToExecutable = timeToActive + proposal.config.votingDuration + proposal.config.executionDelay + 1n;
@@ -50,7 +51,7 @@ export async function executeGovernanceProposal(
   await waitL1Block();
 
   const executeTx = await governance.write.execute([proposalId], { account: privateKey });
-  await publicClient.waitForTransactionReceipt({ hash: executeTx });
+  await l1Client.waitForTransactionReceipt({ hash: executeTx });
   logger.info(`Executed proposal`);
 }
 
@@ -58,12 +59,16 @@ export async function createGovernanceProposal(
   payloadAddress: `0x${string}`,
   addresses: L1ContractAddresses,
   privateKey: PrivateKeyAccount,
-  publicClient: L1Clients['publicClient'],
+  publicClient: ViemPublicClient,
   logger: Logger,
-): Promise<{ governance: GetContractReturnType<typeof GovernanceAbi, L1Clients['publicClient']>; voteAmount: bigint }> {
+): Promise<{
+  governance: GetContractReturnType<typeof GovernanceAbi, ViemPublicClient>;
+  voteAmount: bigint;
+  proposalId: bigint;
+}> {
   const token = getContract({
-    address: addresses.feeJuiceAddress.toString(),
-    abi: FeeJuiceAbi,
+    address: addresses.stakingAssetAddress.toString(),
+    abi: StakingAssetAbi,
     client: publicClient,
   });
 
@@ -92,9 +97,13 @@ export async function createGovernanceProposal(
   await publicClient.waitForTransactionReceipt({ hash: depositTx });
   logger.info(`Deposited tokens`);
 
-  await governance.write.proposeWithLock([payloadAddress, privateKey.address], {
+  const proposeTx = await governance.write.proposeWithLock([payloadAddress, privateKey.address], {
     account: privateKey,
   });
+  const receipt = await publicClient.waitForTransactionReceipt({ hash: proposeTx });
+  logger.info(`Proposed upgrade`);
+
+  const proposalId = extractProposalIdFromLogs(receipt.logs);
 
-  return { governance, voteAmount };
+  return { governance, voteAmount, proposalId };
 }