From de02626a54cb6c1c4803bf3ff9d6ab2e8aebb944 Mon Sep 17 00:00:00 2001 From: LHerskind Date: Thu, 30 May 2024 11:43:23 +0000 Subject: [PATCH] update simulator --- .../access_control.test.ts | 4 +- .../blacklist_token_contract_test.ts | 1 + .../e2e_blacklist_token_contract/burn.test.ts | 2 +- .../minting.test.ts | 4 +- .../shielding.test.ts | 6 +- .../transfer_private.test.ts | 2 +- .../transfer_public.test.ts | 2 +- .../unshielding.test.ts | 2 +- .../src/e2e_lending_contract.test.ts | 6 +- .../src/e2e_prover/e2e_prover.test.ts | 6 +- .../src/e2e_prover/e2e_prover_test.ts | 1 + .../e2e_token_contract/access_control.test.ts | 4 +- .../src/e2e_token_contract/burn.test.ts | 2 +- .../src/e2e_token_contract/minting.test.ts | 2 +- .../reading_constants.test.ts | 5 +- .../src/e2e_token_contract/shielding.test.ts | 6 +- .../e2e_token_contract/token_contract_test.ts | 1 + .../transfer_private.test.ts | 2 +- .../transfer_public.test.ts | 2 +- .../e2e_token_contract/unshielding.test.ts | 2 +- .../src/simulators/lending_simulator.ts | 8 +-- .../src/simulators/token_simulator.ts | 71 ++++++++++++++++--- 22 files changed, 95 insertions(+), 46 deletions(-) diff --git a/yarn-project/end-to-end/src/e2e_blacklist_token_contract/access_control.test.ts b/yarn-project/end-to-end/src/e2e_blacklist_token_contract/access_control.test.ts index 6dfac5a88e50..4c1400caffe0 100644 --- a/yarn-project/end-to-end/src/e2e_blacklist_token_contract/access_control.test.ts +++ b/yarn-project/end-to-end/src/e2e_blacklist_token_contract/access_control.test.ts @@ -4,12 +4,10 @@ import { BlacklistTokenContractTest, Role } from './blacklist_token_contract_tes describe('e2e_blacklist_token_contract access control', () => { const t = new BlacklistTokenContractTest('access_control'); - let { wallets } = t; beforeAll(async () => { await t.applyBaseSnapshots(); await t.setup(); - ({ wallets } = t); }); afterAll(async () => { @@ -17,7 +15,7 @@ describe('e2e_blacklist_token_contract access control', () => { }); afterEach(async () => { - await t.tokenSim.check(wallets[0]); + await t.tokenSim.check(); }); it('grant mint permission to the admin', async () => { diff --git a/yarn-project/end-to-end/src/e2e_blacklist_token_contract/blacklist_token_contract_test.ts b/yarn-project/end-to-end/src/e2e_blacklist_token_contract/blacklist_token_contract_test.ts index 708f1cbfb3dd..8b7db80d2134 100644 --- a/yarn-project/end-to-end/src/e2e_blacklist_token_contract/blacklist_token_contract_test.ts +++ b/yarn-project/end-to-end/src/e2e_blacklist_token_contract/blacklist_token_contract_test.ts @@ -121,6 +121,7 @@ export class BlacklistTokenContractTest { this.tokenSim = new TokenSimulator( this.asset as unknown as TokenContract, + this.wallets[0], this.logger, this.accounts.map(a => a.address), ); diff --git a/yarn-project/end-to-end/src/e2e_blacklist_token_contract/burn.test.ts b/yarn-project/end-to-end/src/e2e_blacklist_token_contract/burn.test.ts index 865a53213222..05ca22f844fa 100644 --- a/yarn-project/end-to-end/src/e2e_blacklist_token_contract/burn.test.ts +++ b/yarn-project/end-to-end/src/e2e_blacklist_token_contract/burn.test.ts @@ -21,7 +21,7 @@ describe('e2e_blacklist_token_contract burn', () => { }); afterEach(async () => { - await t.tokenSim.check(wallets[0]); + await t.tokenSim.check(); }); describe('public', () => { diff --git a/yarn-project/end-to-end/src/e2e_blacklist_token_contract/minting.test.ts b/yarn-project/end-to-end/src/e2e_blacklist_token_contract/minting.test.ts index e2979d991074..ab535e3e61ab 100644 --- a/yarn-project/end-to-end/src/e2e_blacklist_token_contract/minting.test.ts +++ b/yarn-project/end-to-end/src/e2e_blacklist_token_contract/minting.test.ts @@ -21,11 +21,11 @@ describe('e2e_blacklist_token_contract mint', () => { }); beforeEach(async () => { - await t.tokenSim.check(wallets[0]); + await t.tokenSim.check(); }); afterEach(async () => { - await t.tokenSim.check(wallets[0]); + await t.tokenSim.check(); }); describe('Public', () => { diff --git a/yarn-project/end-to-end/src/e2e_blacklist_token_contract/shielding.test.ts b/yarn-project/end-to-end/src/e2e_blacklist_token_contract/shielding.test.ts index ef289f6ffb90..22701ac0cbf7 100644 --- a/yarn-project/end-to-end/src/e2e_blacklist_token_contract/shielding.test.ts +++ b/yarn-project/end-to-end/src/e2e_blacklist_token_contract/shielding.test.ts @@ -20,7 +20,7 @@ describe('e2e_blacklist_token_contract shield + redeem_shield', () => { }); afterEach(async () => { - await t.tokenSim.check(wallets[0]); + await t.tokenSim.check(); }); const secret = Fr.random(); @@ -38,7 +38,7 @@ describe('e2e_blacklist_token_contract shield + redeem_shield', () => { const receipt = await asset.methods.shield(wallets[0].getAddress(), amount, secretHash, 0).send().wait(); tokenSim.shield(wallets[0].getAddress(), amount); - await t.tokenSim.check(wallets[0]); + await t.tokenSim.check(); // Redeem it await t.addPendingShieldNoteToPXE(0, amount, secretHash, receipt.txHash); @@ -60,7 +60,7 @@ describe('e2e_blacklist_token_contract shield + redeem_shield', () => { const receipt = await action.send().wait(); tokenSim.shield(wallets[0].getAddress(), amount); - await t.tokenSim.check(wallets[0]); + await t.tokenSim.check(); // Check that replaying the shield should fail! const txReplay = asset diff --git a/yarn-project/end-to-end/src/e2e_blacklist_token_contract/transfer_private.test.ts b/yarn-project/end-to-end/src/e2e_blacklist_token_contract/transfer_private.test.ts index cfee467902af..e69813b6fedf 100644 --- a/yarn-project/end-to-end/src/e2e_blacklist_token_contract/transfer_private.test.ts +++ b/yarn-project/end-to-end/src/e2e_blacklist_token_contract/transfer_private.test.ts @@ -21,7 +21,7 @@ describe('e2e_blacklist_token_contract transfer private', () => { }); afterEach(async () => { - await t.tokenSim.check(wallets[0]); + await t.tokenSim.check(); }); it('transfer less than balance', async () => { diff --git a/yarn-project/end-to-end/src/e2e_blacklist_token_contract/transfer_public.test.ts b/yarn-project/end-to-end/src/e2e_blacklist_token_contract/transfer_public.test.ts index 8b2ac023d0ec..0b9fd8463764 100644 --- a/yarn-project/end-to-end/src/e2e_blacklist_token_contract/transfer_public.test.ts +++ b/yarn-project/end-to-end/src/e2e_blacklist_token_contract/transfer_public.test.ts @@ -21,7 +21,7 @@ describe('e2e_blacklist_token_contract transfer public', () => { }); afterEach(async () => { - await t.tokenSim.check(wallets[0]); + await t.tokenSim.check(); }); it('transfer less than balance', async () => { diff --git a/yarn-project/end-to-end/src/e2e_blacklist_token_contract/unshielding.test.ts b/yarn-project/end-to-end/src/e2e_blacklist_token_contract/unshielding.test.ts index 346a3ebba3f8..224a26b5f0f2 100644 --- a/yarn-project/end-to-end/src/e2e_blacklist_token_contract/unshielding.test.ts +++ b/yarn-project/end-to-end/src/e2e_blacklist_token_contract/unshielding.test.ts @@ -21,7 +21,7 @@ describe('e2e_blacklist_token_contract unshielding', () => { }); afterEach(async () => { - await t.tokenSim.check(wallets[0]); + await t.tokenSim.check(); }); it('on behalf of self', async () => { diff --git a/yarn-project/end-to-end/src/e2e_lending_contract.test.ts b/yarn-project/end-to-end/src/e2e_lending_contract.test.ts index 98a3ae7854fa..7e22eecf3585 100644 --- a/yarn-project/end-to-end/src/e2e_lending_contract.test.ts +++ b/yarn-project/end-to-end/src/e2e_lending_contract.test.ts @@ -74,15 +74,15 @@ describe('e2e_lending_contract', () => { lendingAccount, rate, lendingContract, - new TokenSimulator(collateralAsset, logger, [lendingContract.address, wallet.getAddress()]), - new TokenSimulator(stableCoin, logger, [lendingContract.address, wallet.getAddress()]), + new TokenSimulator(collateralAsset, wallet, logger, [lendingContract.address, wallet.getAddress()]), + new TokenSimulator(stableCoin, wallet, logger, [lendingContract.address, wallet.getAddress()]), ); }, 300_000); afterAll(() => teardown()); afterEach(async () => { - await lendingSim.check(wallet); + await lendingSim.check(); }); it('Mint assets for later usage', async () => { diff --git a/yarn-project/end-to-end/src/e2e_prover/e2e_prover.test.ts b/yarn-project/end-to-end/src/e2e_prover/e2e_prover.test.ts index ef36701bfae8..8c1441724c35 100644 --- a/yarn-project/end-to-end/src/e2e_prover/e2e_prover.test.ts +++ b/yarn-project/end-to-end/src/e2e_prover/e2e_prover.test.ts @@ -7,14 +7,14 @@ const TIMEOUT = 1_800_000; describe('full_prover', () => { const t = new FullProverTest('full_prover'); - let { provenAssets, accounts, tokenSim, logger, wallets } = t; + let { provenAssets, accounts, tokenSim, logger } = t; beforeAll(async () => { await t.applyBaseSnapshots(); await t.applyMintSnapshot(); await t.setup(); await t.deployVerifier(); - ({ provenAssets, accounts, tokenSim, logger, wallets } = t); + ({ provenAssets, accounts, tokenSim, logger } = t); }); afterAll(async () => { @@ -22,7 +22,7 @@ describe('full_prover', () => { }); afterEach(async () => { - await t.tokenSim.check(wallets[0]); + await t.tokenSim.check(); }); it( diff --git a/yarn-project/end-to-end/src/e2e_prover/e2e_prover_test.ts b/yarn-project/end-to-end/src/e2e_prover/e2e_prover_test.ts index 794f9f3d60a3..f01fd3cca336 100644 --- a/yarn-project/end-to-end/src/e2e_prover/e2e_prover_test.ts +++ b/yarn-project/end-to-end/src/e2e_prover/e2e_prover_test.ts @@ -120,6 +120,7 @@ export class FullProverTest { this.tokenSim = new TokenSimulator( this.fakeProofsAsset, + this.wallets[0], this.logger, this.accounts.map(a => a.address), ); diff --git a/yarn-project/end-to-end/src/e2e_token_contract/access_control.test.ts b/yarn-project/end-to-end/src/e2e_token_contract/access_control.test.ts index 72e3155ab629..8efc7d88277e 100644 --- a/yarn-project/end-to-end/src/e2e_token_contract/access_control.test.ts +++ b/yarn-project/end-to-end/src/e2e_token_contract/access_control.test.ts @@ -2,12 +2,10 @@ import { TokenContractTest } from './token_contract_test.js'; describe('e2e_token_contract access control', () => { const t = new TokenContractTest('access_control'); - let { wallets } = t; beforeAll(async () => { await t.applyBaseSnapshots(); await t.setup(); - ({ wallets } = t); }); afterAll(async () => { @@ -15,7 +13,7 @@ describe('e2e_token_contract access control', () => { }); afterEach(async () => { - await t.tokenSim.check(wallets[0]); + await t.tokenSim.check(); }); it('Set admin', async () => { diff --git a/yarn-project/end-to-end/src/e2e_token_contract/burn.test.ts b/yarn-project/end-to-end/src/e2e_token_contract/burn.test.ts index ed605ac69560..bfe3406329c6 100644 --- a/yarn-project/end-to-end/src/e2e_token_contract/burn.test.ts +++ b/yarn-project/end-to-end/src/e2e_token_contract/burn.test.ts @@ -20,7 +20,7 @@ describe('e2e_token_contract burn', () => { }); afterEach(async () => { - await t.tokenSim.check(wallets[0]); + await t.tokenSim.check(); }); describe('public', () => { diff --git a/yarn-project/end-to-end/src/e2e_token_contract/minting.test.ts b/yarn-project/end-to-end/src/e2e_token_contract/minting.test.ts index c093b5b05dc8..7a833d6a1573 100644 --- a/yarn-project/end-to-end/src/e2e_token_contract/minting.test.ts +++ b/yarn-project/end-to-end/src/e2e_token_contract/minting.test.ts @@ -18,7 +18,7 @@ describe('e2e_token_contract minting', () => { }); afterEach(async () => { - await t.tokenSim.check(wallets[0]); + await t.tokenSim.check(); }); describe('Public', () => { diff --git a/yarn-project/end-to-end/src/e2e_token_contract/reading_constants.test.ts b/yarn-project/end-to-end/src/e2e_token_contract/reading_constants.test.ts index e4fbc22e35c7..7fdfd2b2f829 100644 --- a/yarn-project/end-to-end/src/e2e_token_contract/reading_constants.test.ts +++ b/yarn-project/end-to-end/src/e2e_token_contract/reading_constants.test.ts @@ -17,13 +17,10 @@ const toString = ({ value }: { value: bigint }) => { describe('e2e_token_contract reading constants', () => { const t = new TokenContractTest('reading_constants'); const { TOKEN_DECIMALS, TOKEN_NAME, TOKEN_SYMBOL } = TokenContractTest; - let { wallets } = t; beforeAll(async () => { await t.applyBaseSnapshots(); await t.setup(); - - ({ wallets } = t); }); afterAll(async () => { @@ -33,7 +30,7 @@ describe('e2e_token_contract reading constants', () => { beforeEach(async () => {}); afterEach(async () => { - await t.tokenSim.check(wallets[0]); + await t.tokenSim.check(); }); it('check name private', async () => { diff --git a/yarn-project/end-to-end/src/e2e_token_contract/shielding.test.ts b/yarn-project/end-to-end/src/e2e_token_contract/shielding.test.ts index 2438f1554db3..93ab4e448708 100644 --- a/yarn-project/end-to-end/src/e2e_token_contract/shielding.test.ts +++ b/yarn-project/end-to-end/src/e2e_token_contract/shielding.test.ts @@ -23,7 +23,7 @@ describe('e2e_token_contract shield + redeem shield', () => { }); afterEach(async () => { - await t.tokenSim.check(wallets[0]); + await t.tokenSim.check(); }); it('on behalf of self', async () => { @@ -34,7 +34,7 @@ describe('e2e_token_contract shield + redeem shield', () => { const receipt = await asset.methods.shield(accounts[0].address, amount, secretHash, 0).send().wait(); tokenSim.shield(accounts[0].address, amount); - await tokenSim.check(wallets[0]); + await tokenSim.check(); // Redeem it await t.addPendingShieldNoteToPXE(0, amount, secretHash, receipt.txHash); @@ -56,7 +56,7 @@ describe('e2e_token_contract shield + redeem shield', () => { const receipt = await action.send().wait(); tokenSim.shield(accounts[0].address, amount); - await tokenSim.check(wallets[0]); + await tokenSim.check(); // Check that replaying the shield should fail! const txReplay = asset.withWallet(wallets[1]).methods.shield(accounts[0].address, amount, secretHash, nonce).send(); diff --git a/yarn-project/end-to-end/src/e2e_token_contract/token_contract_test.ts b/yarn-project/end-to-end/src/e2e_token_contract/token_contract_test.ts index 5e7c977c8458..6510e292f5c1 100644 --- a/yarn-project/end-to-end/src/e2e_token_contract/token_contract_test.ts +++ b/yarn-project/end-to-end/src/e2e_token_contract/token_contract_test.ts @@ -86,6 +86,7 @@ export class TokenContractTest { this.tokenSim = new TokenSimulator( this.asset, + this.wallets[0], this.logger, this.accounts.map(a => a.address), ); diff --git a/yarn-project/end-to-end/src/e2e_token_contract/transfer_private.test.ts b/yarn-project/end-to-end/src/e2e_token_contract/transfer_private.test.ts index 647cd6e013b9..fb5394567dc2 100644 --- a/yarn-project/end-to-end/src/e2e_token_contract/transfer_private.test.ts +++ b/yarn-project/end-to-end/src/e2e_token_contract/transfer_private.test.ts @@ -19,7 +19,7 @@ describe('e2e_token_contract transfer private', () => { }); afterEach(async () => { - await t.tokenSim.check(wallets[0]); + await t.tokenSim.check(); }); it('transfer less than balance', async () => { diff --git a/yarn-project/end-to-end/src/e2e_token_contract/transfer_public.test.ts b/yarn-project/end-to-end/src/e2e_token_contract/transfer_public.test.ts index 037c3ce29d94..cb352c57df2d 100644 --- a/yarn-project/end-to-end/src/e2e_token_contract/transfer_public.test.ts +++ b/yarn-project/end-to-end/src/e2e_token_contract/transfer_public.test.ts @@ -20,7 +20,7 @@ describe('e2e_token_contract transfer public', () => { }); afterEach(async () => { - await t.tokenSim.check(wallets[0]); + await t.tokenSim.check(); }); it('transfer less than balance', async () => { diff --git a/yarn-project/end-to-end/src/e2e_token_contract/unshielding.test.ts b/yarn-project/end-to-end/src/e2e_token_contract/unshielding.test.ts index 3dc03f2ece76..d52b3ce214e4 100644 --- a/yarn-project/end-to-end/src/e2e_token_contract/unshielding.test.ts +++ b/yarn-project/end-to-end/src/e2e_token_contract/unshielding.test.ts @@ -20,7 +20,7 @@ describe('e2e_token_contract unshielding', () => { }); afterEach(async () => { - await t.tokenSim.check(wallets[0]); + await t.tokenSim.check(); }); it('on behalf of self', async () => { diff --git a/yarn-project/end-to-end/src/simulators/lending_simulator.ts b/yarn-project/end-to-end/src/simulators/lending_simulator.ts index b9bea8c7a362..95805f004211 100644 --- a/yarn-project/end-to-end/src/simulators/lending_simulator.ts +++ b/yarn-project/end-to-end/src/simulators/lending_simulator.ts @@ -1,5 +1,5 @@ // Convenience struct to hold an account's address and secret that can easily be passed around. -import { type AztecAddress, type CheatCodes, Fr, type Wallet } from '@aztec/aztec.js'; +import { type AztecAddress, type CheatCodes, Fr } from '@aztec/aztec.js'; import { pedersenHash } from '@aztec/foundation/crypto'; import { type LendingContract } from '@aztec/noir-contracts.js/Lending'; @@ -154,10 +154,10 @@ export class LendingSimulator { this.mintedOutside += amount; } - async check(wallet: Wallet) { + async check() { // Run checks on both underlying assets - await this.collateralAsset.check(wallet); - await this.stableCoin.check(wallet); + await this.collateralAsset.check(); + await this.stableCoin.check(); // Check that total collateral equals total holdings by contract. const totalCollateral = Object.values(this.collateral).reduce((a, b) => new Fr(a.value + b.value), Fr.ZERO); diff --git a/yarn-project/end-to-end/src/simulators/token_simulator.ts b/yarn-project/end-to-end/src/simulators/token_simulator.ts index b9ef2a88101f..c4addad4b217 100644 --- a/yarn-project/end-to-end/src/simulators/token_simulator.ts +++ b/yarn-project/end-to-end/src/simulators/token_simulator.ts @@ -9,7 +9,22 @@ export class TokenSimulator { private balancePublic: Map = new Map(); public totalSupply: bigint = 0n; - constructor(protected token: TokenContract, protected logger: DebugLogger, protected accounts: AztecAddress[]) {} + private lookupProvider: Map = new Map(); + + constructor( + protected token: TokenContract, + protected defaultWallet: Wallet, + protected logger: DebugLogger, + protected accounts: AztecAddress[], + ) {} + + public addAccount(account: AztecAddress) { + this.accounts.push(account); + } + + public setLookupProvider(account: AztecAddress, wallet: Wallet) { + this.lookupProvider.set(account.toString(), wallet); + } public mintPrivate(amount: bigint) { this.totalSupply += amount; @@ -82,23 +97,61 @@ export class TokenSimulator { return this.balancesPrivate.get(address.toString()) || 0n; } - public async check(wallet: Wallet) { + async checkPublic() { + // public calls const calls = [this.token.methods.total_supply().request()]; for (const address of this.accounts) { calls.push(this.token.methods.balance_of_public(address).request()); - calls.push(this.token.methods.balance_of_private(address).request()); } - const batchedCalls = chunk(calls, 4); - - const results = (await Promise.all(batchedCalls.map(batch => new BatchCall(wallet, batch).simulate()))).flat(); - + const results = ( + await Promise.all(chunk(calls, 4).map(batch => new BatchCall(this.defaultWallet, batch).simulate())) + ).flat(); expect(results[0]).toEqual(this.totalSupply); // Check that all our balances match for (let i = 0; i < this.accounts.length; i++) { - expect(results[i * 2 + 1]).toEqual(this.balanceOfPublic(this.accounts[i])); - expect(results[i * 2 + 2]).toEqual(this.balanceOfPrivate(this.accounts[i])); + expect(results[i + 1]).toEqual(this.balanceOfPublic(this.accounts[i])); + } + } + + async checkPrivate() { + // Private calls + const defaultLookups = []; + const nonDefaultLookups = []; + + for (const address of this.accounts) { + if (this.lookupProvider.has(address.toString())) { + nonDefaultLookups.push(address); + } else { + defaultLookups.push(address); + } + } + + const defaultCalls = []; + for (const address of defaultLookups) { + defaultCalls.push(this.token.methods.balance_of_private(address).request()); + } + const results = ( + await Promise.all(chunk(defaultCalls, 4).map(batch => new BatchCall(this.defaultWallet, batch).simulate())) + ).flat(); + for (let i = 0; i < defaultLookups.length; i++) { + expect(results[i]).toEqual(this.balanceOfPrivate(defaultLookups[i])); } + + // We are just running individual calls for the non-default lookups + // @todo We should also batch these + for (const address of nonDefaultLookups) { + const wallet = this.lookupProvider.get(address.toString()); + const asset = wallet ? this.token.withWallet(wallet) : this.token; + + const actualPrivateBalance = await asset.methods.balance_of_private({ address }).simulate(); + expect(actualPrivateBalance).toEqual(this.balanceOfPrivate(address)); + } + } + + public async check() { + await this.checkPublic(); + await this.checkPrivate(); } }