Skip to content

Commit

Permalink
refactor: speed up transfer test
Browse files Browse the repository at this point in the history
  • Loading branch information
@LHerskind
LHerskind committed Sep 4, 2024
1 parent ea8e0ee commit 5f8075f
Show file tree
Hide file tree
Showing 3 changed files with 121 additions and 57 deletions.
Empty file modified yarn-project/end-to-end/scripts/create_k8s_dashboard.sh
100644 → 100755
View file
Empty file.
12 changes: 8 additions & 4 deletions yarn-project/end-to-end/src/fixtures/snapshot_manager.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -467,7 +467,7 @@ async function setupFromState(statePath: string, logger: Logger): Promise<Subsys
* The 'restore' function is not provided, as it must be a closure within the test context to capture the results.
*/
export const addAccounts =
(numberOfAccounts: number, logger: DebugLogger) =>
(numberOfAccounts: number, logger: DebugLogger, waitUntilProven = true) =>
async ({ pxe }: { pxe: PXE }) => {
// Generate account keys.
const accountKeys: [Fr, GrumpkinScalar][] = Array.from({ length: numberOfAccounts }).map(_ => [
Expand All @@ -494,7 +494,7 @@ export const addAccounts =

logger.verbose('Deploying accounts...');
const txs = await Promise.all(accountManagers.map(account => account.deploy()));
await Promise.all(txs.map(tx => tx.wait({ interval: 0.1, proven: true })));
await Promise.all(txs.map(tx => tx.wait({ interval: 0.1, proven: waitUntilProven })));

return { accountKeys };
};
Expand All @@ -505,12 +505,16 @@ export const addAccounts =
* @param sender - Wallet to send the deployment tx.
* @param accountsToDeploy - Which accounts to publicly deploy.
*/
export async function publicDeployAccounts(sender: Wallet, accountsToDeploy: (CompleteAddress | AztecAddress)[]) {
export async function publicDeployAccounts(
sender: Wallet,
accountsToDeploy: (CompleteAddress | AztecAddress)[],
waitUntilProven = true,
) {
const accountAddressesToDeploy = accountsToDeploy.map(a => ('address' in a ? a.address : a));
const instances = await Promise.all(accountAddressesToDeploy.map(account => sender.getContractInstance(account)));
const batch = new BatchCall(sender, [
(await registerContractClass(sender, SchnorrAccountContractArtifact)).request(),
...instances.map(instance => deployInstance(sender, instance!).request()),
]);
await batch.send().wait({ proven: true });
await batch.send().wait({ proven: waitUntilProven });
}
166 changes: 113 additions & 53 deletions yarn-project/end-to-end/src/spartan/transfer.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@ import {
type AccountWallet,
type AccountWalletWithSecretKey,
type AztecAddress,
type CompleteAddress,
ExtendedNote,
Fr,
Note,
Expand All @@ -17,7 +16,7 @@ import { TokenContract } from '@aztec/noir-contracts.js';

import { jest } from '@jest/globals';

import { addAccounts, publicDeployAccounts } from '../fixtures/snapshot_manager.js';
import { addAccounts } from '../fixtures/snapshot_manager.js';

const { PXE_URL } = process.env;
if (!PXE_URL) {
Expand Down Expand Up @@ -65,52 +64,93 @@ describe('token transfer test', () => {
const TOKEN_SYMBOL = 'USD';
const TOKEN_DECIMALS = 18n;
const MINT_AMOUNT = 20n;

const WALLET_COUNT = 1;
const ROUNDS = 5n;

let pxe: PXE;
let wallets: AccountWalletWithSecretKey[];
let completeAddresses: CompleteAddress[];
let recipientWallet: AccountWalletWithSecretKey;
let tokenAddress: AztecAddress;
let tokenAtWallet0: TokenContract;

beforeAll(async () => {
expect(ROUNDS).toBeLessThanOrEqual(MINT_AMOUNT);

// My guess is that is never proven so if we are waiting for it to prove we wait forever?
pxe = await createCompatibleClient(PXE_URL, logger);
const { accountKeys } = await addAccounts(3, logger)({ pxe });
{
const { accountKeys } = await addAccounts(1, logger, false)({ pxe });
const accountManagers = accountKeys.map(ak => getSchnorrAccount(pxe, ak[0], ak[1], 1));

const partialAddress = accountManagers[0].getCompleteAddress().partialAddress;
await pxe.registerAccount(accountKeys[0][0], partialAddress);
recipientWallet = await accountManagers[0].getWallet();
logger.verbose(`Recipient Wallet address: ${recipientWallet.getAddress()} registered`);
}

const { accountKeys } = await addAccounts(Number(WALLET_COUNT), logger, false)({ pxe });
const accountManagers = accountKeys.map(ak => getSchnorrAccount(pxe, ak[0], ak[1], 1));
wallets = await Promise.all(accountManagers.map(a => a.getWallet()));
completeAddresses = await pxe.getRegisteredAccounts();
wallets.forEach((w, i) => logger.verbose(`Wallet ${i} address: ${w.getAddress()}`));
await publicDeployAccounts(wallets[0], completeAddresses.slice(0, 2));

wallets = await Promise.all(
accountManagers.map(async (a, i) => {
const partialAddress = a.getCompleteAddress().partialAddress;
await pxe.registerAccount(accountKeys[i][0], partialAddress);
const wallet = await a.getWallet();
logger.verbose(`Wallet ${i} address: ${wallet.getAddress()} registered`);
return wallet;
}),
);

logger.verbose(`Deploying TokenContract...`);
const tokenContract = await TokenContract.deploy(
wallets[0],
completeAddresses[0],
wallets[0].getAddress(),
TOKEN_NAME,
TOKEN_SYMBOL,
TOKEN_DECIMALS,
)
.send()
.deployed();
.deployed({ timeout: 600 });

tokenAddress = tokenContract.address;
tokenAtWallet0 = await TokenContract.at(tokenAddress, wallets[0]);

logger.verbose(`Minting ${MINT_AMOUNT} publicly...`);
await tokenAtWallet0.methods.mint_public(completeAddresses[0].address, MINT_AMOUNT).send().wait();

logger.verbose(`Minting ${MINT_AMOUNT} privately...`);
const secret = Fr.random();
const secretHash = computeSecretHash(secret);
const receipt = await tokenAtWallet0.methods.mint_private(MINT_AMOUNT, secretHash).send().wait();

await addPendingShieldNoteToPXE({
amount: MINT_AMOUNT,
secretHash,
txHash: receipt.txHash,
accountAddress: completeAddresses[0].address,
assetAddress: tokenAddress,
wallet: wallets[0],
logger.verbose(`Minting ${MINT_AMOUNT} public assets to the ${wallets.length} wallets...`);

await Promise.all(
wallets.map(w => tokenAtWallet0.methods.mint_public(w.getAddress(), MINT_AMOUNT).send().wait({ timeout: 600 })),
);

logger.verbose(`Minting ${MINT_AMOUNT} private assets to the ${wallets.length} wallets...`);
const secrets: Fr[] = wallets.map(() => Fr.random());

const txs = await Promise.all(
wallets.map((w, i) =>
tokenAtWallet0.methods.mint_private(MINT_AMOUNT, computeSecretHash(secrets[i])).send().wait({ timeout: 600 }),
),
);

wallets.forEach(async (wallet, i) => {
await addPendingShieldNoteToPXE({
amount: MINT_AMOUNT,
secretHash: computeSecretHash(secrets[i]),
txHash: txs[i].txHash,
accountAddress: wallet.getAddress(),
assetAddress: tokenAddress,
wallet: wallet,
});
});
const txClaim = tokenAtWallet0.methods.redeem_shield(completeAddresses[0].address, MINT_AMOUNT, secret).send();
await txClaim.wait({ debug: true });

await Promise.all(
wallets.map(async (w, i) =>
(await TokenContract.at(tokenAddress, w)).methods
.redeem_shield(w.getAddress(), MINT_AMOUNT, secrets[i])
.send()
.wait({ timeout: 600 }),
),
);

logger.verbose(`Minting complete.`);
});

Expand All @@ -119,34 +159,54 @@ describe('token transfer test', () => {
expect(name).toBe(TOKEN_NAME);
});

it('can transfer 1 publicly', async () => {
it('can transfer 1 token privately and publicly', async () => {
const recipient = recipientWallet.getAddress();
const transferAmount = 1n;
const numTransfers = MINT_AMOUNT / transferAmount;
const initialBalance = await tokenAtWallet0.methods.balance_of_public(completeAddresses[0].address).simulate();
expect(initialBalance).toBeGreaterThanOrEqual(transferAmount);
for (let i = 1n; i <= numTransfers; i++) {
await tokenAtWallet0.methods
.transfer_public(completeAddresses[0].address, completeAddresses[1].address, transferAmount, 0)
.send()
.wait();
}
const finalBalance0 = await tokenAtWallet0.methods.balance_of_public(completeAddresses[0].address).simulate();
expect(finalBalance0).toBe(0n);
const finalBalance1 = await tokenAtWallet0.methods.balance_of_public(completeAddresses[1].address).simulate();
expect(finalBalance1).toBe(MINT_AMOUNT);
});

it('can transfer 1 privately', async () => {
const transferAmount = 1n;
const numTransfers = MINT_AMOUNT / transferAmount;
const initialBalance = await tokenAtWallet0.methods.balance_of_private(completeAddresses[0].address).simulate();
expect(initialBalance).toBeGreaterThanOrEqual(transferAmount);
for (let i = 1n; i <= numTransfers; i++) {
await tokenAtWallet0.methods.transfer(completeAddresses[1].address, transferAmount).send().wait();
wallets.forEach(async w => {
expect(MINT_AMOUNT).toBe(
await (await TokenContract.at(tokenAddress, w)).methods.balance_of_private(w.getAddress()).simulate(),
);
expect(MINT_AMOUNT).toBe(await tokenAtWallet0.methods.balance_of_public(w.getAddress()).simulate());
});

expect(0n).toBe(
await (await TokenContract.at(tokenAddress, recipientWallet)).methods.balance_of_private(recipient).simulate(),
);
expect(0n).toBe(await tokenAtWallet0.methods.balance_of_public(recipient).simulate());

// For each round, make both private and public transfers
for (let i = 1n; i <= ROUNDS; i++) {
const txs = await Promise.all([
...wallets.map(async w =>
(await TokenContract.at(tokenAddress, w)).methods.transfer(recipient, transferAmount),
),
...wallets.map(async w =>
(
await TokenContract.at(tokenAddress, w)
).methods.transfer_public(w.getAddress(), recipient, transferAmount, 0),
),
]);

txs.forEach(async t => await t.prove());

await Promise.all(txs.map(t => t.send().wait({ timeout: 600 })));
}
const finalBalance0 = await tokenAtWallet0.methods.balance_of_private(completeAddresses[0].address).simulate();
expect(finalBalance0).toBe(0n);
const finalBalance1 = await tokenAtWallet0.methods.balance_of_private(completeAddresses[1].address).simulate();
expect(finalBalance1).toBe(MINT_AMOUNT);

wallets.forEach(async w => {
expect(MINT_AMOUNT - ROUNDS * transferAmount).toBe(
await (await TokenContract.at(tokenAddress, w)).methods.balance_of_private(w.getAddress()).simulate(),
);
expect(MINT_AMOUNT - ROUNDS * transferAmount).toBe(
await tokenAtWallet0.methods.balance_of_public(w.getAddress()).simulate(),
);
});

expect(ROUNDS * transferAmount * BigInt(wallets.length)).toBe(
await (await TokenContract.at(tokenAddress, recipientWallet)).methods.balance_of_private(recipient).simulate(),
);
expect(ROUNDS * transferAmount * BigInt(wallets.length)).toBe(
await tokenAtWallet0.methods.balance_of_public(recipient).simulate(),
);
});
});

0 comments on commit 5f8075f

Please sign in to comment.