Skip to main content

miden_testing/mock_chain/
chain_builder.rs

1use alloc::collections::{BTreeMap, BTreeSet};
2use alloc::vec::Vec;
3
4use anyhow::Context;
5
6// CONSTANTS
7// ================================================================================================
8
9/// Default number of decimals for faucets created in tests.
10const DEFAULT_FAUCET_DECIMALS: u8 = 10;
11
12/// Default number of validators committed to by the genesis block of a mock chain.
13///
14/// This is purely a test default -- the protocol does not fix the size of a validator set.
15const DEFAULT_VALIDATOR_COUNT: usize = 3;
16
17// IMPORTS
18// ================================================================================================
19
20use itertools::Itertools;
21use miden_processor::crypto::random::RandomCoin;
22use miden_protocol::account::{
23    Account,
24    AccountBuilder,
25    AccountComponent,
26    AccountId,
27    AccountPatch,
28    AccountType,
29    AccountUpdateDetails,
30    AssetCallbackFlag,
31    StorageSlot,
32};
33use miden_protocol::asset::{Asset, AssetAmount, FungibleAsset, TokenSymbol};
34use miden_protocol::block::account_tree::AccountTree;
35use miden_protocol::block::nullifier_tree::NullifierTree;
36use miden_protocol::block::{
37    BlockAccountUpdate,
38    BlockBody,
39    BlockHeader,
40    BlockNoteTree,
41    BlockNumber,
42    BlockProof,
43    BlockSignatures,
44    Blockchain,
45    FeeParameters,
46    OutputNoteBatch,
47    ProvenBlock,
48    ValidatorKeys,
49};
50use miden_protocol::crypto::dsa::ecdsa_k256_keccak::SigningKey;
51use miden_protocol::crypto::merkle::smt::Smt;
52use miden_protocol::errors::NoteError;
53use miden_protocol::note::{Note, NoteAttachments, NoteDetails, NoteScriptRoot, NoteType};
54use miden_protocol::testing::account_id::ACCOUNT_ID_FEE_FAUCET;
55use miden_protocol::testing::random_secret_key::random_secret_key;
56use miden_protocol::transaction::{OrderedTransactionHeaders, RawOutputNote, TransactionKernel};
57use miden_protocol::{MAX_OUTPUT_NOTES_PER_BATCH, Word};
58use miden_standards::account::access::{AccessControl, Authority, Pausable, PausableManager};
59use miden_standards::account::auth::AuthNetworkAccount;
60use miden_standards::account::faucets::{FungibleFaucet, TokenName};
61use miden_standards::account::fees::{ConstantFeePolicy, FeeManager};
62use miden_standards::account::policies::{
63    BurnPolicy,
64    MintPolicy,
65    TokenPolicyManager,
66    TransferPolicy,
67};
68use miden_standards::account::wallets::{BasicWallet, NoteCreator};
69use miden_standards::note::{BurnNote, MintNote, P2idNote, P2ideNote, SwapNote, TxFeeNote};
70use miden_standards::testing::account_component::MockAccountComponent;
71use rand::RngExt;
72
73use crate::mock_chain::chain::AccountAuthenticator;
74use crate::utils::{create_p2any_note, create_spawn_note};
75use crate::{AccountState, Auth, MockChain};
76
77/// A builder for a [`MockChain`]'s genesis block.
78///
79/// ## Example
80///
81/// ```
82/// # use anyhow::Result;
83/// # use miden_protocol::{
84/// #    asset::{Asset, FungibleAsset},
85/// #    note::NoteType,
86/// # };
87/// # use miden_testing::{Auth, MockChain};
88/// #
89/// # fn main() -> Result<()> {
90/// let mut builder = MockChain::builder();
91/// let existing_wallet =
92///     builder.add_existing_wallet_with_assets(Auth::IncrNonce, [FungibleAsset::mock(500)])?;
93/// let new_wallet = builder.create_new_wallet(Auth::IncrNonce)?;
94///
95/// let existing_note = builder.add_p2id_note(
96///     existing_wallet.id(),
97///     new_wallet.id(),
98///     &[FungibleAsset::mock(100)],
99///     NoteType::Private,
100/// )?;
101/// let chain = builder.build()?;
102///
103/// // The existing wallet and note should be part of the chain state.
104/// assert!(chain.committed_account(existing_wallet.id()).is_ok());
105/// assert!(chain.committed_notes().get(&existing_note.id()).is_some());
106///
107/// // The new wallet should *not* be part of the chain state - it must be created in
108/// // a transaction first.
109/// assert!(chain.committed_account(new_wallet.id()).is_err());
110///
111/// # Ok(())
112/// # }
113/// ```
114///
115/// Note the distinction between `add_` and `create_` APIs. Any `add_` APIs will add something to
116/// the genesis chain state while `create_` APIs do not mutate the genesis state. The latter are
117/// simply convenient for creating accounts or notes that will be created by transactions.
118///
119/// See also the [`MockChain`] docs for examples on using the mock chain.
120#[derive(Debug, Clone)]
121pub struct MockChainBuilder {
122    accounts: BTreeMap<AccountId, Account>,
123    account_authenticators: BTreeMap<AccountId, AccountAuthenticator>,
124    notes: Vec<RawOutputNote>,
125    rng: RandomCoin,
126    // Fee parameters.
127    fee_faucet_id: AccountId,
128    verification_base_fee: u32,
129}
130
131impl MockChainBuilder {
132    // CONSTRUCTORS
133    // ----------------------------------------------------------------------------------------
134
135    /// Initializes a new mock chain builder with an empty state.
136    ///
137    /// By default, the `fee_faucet_id` is set to [`ACCOUNT_ID_FEE_FAUCET`] and can be
138    /// overwritten using [`Self::fee_faucet_id`].
139    ///
140    /// The `verification_base_fee` is initialized to 0 which means no fees are required by default.
141    pub fn new() -> Self {
142        let fee_faucet_id = ACCOUNT_ID_FEE_FAUCET.try_into().expect("account ID should be valid");
143
144        Self {
145            accounts: BTreeMap::new(),
146            account_authenticators: BTreeMap::new(),
147            notes: Vec::new(),
148            rng: RandomCoin::new(Default::default()),
149            fee_faucet_id,
150            verification_base_fee: 0,
151        }
152    }
153
154    /// Initializes a new mock chain builder with the provided accounts.
155    ///
156    /// This method only adds the accounts and cannot not register any authenticators for them.
157    /// Calling [`MockChain::build_tx_context`] on accounts added in this way will not work if the
158    /// account needs an authenticator.
159    ///
160    /// Due to these limitations, prefer using other methods to add accounts to the chain, e.g.
161    /// [`MockChainBuilder::add_account_from_builder`].
162    pub fn with_accounts(accounts: impl IntoIterator<Item = Account>) -> anyhow::Result<Self> {
163        let mut builder = Self::new();
164
165        for account in accounts {
166            builder.add_account(account)?;
167        }
168
169        Ok(builder)
170    }
171
172    // BUILDER METHODS
173    // ----------------------------------------------------------------------------------------
174
175    /// Sets the fee faucet ID of the chain.
176    ///
177    /// This must be a fungible faucet [`AccountId`] and is the asset in which fees will be accepted
178    /// by the transaction kernel.
179    pub fn fee_faucet_id(mut self, fee_faucet_id: AccountId) -> Self {
180        self.fee_faucet_id = fee_faucet_id;
181        self
182    }
183
184    /// Sets the `verification_base_fee` of the chain.
185    ///
186    /// See [`FeeParameters`] for more details.
187    pub fn verification_base_fee(mut self, verification_base_fee: u32) -> Self {
188        self.verification_base_fee = verification_base_fee;
189        self
190    }
191
192    /// Consumes the builder, creates the genesis block of the chain and returns the [`MockChain`].
193    pub fn build(self) -> anyhow::Result<MockChain> {
194        // Create the genesis block, consisting of the provided accounts and notes.
195        let block_account_updates: Vec<BlockAccountUpdate> = self
196            .accounts
197            .into_values()
198            .map(|account| {
199                let account_id = account.id();
200                let account_commitment = account.to_commitment();
201                let account_patch = AccountPatch::try_from(account)
202                    .expect("chain builder should only store existing accounts without seeds");
203                let update_details = AccountUpdateDetails::Public(account_patch);
204
205                BlockAccountUpdate::new(account_id, account_commitment, update_details)
206            })
207            .collect();
208
209        let account_tree = AccountTree::with_entries(
210            block_account_updates
211                .iter()
212                .map(|account| (account.account_id(), account.final_state_commitment())),
213        )
214        .context("failed to create genesis account tree")?;
215
216        // Extract full notes before shrinking for later use in MockChain
217        let full_notes: Vec<Note> = self
218            .notes
219            .iter()
220            .filter_map(|note| match note {
221                RawOutputNote::Full(n) => Some(n.clone()),
222                _ => None,
223            })
224            .collect();
225
226        let proven_notes: Vec<_> = self
227            .notes
228            .into_iter()
229            .map(|note| note.into_output_note().expect("genesis note should be valid"))
230            .collect();
231        let note_chunks = proven_notes.into_iter().chunks(MAX_OUTPUT_NOTES_PER_BATCH);
232        let output_note_batches: Vec<OutputNoteBatch> = note_chunks
233            .into_iter()
234            .map(|batch_notes| batch_notes.into_iter().enumerate().collect::<Vec<_>>())
235            .collect();
236
237        let created_nullifiers = Vec::new();
238        let transactions = OrderedTransactionHeaders::new_unchecked(Vec::new());
239
240        let note_tree = BlockNoteTree::from_note_batches(&output_note_batches)
241            .context("failed to create block note tree")?;
242
243        let version = 0;
244        let prev_block_commitment = Word::empty();
245        let block_num = BlockNumber::from(0u32);
246        let chain_commitment = Blockchain::new().commitment();
247        let account_root = account_tree.root();
248        let nullifier_root = NullifierTree::<Smt>::default().root();
249        let note_root = note_tree.root();
250        let tx_commitment = transactions.commitment();
251        let tx_kernel_commitment = TransactionKernel.to_commitment();
252        let timestamp = MockChain::TIMESTAMP_START_SECS;
253        let fee_parameters = FeeParameters::new(self.fee_faucet_id, self.verification_base_fee);
254        let validator_secret_keys: Vec<SigningKey> =
255            (0..DEFAULT_VALIDATOR_COUNT).map(|_| random_secret_key()).collect();
256        let validator_keys =
257            ValidatorKeys::new(validator_secret_keys.iter().map(|sk| sk.public_key()).collect())
258                .expect("randomly generated genesis validator keys should be distinct");
259
260        let header = BlockHeader::new(
261            version,
262            prev_block_commitment,
263            block_num,
264            chain_commitment,
265            account_root,
266            nullifier_root,
267            note_root,
268            tx_commitment,
269            tx_kernel_commitment,
270            validator_keys.clone(),
271            fee_parameters,
272            timestamp,
273        );
274
275        let body = BlockBody::new_unchecked(
276            block_account_updates,
277            output_note_batches,
278            created_nullifiers,
279            transactions,
280        );
281
282        // The genesis block is the trust root: it is self-signed by the validator set it commits
283        // as the signer of block 1.
284        let signatures = BlockSignatures::new(
285            validator_keys
286                .as_keys()
287                .iter()
288                .map(|key| {
289                    let signer = validator_secret_keys
290                        .iter()
291                        .find(|sk| &sk.public_key() == key)
292                        .expect("a signer should exist for every validator key");
293                    signer.sign(header.commitment())
294                })
295                .collect(),
296        )
297        .expect("signature count same as validator key count");
298        let block_proof = BlockProof::new_dummy();
299        let genesis_block = ProvenBlock::new_unchecked(header, body, signatures, block_proof);
300
301        MockChain::from_genesis_block(
302            genesis_block,
303            account_tree,
304            self.account_authenticators,
305            validator_secret_keys,
306            full_notes,
307        )
308    }
309
310    // ACCOUNT METHODS
311    // ----------------------------------------------------------------------------------------
312
313    /// Creates a new public [`BasicWallet`] account and registers the authenticator (if any) for
314    /// it.
315    ///
316    /// This does not add the account to the chain state, but it can still be used to call
317    /// [`MockChain::build_tx_context`] to automatically add the authenticator.
318    pub fn create_new_wallet(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
319        let account_builder = AccountBuilder::new(self.rng.random())
320            .account_type(AccountType::Public)
321            .with_component(BasicWallet);
322
323        self.add_account_from_builder(auth_method, account_builder, AccountState::New)
324    }
325
326    /// Adds an existing public [`BasicWallet`] account to the initial chain state and registers the
327    /// authenticator (if any).
328    pub fn add_existing_wallet(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
329        self.add_existing_wallet_with_assets(auth_method, [])
330    }
331
332    /// Adds an existing public [`BasicWallet`] account to the initial chain state and registers the
333    /// authenticator (if any).
334    pub fn add_existing_wallet_with_assets(
335        &mut self,
336        auth_method: Auth,
337        assets: impl IntoIterator<Item = Asset>,
338    ) -> anyhow::Result<Account> {
339        let account_builder = Account::builder(self.rng.random())
340            .account_type(AccountType::Public)
341            .with_component(BasicWallet)
342            .with_assets(assets);
343
344        self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)
345    }
346
347    /// Adds an existing public [`NoteCreator`] account to the initial chain state and registers the
348    /// authenticator (if any).
349    ///
350    /// Unlike [`add_existing_wallet`](Self::add_existing_wallet), the account exposes only the
351    /// `create_note` procedure, which is enough for tests that only create output notes.
352    pub fn add_existing_note_creator(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
353        let account_builder = Account::builder(self.rng.random())
354            .account_type(AccountType::Public)
355            .with_component(NoteCreator);
356
357        self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)
358    }
359
360    /// Internal helper: adds an existing network-style fungible faucet (Ownable2Step / Rbac).
361    /// Bundles [`PausableManager`] to match the `create_network_fungible_faucet` factory.
362    fn add_existing_network_fungible_faucet(
363        &mut self,
364        auth: AuthNetworkAccount,
365        faucet: FungibleFaucet,
366        account_type: AccountType,
367        access_control: AccessControl,
368        token_policy_manager: TokenPolicyManager,
369    ) -> anyhow::Result<Account> {
370        // network faucets authenticate with AuthNetworkAccount, which collects sponsored fees and
371        // answers sponsorship fee estimates; both require an active fee policy. The empty schedule
372        // keeps this a no-op on fee-free chains.
373        let mut constant_fee_policy = ConstantFeePolicy::new();
374        for note_script in auth.allowed_notes().allowed_script_roots() {
375            constant_fee_policy = constant_fee_policy.with_fee(*note_script, AssetAmount::ZERO)
376        }
377
378        let fee_manager = FeeManager::builder()
379            .active_fee_policy(constant_fee_policy.into())
380            .fee_faucet_id(self.fee_faucet_id)
381            .build();
382
383        let account_builder = AccountBuilder::new(self.rng.random())
384            .account_type(account_type)
385            .with_component(faucet)
386            .with_components(access_control)
387            .with_asset_callbacks(AssetCallbackFlag::from(
388                token_policy_manager.has_transfer_policy(),
389            ))
390            .with_components(token_policy_manager)
391            .with_component(Pausable::unpaused())
392            .with_component(PausableManager)
393            .with_components(fee_manager);
394
395        let auth = Auth::NetworkAccount {
396            allowed_script_roots: auth.allowed_notes().allowed_script_roots().clone(),
397            allowed_tx_script_roots: auth.allowed_tx_scripts().allowed_script_roots().clone(),
398        };
399
400        self.add_account_from_builder(auth, account_builder, AccountState::Exists)
401    }
402
403    /// Convenience: builds a basic auth-controlled fungible faucet from a token-symbol shorthand
404    /// using default decimals and `AllowAll` policies, then adds it as an existing account with
405    /// [`Authority::AuthControlled`].
406    ///
407    /// The faucet installs only `AllowAll` mint and burn policies and no transfer policy, so its
408    /// account ID has asset callbacks disabled and its assets transfer freely without triggering a
409    /// faucet callback. For a faucet with transfer policies (and thus callbacks), construct a
410    /// [`FungibleFaucet`] with a [`TokenPolicyManager`] manually and use [`AccountBuilder`]
411    /// directly.
412    pub fn add_existing_basic_faucet(
413        &mut self,
414        auth_method: Auth,
415        token_symbol: &str,
416        max_supply: u64,
417        token_supply: Option<u64>,
418    ) -> anyhow::Result<Account> {
419        let token_supply = token_supply.unwrap_or(0);
420        let name = TokenName::new(token_symbol)?;
421        let symbol = TokenSymbol::new(token_symbol)
422            .with_context(|| format!("invalid token symbol: {token_symbol}"))?;
423        let max_supply = AssetAmount::new(max_supply).context("invalid max_supply")?;
424        let token_supply = AssetAmount::new(token_supply).context("invalid token_supply")?;
425        let faucet = FungibleFaucet::builder()
426            .name(name)
427            .symbol(symbol)
428            .decimals(DEFAULT_FAUCET_DECIMALS)
429            .max_supply(max_supply)
430            .token_supply(token_supply)
431            .build()
432            .context("failed to build FungibleFaucet")?;
433
434        let token_policy_manager = TokenPolicyManager::builder()
435            .active_mint_policy(MintPolicy::allow_all())
436            .active_burn_policy(BurnPolicy::allow_all())
437            .build();
438
439        let account_builder = AccountBuilder::new(self.rng.random())
440            .account_type(AccountType::Public)
441            .with_component(faucet)
442            .with_component(Authority::AuthControlled)
443            .with_asset_callbacks(AssetCallbackFlag::Disabled)
444            .with_components(token_policy_manager)
445            .with_component(Pausable::unpaused())
446            .with_component(PausableManager);
447
448        self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)
449    }
450
451    /// Convenience: builds an owner-controlled (network-style) fungible faucet from a
452    /// token-symbol shorthand using default decimals, the given `mint_policy`, and `BurnAllowAll`.
453    ///
454    /// The faucet is added with [`AccountType::Public`] and [`Auth::IncrNonce`].
455    ///
456    /// `mint_policy` selects the initial active mint policy on the faucet. The installed
457    /// [`TokenPolicyManager`] is always owner-controlled.
458    ///
459    /// The [`MintNote`] and [`BurnNote`] script roots are always added to `allowed_script_roots`,
460    /// so callers only need to provide any additional roots their test scripts require.
461    pub fn add_existing_network_faucet(
462        &mut self,
463        token_symbol: &str,
464        max_supply: u64,
465        owner_account_id: AccountId,
466        token_supply: Option<u64>,
467        mint_policy: MintPolicy,
468        allowed_script_roots: impl IntoIterator<Item = NoteScriptRoot>,
469    ) -> anyhow::Result<Account> {
470        let token_supply = token_supply.unwrap_or(0);
471        let name = TokenName::new(token_symbol)?;
472        let symbol = TokenSymbol::new(token_symbol)
473            .with_context(|| format!("invalid token symbol: {token_symbol}"))?;
474        let max_supply = AssetAmount::new(max_supply).context("invalid max_supply")?;
475        let token_supply = AssetAmount::new(token_supply).context("invalid token_supply")?;
476        let faucet = FungibleFaucet::builder()
477            .name(name)
478            .symbol(symbol)
479            .decimals(DEFAULT_FAUCET_DECIMALS)
480            .max_supply(max_supply)
481            .token_supply(token_supply)
482            .build()
483            .context("failed to build FungibleFaucet")?;
484
485        let token_policy_manager = TokenPolicyManager::builder()
486            .active_mint_policy(mint_policy)
487            .active_burn_policy(BurnPolicy::allow_all())
488            .active_send_policy(TransferPolicy::allow_all())
489            .active_receive_policy(TransferPolicy::allow_all())
490            .build();
491
492        let allowed_script_roots: BTreeSet<NoteScriptRoot> = allowed_script_roots
493            .into_iter()
494            .chain([MintNote::script_root(), BurnNote::script_root()])
495            .collect();
496
497        self.add_existing_network_fungible_faucet(
498            AuthNetworkAccount::with_allowed_notes(allowed_script_roots)?,
499            faucet,
500            AccountType::Public,
501            AccessControl::Ownable2Step { owner: owner_account_id },
502            token_policy_manager,
503        )
504    }
505
506    /// Convenience: adds an existing owner-controlled (network-style) fungible faucet whose token
507    /// metadata is fully provided by the caller. Uses `OwnerOnly` mint policy and `AllowAll`
508    /// burn policy by default.
509    ///
510    /// The [`MintNote`] and [`BurnNote`] script roots are always added to `allowed_script_roots`,
511    /// so callers only need to provide any additional roots their test scripts require.
512    pub fn add_existing_network_faucet_with_metadata(
513        &mut self,
514        owner_account_id: AccountId,
515        faucet: FungibleFaucet,
516        allowed_script_roots: impl IntoIterator<Item = NoteScriptRoot>,
517    ) -> anyhow::Result<Account> {
518        let token_policy_manager = TokenPolicyManager::builder()
519            .active_mint_policy(MintPolicy::owner_only())
520            .active_burn_policy(BurnPolicy::allow_all())
521            .active_send_policy(TransferPolicy::allow_all())
522            .active_receive_policy(TransferPolicy::allow_all())
523            .build();
524
525        let allowed_script_roots: BTreeSet<NoteScriptRoot> = allowed_script_roots
526            .into_iter()
527            .chain([MintNote::script_root(), BurnNote::script_root()])
528            .collect();
529
530        self.add_existing_network_fungible_faucet(
531            AuthNetworkAccount::with_allowed_notes(allowed_script_roots)?,
532            faucet,
533            AccountType::Public,
534            AccessControl::Ownable2Step { owner: owner_account_id },
535            token_policy_manager,
536        )
537    }
538
539    /// Convenience: builds a new (uncreated) basic auth-controlled fungible faucet from a
540    /// token-symbol shorthand using default decimals and `AllowAll` mint/burn policies (no transfer
541    /// policy, so asset callbacks are disabled).
542    pub fn create_new_faucet(
543        &mut self,
544        auth_method: Auth,
545        token_symbol: &str,
546        max_supply: u64,
547    ) -> anyhow::Result<Account> {
548        let name = TokenName::new(token_symbol)?;
549        let symbol = TokenSymbol::new(token_symbol)
550            .with_context(|| format!("invalid token symbol: {token_symbol}"))?;
551        let max_supply = AssetAmount::new(max_supply).context("invalid max_supply")?;
552        let faucet = FungibleFaucet::builder()
553            .name(name)
554            .symbol(symbol)
555            .decimals(DEFAULT_FAUCET_DECIMALS)
556            .max_supply(max_supply)
557            .build()
558            .context("failed to build FungibleFaucet")?;
559
560        let token_policy_manager = TokenPolicyManager::builder()
561            .active_mint_policy(MintPolicy::allow_all())
562            .active_burn_policy(BurnPolicy::allow_all())
563            .build();
564
565        let account_builder = AccountBuilder::new(self.rng.random())
566            .account_type(AccountType::Public)
567            .with_component(faucet)
568            .with_component(Authority::AuthControlled)
569            .with_asset_callbacks(AssetCallbackFlag::Disabled)
570            .with_components(token_policy_manager)
571            .with_component(Pausable::unpaused())
572            .with_component(PausableManager);
573
574        self.add_account_from_builder(auth_method, account_builder, AccountState::New)
575    }
576
577    /// Creates a new public account with an [`MockAccountComponent`] and registers the
578    /// authenticator (if any).
579    pub fn create_new_mock_account(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
580        let account_builder = Account::builder(self.rng.random())
581            .account_type(AccountType::Public)
582            .with_component(MockAccountComponent::with_empty_slots());
583
584        self.add_account_from_builder(auth_method, account_builder, AccountState::New)
585    }
586
587    /// Adds an existing public account with an [`MockAccountComponent`] to the initial chain state
588    /// and registers the authenticator (if any).
589    pub fn add_existing_mock_account(&mut self, auth_method: Auth) -> anyhow::Result<Account> {
590        self.add_existing_mock_account_with_storage_and_assets(auth_method, [], [])
591    }
592
593    /// Adds an existing public account with an [`MockAccountComponent`] to the initial chain state
594    /// and registers the authenticator (if any).
595    pub fn add_existing_mock_account_with_storage(
596        &mut self,
597        auth_method: Auth,
598        slots: impl IntoIterator<Item = StorageSlot>,
599    ) -> anyhow::Result<Account> {
600        self.add_existing_mock_account_with_storage_and_assets(auth_method, slots, [])
601    }
602
603    /// Adds an existing public account with an [`MockAccountComponent`] to the initial chain state
604    /// and registers the authenticator (if any).
605    pub fn add_existing_mock_account_with_assets(
606        &mut self,
607        auth_method: Auth,
608        assets: impl IntoIterator<Item = Asset>,
609    ) -> anyhow::Result<Account> {
610        self.add_existing_mock_account_with_storage_and_assets(auth_method, [], assets)
611    }
612
613    /// Adds an existing public account with an [`MockAccountComponent`] to the initial chain state
614    /// and registers the authenticator (if any).
615    pub fn add_existing_mock_account_with_storage_and_assets(
616        &mut self,
617        auth_method: Auth,
618        slots: impl IntoIterator<Item = StorageSlot>,
619        assets: impl IntoIterator<Item = Asset>,
620    ) -> anyhow::Result<Account> {
621        let account_builder = Account::builder(self.rng.random())
622            .account_type(AccountType::Public)
623            .with_component(MockAccountComponent::with_slots(slots.into_iter().collect()))
624            .with_assets(assets);
625
626        self.add_account_from_builder(auth_method, account_builder, AccountState::Exists)
627    }
628
629    /// Builds the provided [`AccountBuilder`] with the provided auth method and registers the
630    /// authenticator (if any).
631    ///
632    /// - If [`AccountState::Exists`] is given the account is built as an existing account and added
633    ///   to the initial chain state. It can then be used in a transaction without having to
634    ///   validate its seed.
635    /// - If [`AccountState::New`] is given the account is built as a new account and is **not**
636    ///   added to the chain. Its authenticator is registered (if present). Its first transaction
637    ///   will be its creation transaction. [`MockChain::build_tx_context`] can be called with the
638    ///   account to automatically add the authenticator.
639    pub fn add_account_from_builder(
640        &mut self,
641        auth_method: Auth,
642        mut account_builder: AccountBuilder,
643        account_state: AccountState,
644    ) -> anyhow::Result<Account> {
645        let (auth_component, authenticator) = auth_method.build_component();
646        account_builder = account_builder.with_auth_component(auth_component);
647
648        let account = if let AccountState::New = account_state {
649            account_builder.build().context("failed to build account from builder")?
650        } else {
651            account_builder
652                .build_existing()
653                .context("failed to build account from builder")?
654        };
655
656        self.account_authenticators
657            .insert(account.id(), AccountAuthenticator::new(authenticator));
658
659        if let AccountState::Exists = account_state {
660            self.accounts.insert(account.id(), account.clone());
661        }
662
663        Ok(account)
664    }
665    pub fn add_existing_account_from_components(
666        &mut self,
667        auth: Auth,
668        components: impl IntoIterator<Item = AccountComponent>,
669    ) -> anyhow::Result<Account> {
670        let mut account_builder =
671            Account::builder(rand::rng().random()).account_type(AccountType::Public);
672
673        for component in components {
674            account_builder = account_builder.with_component(component);
675        }
676
677        self.add_account_from_builder(auth, account_builder, AccountState::Exists)
678    }
679
680    /// Adds the provided account to the list of genesis accounts.
681    ///
682    /// This method only adds the account and does not store its account authenticator for it.
683    /// Calling [`MockChain::build_tx_context`] on accounts added in this way will not work if
684    /// the account needs an authenticator.
685    ///
686    /// Due to these limitations, prefer using other methods to add accounts to the chain, e.g.
687    /// [`MockChainBuilder::add_account_from_builder`].
688    pub fn add_account(&mut self, account: Account) -> anyhow::Result<()> {
689        self.accounts.insert(account.id(), account);
690
691        // This returns a Result to be conservative in case we need to return an error in the future
692        // and do not want to break this API.
693        Ok(())
694    }
695
696    // NOTE ADD METHODS
697    // ----------------------------------------------------------------------------------------
698
699    /// Adds the provided note to the initial chain state.
700    pub fn add_output_note(&mut self, note: impl Into<RawOutputNote>) {
701        self.notes.push(note.into());
702    }
703
704    /// Creates a new P2ANY note from the provided parameters and adds it to the list of
705    /// genesis notes.
706    ///
707    /// This note is similar to a P2ID note but can be consumed by any account.
708    pub fn add_p2any_note(
709        &mut self,
710        sender_account_id: AccountId,
711        note_type: NoteType,
712        assets: impl IntoIterator<Item = Asset>,
713    ) -> anyhow::Result<Note> {
714        let note = create_p2any_note(sender_account_id, note_type, assets, &mut self.rng);
715        self.add_output_note(RawOutputNote::Full(note.clone()));
716
717        Ok(note)
718    }
719
720    /// Creates a new P2ID note from the provided parameters and adds it to the list of genesis
721    /// notes.
722    ///
723    /// In the created [`MockChain`], the note will be immediately spendable by `target_account_id`
724    /// and carries no additional reclaim or timelock conditions.
725    pub fn add_p2id_note(
726        &mut self,
727        sender_account_id: AccountId,
728        target_account_id: AccountId,
729        asset: &[Asset],
730        note_type: NoteType,
731    ) -> Result<Note, NoteError> {
732        let note: Note = P2idNote::builder()
733            .sender(sender_account_id)
734            .target(target_account_id)
735            .assets(asset.iter().copied())
736            .note_type(note_type)
737            .generate_serial_number(&mut self.rng)
738            .build()?
739            .into();
740        self.add_output_note(RawOutputNote::Full(note.clone()));
741
742        Ok(note)
743    }
744
745    /// Creates a new TX_FEE note from the provided parameters and adds it to the list of genesis
746    /// notes.
747    ///
748    /// In the created [`MockChain`], the note will be immediately spendable by any account.
749    pub fn add_tx_fee_note(
750        &mut self,
751        sender_account_id: AccountId,
752        assets: &[Asset],
753    ) -> Result<Note, NoteError> {
754        let note: Note = TxFeeNote::builder()
755            .sender(sender_account_id)
756            .assets(assets.iter().copied())
757            .generate_serial_number(&mut self.rng)
758            .build()?
759            .into();
760        self.add_output_note(RawOutputNote::Full(note.clone()));
761
762        Ok(note)
763    }
764
765    /// Adds a P2IDE note (pay‑to‑ID‑extended) to the list of genesis notes.
766    ///
767    /// A P2IDE note can include an optional `timelock_height` and/or an optional
768    /// `reclaim_height` after which the note's reclaimer may reclaim the funds.
769    ///
770    /// The `reclaimer` is the account allowed to reclaim the note; when `None` it
771    /// defaults to `sender_account_id`.
772    pub fn add_p2ide_note(
773        &mut self,
774        sender_account_id: AccountId,
775        target_account_id: AccountId,
776        reclaimer: Option<AccountId>,
777        asset: &[Asset],
778        note_type: NoteType,
779        reclaim_height: Option<BlockNumber>,
780        timelock_height: Option<BlockNumber>,
781    ) -> Result<Note, NoteError> {
782        let note: Note = P2ideNote::builder()
783            .sender(sender_account_id)
784            .target(target_account_id)
785            .maybe_reclaimer(reclaimer)
786            .assets(asset.iter().copied())
787            .note_type(note_type)
788            .maybe_reclaim_height(reclaim_height)
789            .maybe_timelock_height(timelock_height)
790            .generate_serial_number(&mut self.rng)
791            .build()?
792            .into();
793
794        self.add_output_note(RawOutputNote::Full(note.clone()));
795
796        Ok(note)
797    }
798
799    /// Adds a public SWAP note to the list of genesis notes.
800    pub fn add_swap_note(
801        &mut self,
802        sender: AccountId,
803        offered_asset: Asset,
804        requested_asset: Asset,
805        payback_note_type: NoteType,
806    ) -> anyhow::Result<(Note, NoteDetails)> {
807        let (swap_note, payback_note) = SwapNote::create(
808            sender,
809            offered_asset,
810            requested_asset,
811            NoteType::Public,
812            NoteAttachments::default(),
813            payback_note_type,
814            &mut self.rng,
815        )?;
816
817        self.add_output_note(RawOutputNote::Full(swap_note.clone()));
818
819        Ok((swap_note, payback_note))
820    }
821
822    /// Adds a public `SPAWN` note to the list of genesis notes.
823    ///
824    /// A `SPAWN` note contains a note script that creates all `output_notes` that get passed as a
825    /// parameter.
826    ///
827    /// # Errors
828    ///
829    /// Returns an error if:
830    /// - the sender account ID of the provided output notes is not consistent or does not match the
831    ///   transaction's sender.
832    pub fn add_spawn_note<'note, I>(
833        &mut self,
834        output_notes: impl IntoIterator<Item = &'note Note, IntoIter = I>,
835    ) -> anyhow::Result<Note>
836    where
837        I: ExactSizeIterator<Item = &'note Note>,
838    {
839        let note = create_spawn_note(output_notes)?;
840        self.add_output_note(RawOutputNote::Full(note.clone()));
841
842        Ok(note)
843    }
844
845    /// Creates a new P2ID note with the provided amount of the fee asset of the chain.
846    ///
847    /// The fee faucet ID of the asset can be set using [`Self::fee_faucet_id`]. By default it
848    /// is [`ACCOUNT_ID_FEE_FAUCET`].
849    ///
850    /// In the created [`MockChain`], the note will be immediately spendable by `target_account_id`.
851    pub fn add_p2id_note_with_fee(
852        &mut self,
853        target_account_id: AccountId,
854        amount: u64,
855    ) -> anyhow::Result<Note> {
856        let fee_asset = self.fee_asset(amount)?;
857        let note = self.add_p2id_note(
858            self.fee_faucet_id,
859            target_account_id,
860            &[Asset::from(fee_asset)],
861            NoteType::Public,
862        )?;
863
864        Ok(note)
865    }
866
867    // HELPER FUNCTIONS
868    // ----------------------------------------------------------------------------------------
869
870    /// Returns a mutable reference to the builder's RNG.
871    ///
872    /// This can be used when creating accounts or notes and randomness is required.
873    pub fn rng_mut(&mut self) -> &mut RandomCoin {
874        &mut self.rng
875    }
876
877    /// Constructs a fungible asset based on the fee faucet ID and the provided amount.
878    fn fee_asset(&self, amount: u64) -> anyhow::Result<FungibleAsset> {
879        FungibleAsset::new(self.fee_faucet_id, amount).context("failed to create fee asset")
880    }
881}
882
883impl Default for MockChainBuilder {
884    fn default() -> Self {
885        Self::new()
886    }
887}