test: cleaning up token test utils

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);

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

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
+}