test: dual pxe tests (#12)

To run these tests you need to run two separate of the aztec sandboxes

- First you need to run the sandbox normally
`aztec start --sandbox`

- Then you need to run a PXE instance pointing to the previously started
sandbox:
`aztec start --port 8081 --pxe
--pxe.nodeUrl=http://host.docker.internal:8080`

Currently blocked because although I've sent the tokens to the escrow
contract I'm unable to withdraw from it.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **New Features**
- Launched an additional container service with configurable settings
for improved resilience.
- **Tests**
- Expanded validation scenarios for token exchanges and escrow
operations across multiple instances.
- **Chores**
- Streamlined the automated workflow to ensure clearer separation of
environment setup and testing, with a defined job timeout for enhanced
process efficiency.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Weißer Hase <wei3erhase@defi.sucks>

.github/workflows/tests.yaml

@@ -3,13 +3,11 @@ on:
       branches:
         - main
     pull_request:
-      branches:
-        - main
-        - dev
-
+
 jobs:
     setup-and-run:
       runs-on: ubuntu-latest
+      timeout-minutes: 30
 
       steps:
         - name: Checkout repository
@@ -26,11 +24,14 @@ jobs:
         - name: Update path
           run: echo "/home/runner/.aztec/bin" >> $GITHUB_PATH
 
-        - name: Set Aztec version and start sandbox
+        - name: Set Aztec version
           run: |
             VERSION=0.72.1 aztec-up
-            aztec start --sandbox &
-  
+        
+        - name: Start sandbox
+          run: |
+            docker compose -p sandbox -f ~/.aztec/docker-compose.sandbox.yml -f docker-compose.override.yml up &
+
         - name: Install project dependencies
           run: yarn
 
@@ -39,7 +40,11 @@ jobs:
 
         - name: Codegen
           run: script -e -c "aztec codegen target --outdir src/artifacts"
+
+        - name: Run js tests
+          run: script -e -c "BASE_PXE_URL=http://localhost NODE_NO_WARNINGS=1 node --experimental-vm-modules $(yarn bin jest) --no-cache --runInBand --config jest.integration.config.json"
 
-        - name: Run tests
-          run: script -e -c "NODE_NO_WARNINGS=1 node --experimental-vm-modules $(yarn bin jest) --no-cache --runInBand --config jest.integration.config.json && aztec test"
-
+        - name: Run nr tests
+          run: |
+            script -e -c "aztec test"
+

docker-compose.override.yml

@@ -0,0 +1,18 @@
+services:
+  alternative-pxe:
+    image: "aztecprotocol/aztec:latest"
+    ports:
+      - "${PXE_PORT:-8081}:${PXE_PORT:-8081}"
+    environment:
+      LOG_LEVEL: '${LOG_LEVEL:-info; verbose: simulator:avm:debug_log}'
+      HOST_WORKDIR: "${PWD}"
+      VERSION: latest
+    volumes:
+      - ./pxe/log:/usr/src/yarn-project/aztec/log:rw
+    # TODO: Use `condition` to start only after the sandbox is healthy
+    depends_on:
+      - ethereum
+      - aztec
+    command: "start --port 8081 --pxe --pxe.nodeUrl=http://aztec:8080"
+    # TODO: If started at the same time it usually takes three attempts to connect to the sandbox's node
+    restart: unless-stopped

src/escrow_contract/src/test/escrow.test.ts

@@ -0,0 +1,229 @@
+import { TokenContractArtifact, TokenContract, Transfer } from '../../../artifacts/Token.js';
+import {
+  AccountWallet,
+  createLogger,
+  Fr,
+  PXE,
+  waitForPXE,
+  TxStatus,
+  createPXEClient,
+  getContractInstanceFromDeployParams,
+  Logger,
+  Contract,
+  AztecAddress,
+  AccountWalletWithSecretKey,
+  Wallet,
+  UniqueNote,
+} from '@aztec/aztec.js';
+import { createAccount } from '@aztec/accounts/testing';
+import { computePartialAddress, deriveKeys } from '@aztec/circuits.js';
+import { EscrowContract, EscrowContractArtifact } from '@aztec/noir-contracts.js/Escrow';
+
+const createPXE = async (id: number = 0) => {
+  // TODO: we should probably define testing fixtures for this kind of configuration
+  const { BASE_PXE_URL = `http://localhost` } = process.env;
+  const url = `${BASE_PXE_URL}:${8080 + id}`;
+  const pxe = createPXEClient(url);
+  await waitForPXE(pxe);
+  return pxe;
+};
+
+const setupSandbox = async () => {
+  return createPXE();
+};
+
+async function deployToken(deployer: AccountWallet, minter: AztecAddress) {
+  const contract = await Contract.deploy(deployer, TokenContractArtifact, [minter, 'PrivateToken', 'PT', 18])
+    .send()
+    .deployed();
+  console.log('Token contract deployed at', contract.address);
+  return contract;
+}
+
+async function deployEscrow(pxes: PXE[], wallet: Wallet, owner: AztecAddress) {
+  const escrowSecretKey = Fr.random();
+  const escrowPublicKeys = (await deriveKeys(escrowSecretKey)).publicKeys;
+  const escrowDeployment = EscrowContract.deployWithPublicKeys(escrowPublicKeys, wallet, owner);
+  const escrowInstance = await escrowDeployment.getInstance();
+
+  await Promise.all(
+    pxes.map(async (pxe) => pxe.registerAccount(escrowSecretKey, await computePartialAddress(escrowInstance))),
+  );
+
+  const escrowContract = await escrowDeployment.send().deployed();
+  console.log(`Escrow contract deployed at ${escrowContract.address}`);
+
+  return escrowContract;
+}
+
+describe('Multi PXE', () => {
+  let alicePXE: PXE;
+  let bobPXE: PXE;
+
+  let aliceWallet: AccountWalletWithSecretKey;
+  let bobWallet: AccountWalletWithSecretKey;
+
+  let alice: AccountWallet;
+  let bob: AccountWallet;
+  let carl: AccountWallet;
+
+  let token: TokenContract;
+  let escrow: EscrowContract;
+  const AMOUNT = 1000n;
+
+  let logger: Logger;
+
+  beforeAll(async () => {
+    logger = createLogger('aztec:aztec-starter');
+    logger.info('Aztec-Starter tests running.');
+
+    alicePXE = await createPXE(0);
+    bobPXE = await createPXE(1);
+
+    // TODO: assert that the used PXEs are actually separate instances?
+
+    aliceWallet = await createAccount(alicePXE);
+    bobWallet = await createAccount(bobPXE);
+
+    alice = aliceWallet;
+    bob = bobWallet;
+    console.log({
+      alice: aliceWallet.getAddress(),
+      bob: bobWallet.getAddress(),
+    });
+  });
+
+  beforeEach(async () => {
+    token = (await deployToken(alice, alice.getAddress())) as TokenContract;
+
+    await bobPXE.registerContract(token);
+
+    escrow = await deployEscrow([alicePXE, bobPXE], alice, bob.getAddress());
+    await bobPXE.registerContract({
+      instance: escrow.instance,
+      artifact: EscrowContractArtifact,
+    });
+    await alicePXE.registerContract({
+      instance: escrow.instance,
+      artifact: EscrowContractArtifact,
+    });
+
+    // alice knows bob
+    await alicePXE.registerAccount(bobWallet.getSecretKey(), bob.getCompleteAddress().partialAddress);
+    alicePXE.registerSender(bob.getAddress());
+    alice.setScopes([
+      alice.getAddress(),
+      bob.getAddress(),
+      // token.address,
+    ]);
+    // bob knows alice
+    await bobPXE.registerAccount(aliceWallet.getSecretKey(), alice.getCompleteAddress().partialAddress);
+    bobPXE.registerSender(alice.getAddress());
+
+    bob.setScopes([
+      bob.getAddress(),
+      alice.getAddress(),
+      // token.address
+      escrow.address,
+    ]);
+  });
+
+  const expectAddressNote = (note: UniqueNote, address: AztecAddress, owner: AztecAddress) => {
+    logger.info('checking address note {} {}', [address, owner]);
+    expect(note.note.items[0]).toEqual(new Fr(address.toBigInt()));
+    expect(note.note.items[1]).toEqual(new Fr(owner.toBigInt()));
+  };
+
+  const expectNote = (note: UniqueNote, amount: bigint, owner: AztecAddress) => {
+    // 4th element of items is randomness, so we slice the first 3
+    // dev: why the second element is always 0?
+    expect(note.note.items.slice(0, 3)).toStrictEqual([new Fr(amount), new Fr(0), new Fr(owner.toBigInt())]);
+  };
+
+  const expectBalances = async (address: AztecAddress, publicBalance: bigint, privateBalance: bigint) => {
+    logger.info('checking balances for', address.toString());
+    expect(await token.methods.balance_of_public(address).simulate()).toBe(publicBalance);
+    expect(await token.methods.balance_of_private(address).simulate()).toBe(privateBalance);
+  };
+
+  const wad = (n: number = 1) => AMOUNT * BigInt(n);
+
+  it('escrow', async () => {
+    let events, notes;
+
+    // this is here because the note is created in the constructor
+    await escrow.withWallet(alice).methods.sync_notes().simulate({});
+    await escrow.withWallet(bob).methods.sync_notes().simulate({});
+
+    // alice should have no notes (But it has because I gave it access to Bob's notes)
+    notes = await alice.getNotes({ contractAddress: escrow.address });
+    expect(notes.length).toBe(1);
+    expectAddressNote(notes[0], bob.getAddress(), bob.getAddress());
+
+    // bob should have a note with himself as owner, encrypted by alice
+    notes = await bob.getNotes({ contractAddress: escrow.address });
+    expect(notes.length).toBe(1);
+    expectAddressNote(notes[0], bob.getAddress(), bob.getAddress());
+
+    // mint initial amount
+    await token.withWallet(alice).methods.mint_to_public(alice.getAddress(), wad(10)).send().wait();
+
+    await token.withWallet(alice).methods.transfer_to_private(alice.getAddress(), wad(5)).send().wait();
+    await token.withWallet(alice).methods.sync_notes().simulate({});
+
+    // assert balances
+    await expectBalances(alice.getAddress(), wad(5), wad(5));
+    await expectBalances(bob.getAddress(), wad(0), wad(0));
+
+    // Transfer both in private and public
+    const fundEscrowTx = await token
+      .withWallet(alice)
+      .methods.transfer_in_private(alice.getAddress(), escrow.address, wad(5), 0)
+      .send()
+      .wait({
+        debug: true,
+      });
+
+    const fundEscrowTx2 = await token
+      .withWallet(alice)
+      .methods.transfer_in_public(alice.getAddress(), escrow.address, wad(5), 0)
+      .send()
+      .wait({
+        debug: true,
+      });
+
+    await token.withWallet(alice).methods.sync_notes().simulate({});
+
+    // assert balances, alice 0 and 0, escrow 5 and 5
+    await expectBalances(alice.getAddress(), wad(0), wad(0));
+    await expectBalances(escrow.address, wad(5), wad(5));
+
+    // alice should have a note with escrow as owner (why alice can see the escrow's note?)
+    notes = await alice.getNotes({ contractAddress: token.address });
+    expect(notes.length).toBe(1);
+    expectNote(notes[0], wad(5), escrow.address);
+
+    await escrow.withWallet(alice).methods.sync_notes().simulate({});
+    await escrow.withWallet(bob).methods.sync_notes().simulate({});
+
+    // Q: why only alice can see the escrow's notes if both have the escrow registered?
+    notes = await alice.getNotes({ owner: escrow.address });
+    expect(notes.length).toBe(1);
+    expectNote(notes[0], wad(5), escrow.address);
+
+    notes = await bob.getNotes({ owner: escrow.address });
+    expect(notes.length).toBe(0);
+
+    // withdraw 1 from the escrow
+    const withdrawTx = await escrow
+      .withWallet(bob)
+      .methods.withdraw(token.address, wad(1), bob.getAddress())
+      .send()
+      .wait({
+        debug: true,
+      });
+
+    await expectBalances(escrow.address, wad(5), wad(4));
+    await expectBalances(bob.getAddress(), wad(0), wad(1));
+  }, 300_000);
+});