diff --git a/noir-projects/noir-contracts/contracts/token_contract/src/test/refunds.nr b/noir-projects/noir-contracts/contracts/token_contract/src/test/refunds.nr index e52d6fb44fa..b4238459484 100644 --- a/noir-projects/noir-contracts/contracts/token_contract/src/test/refunds.nr +++ b/noir-projects/noir-contracts/contracts/token_contract/src/test/refunds.nr @@ -1,8 +1,6 @@ use crate::{test::utils, Token}; use dep::authwit::cheatcodes as authwit_cheatcodes; -use dep::aztec::{prelude::NoteHeader, protocol_types::storage::map::derive_storage_slot_in_map}; -use dep::uint_note::uint_note::UintNote; use std::test::OracleMock; #[test] @@ -35,34 +33,23 @@ unconstrained fn setup_refund_success() { setup_refund_from_call_interface.call(&mut env.private()); - let fee_payer_balances_slot = - derive_storage_slot_in_map(Token::storage_layout().balances.slot, fee_payer); - let user_balances_slot = - derive_storage_slot_in_map(Token::storage_layout().balances.slot, user); - // When the refund was set up, we would've spent the note worth mint_amount, and inserted a note worth //`mint_amount - funded_amount`. When completing the refund, we would've constructed a hash corresponding to a note // worth `funded_amount - transaction_fee`. We "know" the transaction fee was 1 (it is hardcoded in // `executePublicFunction` TXE oracle) but we need to notify TXE of the note (preimage). - env.add_note( - &mut UintNote { - value: U128::from_integer(funded_amount - 1), - owner: user, - randomness: user_randomness, - header: NoteHeader::empty(), - }, - user_balances_slot, + utils::add_token_note( + env, token_contract_address, + fee_payer, + 1, + fee_payer_randomness, ); - env.add_note( - &mut UintNote { - value: U128::from_integer(1), - owner: fee_payer, - randomness: fee_payer_randomness, - header: NoteHeader::empty(), - }, - fee_payer_balances_slot, + utils::add_token_note( + env, token_contract_address, + user, + funded_amount - 1, + user_randomness, ); utils::check_private_balance(token_contract_address, user, mint_amount - 1); diff --git a/noir-projects/noir-contracts/contracts/token_contract/src/test/transfer_to_private.nr b/noir-projects/noir-contracts/contracts/token_contract/src/test/transfer_to_private.nr index e5b3463d83b..a02f4f3d4a2 100644 --- a/noir-projects/noir-contracts/contracts/token_contract/src/test/transfer_to_private.nr +++ b/noir-projects/noir-contracts/contracts/token_contract/src/test/transfer_to_private.nr @@ -1,9 +1,8 @@ use crate::{test::utils, Token}; use dep::aztec::{ - keys::getters::get_public_keys, oracle::random::random, prelude::NoteHeader, + keys::getters::get_public_keys, oracle::random::random, protocol_types::storage::map::derive_storage_slot_in_map, }; -use dep::uint_note::uint_note::UintNote; use std::test::OracleMock; /// Internal orchestration means that the calls to `prepare_transfer_to_private` @@ -45,20 +44,13 @@ unconstrained fn transfer_to_private_external_orchestration() { &mut env.public(), ); - // TODO(#8771): We need to manually add the note because in the partial notes flow `notify_created_note_oracle` - // is not called and we don't have a `NoteProcessor` in TXE. - let balances_owner_slot = - derive_storage_slot_in_map(Token::storage_layout().balances.slot, recipient); - - env.add_note( - &mut UintNote { - value: U128::from_integer(amount), - owner: recipient, - randomness: note_randomness, - header: NoteHeader::empty(), - }, - balances_owner_slot, + // We need to manually add the note because #8771 has not yet been implemented + utils::add_token_note( + env, token_contract_address, + recipient, + amount, + note_randomness, ); // Recipient's private balance should be equal to the amount diff --git a/noir-projects/noir-contracts/contracts/token_contract/src/test/utils.nr b/noir-projects/noir-contracts/contracts/token_contract/src/test/utils.nr index d1a77dd924a..53f894c85b5 100644 --- a/noir-projects/noir-contracts/contracts/token_contract/src/test/utils.nr +++ b/noir-projects/noir-contracts/contracts/token_contract/src/test/utils.nr @@ -1,16 +1,16 @@ -use crate::{Token, types::transparent_note::TransparentNote}; +use crate::Token; use dep::uint_note::uint_note::UintNote; use aztec::{ - keys::getters::get_public_keys, oracle::{ execution::{get_block_number, get_contract_address}, random::random, storage::storage_read, }, - prelude::AztecAddress, + prelude::{AztecAddress, NoteHeader}, protocol_types::storage::map::derive_storage_slot_in_map, test::helpers::{cheatcodes, test_environment::TestEnvironment}, }; +use std::test::OracleMock; pub unconstrained fn setup( with_account_contracts: bool, @@ -60,27 +60,33 @@ pub unconstrained fn setup_and_mint_private( with_account_contracts: bool, ) -> (&mut TestEnvironment, AztecAddress, AztecAddress, AztecAddress, Field) { // Setup the tokens and mint public balance - let (env, token_contract_address, owner, recipient) = - setup_and_mint_public(with_account_contracts); + let (env, token_contract_address, owner, recipient) = setup(with_account_contracts); + + // Mint some tokens let mint_amount = 10000; - // Transfer the public balance to private - Token::at(token_contract_address).transfer_to_private(owner, mint_amount).call( - &mut env.private(), - ); + mint_private(env, token_contract_address, owner, mint_amount); - // docs:start:txe_test_add_note - // TODO(#8771): We need to manually add the note because in the partial notes flow `notify_created_note_oracle` - // is not called and we don't have a `NoteProcessor` in TXE. - let balances_owner_slot = - derive_storage_slot_in_map(Token::storage_layout().balances.slot, owner); + (env, token_contract_address, owner, recipient, mint_amount) +} - env.add_note( - &mut UintNote::new(U128::from_integer(mint_amount), owner), - balances_owner_slot, +pub unconstrained fn mint_private( + env: &mut TestEnvironment, + token_contract_address: AztecAddress, + recipient: AztecAddress, + amount: Field, +) { + let note_randomness = random(); + let _ = OracleMock::mock("getRandomField").returns(note_randomness); + + Token::at(token_contract_address).mint_to_private(recipient, amount).call(&mut env.private()); + + add_token_note( + env, token_contract_address, + recipient, + amount, + note_randomness, ); - // docs:end:txe_test_add_note - (env, token_contract_address, owner, recipient, mint_amount) } // docs:start:txe_test_read_public @@ -115,3 +121,29 @@ pub unconstrained fn check_private_balance( cheatcodes::set_contract_address(current_contract_address); } // docs:end:txe_test_call_unconstrained + +// TODO(#8771): We need to manually add the note because in the partial notes flow `notify_created_note_oracle` +// is not called and we don't have a `NoteProcessor` in TXE. +pub unconstrained fn add_token_note( + env: &mut TestEnvironment, + token_contract_address: AztecAddress, + owner: AztecAddress, + amount: Field, + note_randomness: Field, +) { + // docs:start:txe_test_add_note + let balances_owner_slot = + derive_storage_slot_in_map(Token::storage_layout().balances.slot, owner); + + env.add_note( + &mut UintNote { + value: U128::from_integer(amount), + owner: owner, + randomness: note_randomness, + header: NoteHeader::empty(), + }, + balances_owner_slot, + token_contract_address, + ); + // docs:end:txe_test_add_note +}