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