Skip to content

Commit

Permalink
token utils
Browse files Browse the repository at this point in the history
  • Loading branch information
benesjan committed Oct 29, 2024
1 parent 6d16091 commit 618cd24
Show file tree
Hide file tree
Showing 3 changed files with 95 additions and 105 deletions.
138 changes: 33 additions & 105 deletions yarn-project/end-to-end/src/e2e_2_pxes.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,20 +4,19 @@ import {
type AztecAddress,
type AztecNode,
type DebugLogger,
ExtendedNote,
type ExtendedNote,
Fr,
Note,
type PXE,
type Wallet,
computeSecretHash,
retryUntil,
sleep,
} from '@aztec/aztec.js';
import { ChildContract, TestContract, TokenContract } from '@aztec/noir-contracts.js';

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

import { expectsNumOfNoteEncryptedLogsInTheLastBlockToBe, setup, setupPXEService } from './fixtures/utils.js';
import { deployToken, expectTokenBalance, mintTokensToPrivate } from './fixtures/token_utils.js';
import { setup, setupPXEService } from './fixtures/utils.js';

const TIMEOUT = 120_000;

Expand Down Expand Up @@ -55,78 +54,12 @@ describe('e2e_2_pxes', () => {
await teardownA();
});

const awaitUserSynchronized = async (wallet: Wallet, owner: AztecAddress) => {
const isUserSynchronized = async () => {
return await wallet.isAccountStateSynchronized(owner);
};
await retryUntil(isUserSynchronized, `synch of user ${owner.toString()}`, 10);
};

const expectTokenBalance = async (
wallet: Wallet,
tokenAddress: AztecAddress,
owner: AztecAddress,
expectedBalance: bigint,
checkIfSynchronized = true,
) => {
if (checkIfSynchronized) {
// First wait until the corresponding PXE has synchronized the account
await awaitUserSynchronized(wallet, owner);
}

// Then check the balance
const contractWithWallet = await TokenContract.at(tokenAddress, wallet);
const balance = await contractWithWallet.methods.balance_of_private(owner).simulate({ from: owner });
logger.info(`Account ${owner} balance: ${balance}`);
expect(balance).toBe(expectedBalance);
};

const deployTokenContract = async (initialAdminBalance: bigint, admin: AztecAddress, pxe: PXE) => {
logger.info(`Deploying Token contract...`);
const contract = await TokenContract.deploy(walletA, admin, 'TokenName', 'TokenSymbol', 18).send().deployed();

if (initialAdminBalance > 0n) {
// Minter is minting to herself so contract as minter is the same as contract as recipient
await mintTokens(contract, contract, admin, initialAdminBalance, pxe);
}

logger.info('L2 contract deployed');

return contract;
};

const mintTokens = async (
contractAsMinter: TokenContract,
contractAsRecipient: TokenContract,
recipient: AztecAddress,
balance: bigint,
recipientPxe: PXE,
) => {
const secret = Fr.random();
const secretHash = computeSecretHash(secret);

const receipt = await contractAsMinter.methods.mint_private(balance, secretHash).send().wait();

const note = new Note([new Fr(balance), secretHash]);
const extendedNote = new ExtendedNote(
note,
recipient,
contractAsMinter.address,
TokenContract.storage.pending_shields.slot,
TokenContract.notes.TransparentNote.id,
receipt.txHash,
);
await recipientPxe.addNote(extendedNote, recipient);

await contractAsRecipient.methods.redeem_shield(recipient, balance, secret).send().wait();
};

it('transfers funds from user A to B via PXE A followed by transfer from B to A via PXE B', async () => {
const initialBalance = 987n;
const transferAmount1 = 654n;
const transferAmount2 = 323n;

const token = await deployTokenContract(initialBalance, walletA.getAddress(), pxeA);
const token = await deployToken(walletA, initialBalance, logger);

// Add account B to wallet A
await pxeA.registerRecipient(walletB.getCompleteAddress());
Expand All @@ -136,33 +69,31 @@ describe('e2e_2_pxes', () => {
// Add token to PXE B (PXE A already has it because it was deployed through it)
await pxeB.registerContract(token);

// Check initial balances and logs are as expected
await expectTokenBalance(walletA, token.address, walletA.getAddress(), initialBalance);
await expectTokenBalance(walletB, token.address, walletB.getAddress(), 0n);
await expectsNumOfNoteEncryptedLogsInTheLastBlockToBe(aztecNode, 1);
// Check initial balances are as expected
await expectTokenBalance(walletA, token, walletA.getAddress(), initialBalance, logger);
await expectTokenBalance(walletB, token, walletB.getAddress(), 0n, logger);

// Transfer funds from A to B via PXE A
const contractWithWalletA = await TokenContract.at(token.address, walletA);
await contractWithWalletA.methods.transfer(walletB.getAddress(), transferAmount1).send().wait();

// Check balances and logs are as expected
await expectTokenBalance(walletA, token.address, walletA.getAddress(), initialBalance - transferAmount1);
await expectTokenBalance(walletB, token.address, walletB.getAddress(), transferAmount1);
await expectsNumOfNoteEncryptedLogsInTheLastBlockToBe(aztecNode, 2);
// Check balances are as expected
await expectTokenBalance(walletA, token, walletA.getAddress(), initialBalance - transferAmount1, logger);
await expectTokenBalance(walletB, token, walletB.getAddress(), transferAmount1, logger);

// Transfer funds from B to A via PXE B
const contractWithWalletB = await TokenContract.at(token.address, walletB);
await contractWithWalletB.methods.transfer(walletA.getAddress(), transferAmount2).send().wait({ interval: 0.1 });

// Check balances and logs are as expected
// Check balances are as expected
await expectTokenBalance(
walletA,
token.address,
token,
walletA.getAddress(),
initialBalance - transferAmount1 + transferAmount2,
logger,
);
await expectTokenBalance(walletB, token.address, walletB.getAddress(), transferAmount1 - transferAmount2);
await expectsNumOfNoteEncryptedLogsInTheLastBlockToBe(aztecNode, 2);
await expectTokenBalance(walletB, token, walletB.getAddress(), transferAmount1 - transferAmount2, logger);
});

const deployChildContractViaServerA = async () => {
Expand Down Expand Up @@ -212,8 +143,7 @@ describe('e2e_2_pxes', () => {
const userABalance = 100n;
const userBBalance = 150n;

const token = await deployTokenContract(userABalance, walletA.getAddress(), pxeA);
const contractWithWalletA = await TokenContract.at(token.address, walletA);
const token = await deployToken(walletA, userABalance, logger);

// Add account B to wallet A
await pxeA.registerRecipient(walletB.getCompleteAddress());
Expand All @@ -224,21 +154,20 @@ describe('e2e_2_pxes', () => {
await pxeB.registerContract(token);

// Mint tokens to user B
const contractWithWalletB = await TokenContract.at(token.address, walletB);
await mintTokens(contractWithWalletA, contractWithWalletB, walletB.getAddress(), userBBalance, pxeB);
await mintTokensToPrivate(token, walletA, walletB.getAddress(), userBBalance);

// Check that user A balance is 100 on server A
await expectTokenBalance(walletA, token.address, walletA.getAddress(), userABalance);
await expectTokenBalance(walletA, token, walletA.getAddress(), userABalance, logger);
// Check that user B balance is 150 on server B
await expectTokenBalance(walletB, token.address, walletB.getAddress(), userBBalance);
await expectTokenBalance(walletB, token, walletB.getAddress(), userBBalance, logger);

// CHECK THAT PRIVATE BALANCES ARE 0 WHEN ACCOUNT'S SECRET KEYS ARE NOT REGISTERED
// Note: Not checking if the account is synchronized because it is not registered as an account (it would throw).
const checkIfSynchronized = false;
// Check that user A balance is 0 on server B
await expectTokenBalance(walletB, token.address, walletA.getAddress(), 0n, checkIfSynchronized);
await expectTokenBalance(walletB, token, walletA.getAddress(), 0n, logger, checkIfSynchronized);
// Check that user B balance is 0 on server A
await expectTokenBalance(walletA, token.address, walletB.getAddress(), 0n, checkIfSynchronized);
await expectTokenBalance(walletA, token, walletB.getAddress(), 0n, logger, checkIfSynchronized);
});

it('permits migrating an account from one PXE to another', async () => {
Expand All @@ -263,28 +192,25 @@ describe('e2e_2_pxes', () => {
const initialBalance = 987n;
const transferAmount1 = 654n;

const token = await deployTokenContract(initialBalance, walletA.getAddress(), pxeA);
const tokenAddress = token.address;
const token = await deployToken(walletA, initialBalance, logger);

// Add account B to wallet A
await pxeA.registerRecipient(walletB.getCompleteAddress());
// Add account A to wallet B
await pxeB.registerRecipient(walletA.getCompleteAddress());

// Check initial balances and logs are as expected
await expectTokenBalance(walletA, tokenAddress, walletA.getAddress(), initialBalance);
// Check initial balances are as expected
await expectTokenBalance(walletA, token, walletA.getAddress(), initialBalance, logger);
// don't check userB yet

await expectsNumOfNoteEncryptedLogsInTheLastBlockToBe(aztecNode, 1);

// Transfer funds from A to B via PXE A
const contractWithWalletA = await TokenContract.at(tokenAddress, walletA);
const contractWithWalletA = await TokenContract.at(token.address, walletA);
await contractWithWalletA.methods.transfer(walletB.getAddress(), transferAmount1).send().wait();

// now add the contract and check balances
await pxeB.registerContract(token);
await expectTokenBalance(walletA, tokenAddress, walletA.getAddress(), initialBalance - transferAmount1);
await expectTokenBalance(walletB, tokenAddress, walletB.getAddress(), transferAmount1);
await expectTokenBalance(walletA, token, walletA.getAddress(), initialBalance - transferAmount1, logger);
await expectTokenBalance(walletB, token, walletB.getAddress(), transferAmount1, logger);
});

it('permits sending funds to a user, and spending them, before they have registered the contract', async () => {
Expand All @@ -307,7 +233,7 @@ describe('e2e_2_pxes', () => {
await pxeA.registerRecipient(walletB.getCompleteAddress());

// deploy the contract on PXE A
const token = await deployTokenContract(initialBalance, walletA.getAddress(), pxeA);
const token = await deployToken(walletA, initialBalance, logger);

// Transfer funds from A to Shared Wallet via PXE A
const contractWithWalletA = await TokenContract.at(token.address, walletA);
Expand All @@ -318,12 +244,13 @@ describe('e2e_2_pxes', () => {
await contractWithSharedWalletA.methods.transfer(walletB.getAddress(), transferAmount2).send().wait();

// check balances from PXE-A's perspective
await expectTokenBalance(walletA, token.address, walletA.getAddress(), initialBalance - transferAmount1);
await expectTokenBalance(walletA, token, walletA.getAddress(), initialBalance - transferAmount1, logger);
await expectTokenBalance(
sharedWalletOnA,
token.address,
token,
sharedAccountAddress.address,
transferAmount1 - transferAmount2,
logger,
);

// now add the contract and check balances from PXE-B's perspective.
Expand All @@ -332,13 +259,14 @@ describe('e2e_2_pxes', () => {
// PXE-B adds the contract
// PXE-B reprocesses the deferred notes, and sees the nullifier for A -> Shared
await pxeB.registerContract(token);
await expectTokenBalance(walletB, token.address, walletB.getAddress(), transferAmount2);
await expectTokenBalance(walletB, token, walletB.getAddress(), transferAmount2, logger);
await expect(sharedWalletOnB.isAccountStateSynchronized(sharedAccountAddress.address)).resolves.toBe(true);
await expectTokenBalance(
sharedWalletOnB,
token.address,
token,
sharedAccountAddress.address,
transferAmount1 - transferAmount2,
logger,
);
});

Expand Down
1 change: 1 addition & 0 deletions yarn-project/end-to-end/src/fixtures/index.ts
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
export * from './fixtures.js';
export * from './logging.js';
export * from './utils.js';
export * from './token_utils.js';
61 changes: 61 additions & 0 deletions yarn-project/end-to-end/src/fixtures/token_utils.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
import { type AztecAddress, BatchCall, type DebugLogger, type Wallet, retryUntil } from '@aztec/aztec.js';
import { TokenContract } from '@aztec/noir-contracts.js';

export async function deployToken(adminWallet: Wallet, initialAdminBalance: bigint, logger: DebugLogger) {
logger.info(`Deploying Token contract...`);
const contract = await TokenContract.deploy(adminWallet, adminWallet.getAddress(), 'TokenName', 'TokenSymbol', 18)
.send()
.deployed();

if (initialAdminBalance > 0n) {
// Minter is minting to herself so contract as minter is the same as contract as recipient
await mintTokensToPrivate(contract, adminWallet, adminWallet.getAddress(), initialAdminBalance);
}

logger.info('L2 contract deployed');

return contract;
}

export async function mintTokensToPrivate(
token: TokenContract,
minterWallet: Wallet,
recipient: AztecAddress,
amount: bigint,
) {
// We don't have the functionality to mint to private so we mint to the minter address in public and transfer
// the tokens to the recipient in private. We use BatchCall to speed the process up.
await new BatchCall(minterWallet, [
token.methods.mint_public(minterWallet.getAddress(), amount).request(),
token.methods.transfer_to_private(recipient, amount).request(),
])
.send()
.wait();
}

const awaitUserSynchronized = async (wallet: Wallet, owner: AztecAddress) => {
const isUserSynchronized = async () => {
return await wallet.isAccountStateSynchronized(owner);
};
await retryUntil(isUserSynchronized, `synch of user ${owner.toString()}`, 10);
};

export async function expectTokenBalance(
wallet: Wallet,
token: TokenContract,
owner: AztecAddress,
expectedBalance: bigint,
logger: DebugLogger,
checkIfSynchronized = true,
) {
if (checkIfSynchronized) {
// First wait until the corresponding PXE has synchronized the account
await awaitUserSynchronized(wallet, owner);
}

// Then check the balance
const contractWithWallet = await TokenContract.at(token.address, wallet);
const balance = await contractWithWallet.methods.balance_of_private(owner).simulate({ from: owner });
logger.info(`Account ${owner} balance: ${balance}`);
expect(balance).toBe(expectedBalance);
}

0 comments on commit 618cd24

Please sign in to comment.