chore: avm-proving and avm-integration tests do not require simulator to export function with jest mocks (#10228)

Subclass ContractDataSource instead of mocking. Use real WorldStateDB.
Helper function exported by simulator no longer requires
jest-mock-extended.

Author: dbanks12

noir-projects/noir-contracts/contracts/avm_test_contract/src/main.nr

@@ -619,7 +619,7 @@ contract AvmTest {
         dep::aztec::oracle::debug_log::debug_log("pedersen_hash_with_index");
         let _ = pedersen_hash_with_index(args_field);
         dep::aztec::oracle::debug_log::debug_log("test_get_contract_instance");
-        test_get_contract_instance(context.this_address());
+        test_get_contract_instance(AztecAddress::from_field(args_field[0]));
         dep::aztec::oracle::debug_log::debug_log("get_address");
         let _ = get_address();
         dep::aztec::oracle::debug_log::debug_log("get_sender");

yarn-project/bb-prover/src/avm_proving.test.ts

@@ -10,19 +10,13 @@ import path from 'path';
 import { type BBSuccess, BB_RESULT, generateAvmProof, verifyAvmProof } from './bb/execute.js';
 import { extractAvmVkData } from './verification_key/verification_key_data.js';
 
-const TIMEOUT = 180_000;
-
 describe('AVM WitGen, proof generation and verification', () => {
-  it(
-    'Should prove and verify bulk_testing',
-    async () => {
-      await proveAndVerifyAvmTestContract(
-        'bulk_testing',
-        [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10].map(x => new Fr(x)),
-      );
-    },
-    TIMEOUT,
-  );
+  it('Should prove and verify bulk_testing', async () => {
+    await proveAndVerifyAvmTestContract(
+      'bulk_testing',
+      [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10].map(x => new Fr(x)),
+    );
+  }, 180_000);
 });
 
 async function proveAndVerifyAvmTestContract(functionName: string, calldata: Fr[] = []) {

yarn-project/ivc-integration/src/avm_integration.test.ts

@@ -12,7 +12,6 @@ import { BufferReader } from '@aztec/foundation/serialize';
 import { type FixedLengthArray } from '@aztec/noir-protocol-circuits-types/types';
 import { simulateAvmTestContractGenerateCircuitInputs } from '@aztec/simulator/public/fixtures';
 
-import { jest } from '@jest/globals';
 import fs from 'fs/promises';
 import { tmpdir } from 'node:os';
 import os from 'os';
@@ -23,9 +22,6 @@ import { MockPublicBaseCircuit, witnessGenMockPublicBaseCircuit } from './index.
 
 // Auto-generated types from noir are not in camel case.
 /* eslint-disable camelcase */
-
-jest.setTimeout(240_000);
-
 const logger = createDebugLogger('aztec:avm-integration');
 
 describe('AVM Integration', () => {
@@ -120,7 +116,7 @@ describe('AVM Integration', () => {
     );
 
     expect(verifyResult.status).toBe(BB_RESULT.SUCCESS);
-  });
+  }, 240_000);
 });
 
 async function proveAvmTestContract(functionName: string, calldata: Fr[] = []): Promise<BBSuccess> {

yarn-project/simulator/src/public/fixtures/index.ts

@@ -2,7 +2,10 @@ import { PublicExecutionRequest, Tx } from '@aztec/circuit-types';
 import {
   type AvmCircuitInputs,
   CallContext,
+  type ContractClassPublic,
+  type ContractInstanceWithAddress,
   DEFAULT_GAS_LIMIT,
+  type FunctionSelector,
   Gas,
   GasFees,
   GasSettings,
@@ -20,19 +23,16 @@ import {
   computePublicBytecodeCommitment,
 } from '@aztec/circuits.js';
 import { makeContractClassPublic, makeContractInstanceFromClassId } from '@aztec/circuits.js/testing';
+import { type ContractArtifact } from '@aztec/foundation/abi';
 import { AztecAddress } from '@aztec/foundation/aztec-address';
 import { Fr, Point } from '@aztec/foundation/fields';
 import { openTmpStore } from '@aztec/kv-store/utils';
-import { PublicTxSimulator, type WorldStateDB } from '@aztec/simulator';
+import { PublicTxSimulator, WorldStateDB } from '@aztec/simulator';
 import { NoopTelemetryClient } from '@aztec/telemetry-client/noop';
 import { MerkleTrees } from '@aztec/world-state';
 
-import { mock } from 'jest-mock-extended';
-
 import { getAvmTestContractBytecode, getAvmTestContractFunctionSelector } from '../../avm/fixtures/index.js';
 
-const TIMESTAMP = new Fr(99833);
-
 /**
  * If assertionErrString is set, we expect a (non exceptional halting) revert due to a failing assertion and
  * we check that the revert reason error contains this string. However, the circuit must correctly prove the
@@ -49,46 +49,14 @@ export async function simulateAvmTestContractGenerateCircuitInputs(
 
   const globalVariables = GlobalVariables.empty();
   globalVariables.gasFees = GasFees.empty();
-  globalVariables.timestamp = TIMESTAMP;
+  globalVariables.timestamp = new Fr(99833);
 
-  const worldStateDB = mock<WorldStateDB>();
   const telemetry = new NoopTelemetryClient();
   const merkleTrees = await (await MerkleTrees.new(openTmpStore(), telemetry)).fork();
-  worldStateDB.getMerkleInterface.mockReturnValue(merkleTrees);
-
-  // Top level contract call
-  const bytecode = getAvmTestContractBytecode('public_dispatch');
-  const dispatchSelector = getAvmTestContractFunctionSelector('public_dispatch');
-  const publicFn: PublicFunction = { bytecode, selector: dispatchSelector };
-  const contractClass = makeContractClassPublic(0, publicFn);
-  const contractInstance = makeContractInstanceFromClassId(contractClass.id);
-
-  // The values here should match those in `avm_simulator.test.ts`
-  const instanceGet = new SerializableContractInstance({
-    version: 1,
-    salt: new Fr(0x123),
-    deployer: new AztecAddress(new Fr(0x456)),
-    contractClassId: new Fr(0x789),
-    initializationHash: new Fr(0x101112),
-    publicKeys: new PublicKeys(
-      new Point(new Fr(0x131415), new Fr(0x161718), false),
-      new Point(new Fr(0x192021), new Fr(0x222324), false),
-      new Point(new Fr(0x252627), new Fr(0x282930), false),
-      new Point(new Fr(0x313233), new Fr(0x343536), false),
-    ),
-  }).withAddress(contractInstance.address);
-  worldStateDB.getContractInstance
-    .mockResolvedValueOnce(contractInstance)
-    .mockResolvedValueOnce(instanceGet) // test gets deployer
-    .mockResolvedValueOnce(instanceGet) // test gets class id
-    .mockResolvedValueOnce(instanceGet) // test gets init hash
-    .mockResolvedValue(contractInstance);
-  worldStateDB.getContractClass.mockResolvedValue(contractClass);
-  worldStateDB.getBytecode.mockResolvedValue(bytecode);
-  worldStateDB.getBytecodeCommitment.mockResolvedValue(computePublicBytecodeCommitment(bytecode));
-
-  const storageValue = new Fr(5);
-  worldStateDB.storageRead.mockResolvedValue(Promise.resolve(storageValue));
+  const contractDataSource = new MockedAvmTestContractDataSource();
+  const worldStateDB = new WorldStateDB(merkleTrees, contractDataSource);
+
+  const contractInstance = contractDataSource.contractInstance;
 
   const simulator = new PublicTxSimulator(
     merkleTrees,
@@ -99,7 +67,12 @@ export async function simulateAvmTestContractGenerateCircuitInputs(
     /*doMerkleOperations=*/ true,
   );
 
-  const callContext = new CallContext(sender, contractInstance.address, dispatchSelector, /*isStaticCall=*/ false);
+  const callContext = new CallContext(
+    sender,
+    contractInstance.address,
+    contractDataSource.fnSelector,
+    /*isStaticCall=*/ false,
+  );
   const executionRequest = new PublicExecutionRequest(callContext, calldata);
 
   const tx: Tx = createTxForPublicCall(executionRequest);
@@ -159,3 +132,81 @@ export function createTxForPublicCall(
 
   return tx;
 }
+
+class MockedAvmTestContractDataSource {
+  private fnName = 'public_dispatch';
+  private bytecode: Buffer;
+  public fnSelector: FunctionSelector;
+  private publicFn: PublicFunction;
+  private contractClass: ContractClassPublic;
+  public contractInstance: ContractInstanceWithAddress;
+  private bytecodeCommitment: Fr;
+  private otherContractInstance: ContractInstanceWithAddress;
+
+  constructor() {
+    this.bytecode = getAvmTestContractBytecode(this.fnName);
+    this.fnSelector = getAvmTestContractFunctionSelector(this.fnName);
+    this.publicFn = { bytecode: this.bytecode, selector: this.fnSelector };
+    this.contractClass = makeContractClassPublic(0, this.publicFn);
+    this.contractInstance = makeContractInstanceFromClassId(this.contractClass.id);
+    this.bytecodeCommitment = computePublicBytecodeCommitment(this.bytecode);
+    // The values here should match those in `avm_simulator.test.ts`
+    this.otherContractInstance = new SerializableContractInstance({
+      version: 1,
+      salt: new Fr(0x123),
+      deployer: new AztecAddress(new Fr(0x456)),
+      contractClassId: new Fr(0x789),
+      initializationHash: new Fr(0x101112),
+      publicKeys: new PublicKeys(
+        new Point(new Fr(0x131415), new Fr(0x161718), false),
+        new Point(new Fr(0x192021), new Fr(0x222324), false),
+        new Point(new Fr(0x252627), new Fr(0x282930), false),
+        new Point(new Fr(0x313233), new Fr(0x343536), false),
+      ),
+    }).withAddress(this.contractInstance.address);
+  }
+
+  getPublicFunction(_address: AztecAddress, _selector: FunctionSelector): Promise<PublicFunction> {
+    return Promise.resolve(this.publicFn);
+  }
+
+  getBlockNumber(): Promise<number> {
+    throw new Error('Method not implemented.');
+  }
+
+  getContractClass(_id: Fr): Promise<ContractClassPublic> {
+    return Promise.resolve(this.contractClass);
+  }
+
+  getBytecodeCommitment(_id: Fr): Promise<Fr> {
+    return Promise.resolve(this.bytecodeCommitment);
+  }
+
+  addContractClass(_contractClass: ContractClassPublic): Promise<void> {
+    return Promise.resolve();
+  }
+
+  getContract(address: AztecAddress): Promise<ContractInstanceWithAddress> {
+    if (address.equals(this.contractInstance.address)) {
+      return Promise.resolve(this.contractInstance);
+    } else {
+      return Promise.resolve(this.otherContractInstance);
+    }
+  }
+
+  getContractClassIds(): Promise<Fr[]> {
+    throw new Error('Method not implemented.');
+  }
+
+  getContractArtifact(_address: AztecAddress): Promise<ContractArtifact | undefined> {
+    throw new Error('Method not implemented.');
+  }
+
+  getContractFunctionName(_address: AztecAddress, _selector: FunctionSelector): Promise<string> {
+    return Promise.resolve(this.fnName);
+  }
+
+  addContractArtifact(_address: AztecAddress, _contract: ContractArtifact): Promise<void> {
+    return Promise.resolve();
+  }
+}

yarn-project/simulator/src/public/public_tx_simulator.test.ts

@@ -204,7 +204,7 @@ describe('public_tx_simulator', () => {
         );
       },
     );
-  });
+  }, 30_000);
 
   afterEach(async () => {
     await treeStore.delete();