Skip to main content

miden_client/test_utils/
mock.rs

1use alloc::boxed::Box;
2use alloc::collections::{BTreeMap, BTreeSet};
3use alloc::sync::Arc;
4use alloc::vec::Vec;
5use core::sync::atomic::{AtomicUsize, Ordering};
6
7use miden_protocol::Word;
8use miden_protocol::account::{
9    AccountId,
10    AccountUpdateDetails,
11    AccountVaultPatch,
12    StorageMapKey,
13    StorageMapPatchEntries,
14    StorageSlot,
15    StorageSlotContent,
16    StorageSlotName,
17    StorageSlotType,
18};
19use miden_protocol::address::NetworkId;
20use miden_protocol::batch::{ProposedBatch, ProvenBatch};
21use miden_protocol::block::{BlockHeader, BlockNumber, ProvenBlock};
22use miden_protocol::crypto::merkle::MerklePath;
23use miden_protocol::crypto::merkle::mmr::{Forest, Mmr, MmrProof};
24use miden_protocol::crypto::merkle::smt::PartialSmt;
25use miden_protocol::note::{NoteAttachments, NoteHeader, NoteId, NoteScript, NoteTag};
26use miden_protocol::transaction::{OutputNote, ProvenTransaction};
27use miden_testing::{MockChain, MockChainNote};
28use miden_tx::utils::sync::RwLock;
29
30use crate::Client;
31use crate::rpc::domain::account::{
32    AccountDetails,
33    AccountProof,
34    AccountStorageDetails,
35    AccountStorageMapDetails,
36    AccountVaultDetails,
37    GetAccountRequest,
38    StorageMapEntries,
39    StorageMapEntry,
40    StorageMapFetch,
41    VaultFetch,
42};
43use crate::rpc::domain::account_vault::AccountVaultInfo;
44use crate::rpc::domain::note::{CommittedNote, FetchedNote, SyncNotesBlock};
45use crate::rpc::domain::nullifier::NullifierUpdate;
46use crate::rpc::domain::status::NetworkNoteStatusInfo;
47use crate::rpc::domain::storage_map::StorageMapInfo;
48use crate::rpc::domain::sync::{ChainMmrInfo, SyncTarget};
49use crate::rpc::domain::transaction::TransactionRecord;
50use crate::rpc::encryption::{AttestedTransactionEncryptionKey, SealedTransactionInputs};
51use crate::rpc::{AccountStateAt, NodeRpcClient, RpcEndpoint, RpcError, RpcStatusInfo};
52
53pub type MockClient<AUTH> = Client<AUTH>;
54
55/// Mock RPC API
56///
57/// This struct implements the RPC API used by the client to communicate with the node. It simulates
58/// most of the functionality of the actual node, with some small differences:
59/// - It uses a [`MockChain`] to simulate the blockchain state.
60/// - Blocks are not automatically created after time passes, but rather new blocks are created when
61///   calling the `prove_block` method.
62/// - Network account and transactions aren't supported in the current version.
63/// - Account update block numbers aren't tracked, so any endpoint that returns when certain account
64///   updates were made will return the chain tip block number instead.
65#[derive(Clone)]
66pub struct MockRpcApi {
67    account_commitment_updates: Arc<RwLock<BTreeMap<BlockNumber, BTreeMap<AccountId, Word>>>>,
68    pub mock_chain: Arc<RwLock<MockChain>>,
69    /// Chain snapshots used to answer block-pinned account queries.
70    historical_chains: Arc<RwLock<BTreeMap<BlockNumber, Arc<MockChain>>>>,
71    oversize_threshold: usize,
72    /// Note headers to report as erased in sync transaction responses.
73    erased_notes: Arc<RwLock<Vec<NoteHeader>>>,
74    /// Attachment content `get_notes_by_id` serves for private notes, populated by
75    /// `submit_proven_transaction` and by `register_private_note_attachments`. A note absent here
76    /// is served with empty attachments, which is how a test simulates a withholding node.
77    private_note_attachments: Arc<RwLock<BTreeMap<NoteId, NoteAttachments>>>,
78    /// Test overrides for the MMR paths returned by `sync_notes`, keyed by block number.
79    sync_notes_mmr_path_overrides: Arc<RwLock<BTreeMap<BlockNumber, MerklePath>>>,
80    /// Number of `get_notes_by_id` requests served, so a test can assert that a flow avoided the
81    /// round trip.
82    get_notes_by_id_calls: Arc<AtomicUsize>,
83    /// Failures to serve instead of answering, keyed by [`RpcEndpoint::proto_name`] and set by
84    /// [`MockRpcApi::fail_next_call`]. An entry is removed when served, so the call after it
85    /// answers normally and a test can exercise a retry.
86    next_call_failures: Arc<RwLock<BTreeMap<&'static str, RpcError>>>,
87}
88
89impl Default for MockRpcApi {
90    fn default() -> Self {
91        Self::new(MockChain::new())
92    }
93}
94
95impl MockRpcApi {
96    // Constant to use in mocked pagination.
97    const PAGINATION_BLOCK_LIMIT: u32 = 5;
98
99    /// Creates a new [`MockRpcApi`] instance with the state of the provided [`MockChain`].
100    pub fn new(mock_chain: MockChain) -> Self {
101        Self {
102            account_commitment_updates: Arc::new(RwLock::new(build_account_updates(&mock_chain))),
103            mock_chain: Arc::new(RwLock::new(mock_chain)),
104            historical_chains: Arc::new(RwLock::new(BTreeMap::new())),
105            oversize_threshold: 1000,
106            erased_notes: Arc::new(RwLock::new(Vec::new())),
107            private_note_attachments: Arc::new(RwLock::new(BTreeMap::new())),
108            sync_notes_mmr_path_overrides: Arc::new(RwLock::new(BTreeMap::new())),
109            get_notes_by_id_calls: Arc::new(AtomicUsize::new(0)),
110            next_call_failures: Arc::new(RwLock::new(BTreeMap::new())),
111        }
112    }
113
114    /// Makes the next call to `endpoint` fail with `error` instead of answering. The failure is
115    /// consumed, so the call after it answers normally and a test can exercise a retry.
116    ///
117    /// Staging a failure for an endpoint whose mock implementation does not look for one is a
118    /// silent no-op.
119    pub fn fail_next_call(&self, endpoint: RpcEndpoint, error: RpcError) {
120        self.next_call_failures.write().insert(endpoint.proto_name(), error);
121    }
122
123    /// Returns the failure staged for `endpoint`, removing it so it is served once.
124    fn take_failure(&self, endpoint: RpcEndpoint) -> Option<RpcError> {
125        self.next_call_failures.write().remove(endpoint.proto_name())
126    }
127
128    /// Registers the attachment content for a private note so that subsequent `get_notes_by_id`
129    /// responses include it, mirroring a node that stores private-note attachments on-chain.
130    pub fn register_private_note_attachments(&self, note_id: NoteId, attachments: NoteAttachments) {
131        self.private_note_attachments.write().insert(note_id, attachments);
132    }
133
134    /// Returns how many `get_notes_by_id` requests this API has served.
135    pub fn get_notes_by_id_call_count(&self) -> usize {
136        self.get_notes_by_id_calls.load(Ordering::Relaxed)
137    }
138
139    /// Overrides the MMR path returned by `sync_notes` for the specified block.
140    pub fn set_sync_notes_mmr_path(&self, block_num: BlockNumber, path: MerklePath) {
141        self.sync_notes_mmr_path_overrides.write().insert(block_num, path);
142    }
143
144    /// Sets the oversize threshold for `get_account`. A storage map whose entries were requested
145    /// in full comes back as `StorageMapEntries::LimitExceeded` past this threshold, and a vault
146    /// with more assets than it comes back with the `too_many_assets` flag set.
147    #[must_use]
148    pub fn with_oversize_threshold(mut self, threshold: usize) -> Self {
149        self.oversize_threshold = threshold;
150        self
151    }
152
153    /// Registers a note header to be reported as erased in subsequent sync transaction responses.
154    pub fn mark_note_as_erased(&self, header: NoteHeader) {
155        self.erased_notes.write().push(header);
156    }
157
158    /// Returns the current MMR of the blockchain.
159    pub fn get_mmr(&self) -> Mmr {
160        self.mock_chain.read().blockchain().as_mmr().clone()
161    }
162
163    /// Returns the chain tip block number.
164    pub fn get_chain_tip_block_num(&self) -> BlockNumber {
165        self.mock_chain.read().latest_block_header().block_num()
166    }
167
168    /// Advances the mock chain by proving the next block, committing all pending objects to the
169    /// chain in the process.
170    pub fn prove_block(&self) {
171        let proven_block = {
172            let mut mock_chain = self.mock_chain.write();
173            let historical_block_num = mock_chain.latest_block_header().block_num();
174            let snapshot = Arc::new(mock_chain.clone());
175            let proven_block = mock_chain.prove_next_block().unwrap();
176            self.historical_chains.write().insert(historical_block_num, snapshot);
177            proven_block
178        };
179        let block_num = proven_block.header().block_num();
180        let mut account_commitment_updates = self.account_commitment_updates.write();
181        let updates: BTreeMap<AccountId, Word> = proven_block
182            .body()
183            .updated_accounts()
184            .iter()
185            .map(|update| (update.account_id(), update.final_state_commitment()))
186            .collect();
187
188        if !updates.is_empty() {
189            account_commitment_updates.insert(block_num, updates);
190        }
191    }
192
193    /// Retrieves a block by its block number.
194    fn get_block_by_num(&self, block_num: BlockNumber) -> BlockHeader {
195        self.mock_chain.read().block_header(block_num.as_usize())
196    }
197
198    /// Retrieves account vault updates in a given block range.
199    /// This method tries to simulate pagination by limiting the number of blocks processed per
200    /// request.
201    fn get_sync_account_vault_request(
202        &self,
203        block_from: BlockNumber,
204        block_to: BlockNumber,
205        account_id: AccountId,
206    ) -> (BlockNumber, BlockNumber, AccountVaultPatch) {
207        let chain_tip = self.get_chain_tip_block_num();
208        let target_block = block_to.min(chain_tip);
209
210        let page_end_block: BlockNumber = (block_from.as_u32() + Self::PAGINATION_BLOCK_LIMIT)
211            .min(target_block.as_u32())
212            .into();
213
214        // Blocks are iterated in ascending order, so later blocks win per asset ID.
215        let mut vault_patch = AccountVaultPatch::default();
216        for block in self.mock_chain.read().proven_blocks() {
217            let block_number = block.header().block_num();
218            // Only include blocks in range [block_from, page_end_block]
219            if block_number < block_from || block_number > page_end_block {
220                continue;
221            }
222
223            for update in block
224                .body()
225                .updated_accounts()
226                .iter()
227                .filter(|block_acc_update| block_acc_update.account_id() == account_id)
228            {
229                let AccountUpdateDetails::Public(patch) = update.details().clone() else {
230                    continue;
231                };
232
233                vault_patch.merge(patch.vault().clone());
234            }
235        }
236
237        (chain_tip, page_end_block, vault_patch)
238    }
239
240    /// Retrieves transactions in a given block range that match the provided account IDs
241    fn get_sync_transactions_request(
242        &self,
243        block_from: BlockNumber,
244        block_to: BlockNumber,
245        account_ids: &[AccountId],
246    ) -> Vec<TransactionRecord> {
247        let mut transactions = Vec::new();
248        for block in self.mock_chain.read().proven_blocks() {
249            let block_number = block.header().block_num();
250            if block_number < block_from || block_number > block_to {
251                continue;
252            }
253
254            for transaction_header in block.body().transactions().as_slice() {
255                if !account_ids.contains(&transaction_header.account_id()) {
256                    continue;
257                }
258
259                let erased_output_notes = self.erased_notes.read().clone();
260
261                transactions.push(TransactionRecord {
262                    block_num: block_number,
263                    transaction_header: transaction_header.clone(),
264                    output_notes: vec![],
265                    erased_output_notes,
266                    consumed_note_refs: vec![],
267                });
268            }
269        }
270
271        transactions
272    }
273
274    /// Retrieves storage map updates in a given block range.
275    ///
276    /// This method tries to simulate pagination of the real node.
277    fn get_sync_storage_maps_request(
278        &self,
279        block_from: BlockNumber,
280        block_to: BlockNumber,
281        account_id: AccountId,
282    ) -> (BlockNumber, BlockNumber, BTreeMap<StorageSlotName, StorageMapPatchEntries>) {
283        let chain_tip = self.get_chain_tip_block_num();
284        let target_block = block_to.min(chain_tip);
285
286        let page_end_block: BlockNumber = (block_from.as_u32() + Self::PAGINATION_BLOCK_LIMIT)
287            .min(target_block.as_u32())
288            .into();
289
290        // Blocks are iterated in ascending order, so later blocks win per `(slot, key)`.
291        let mut map_entries: BTreeMap<StorageSlotName, StorageMapPatchEntries> = BTreeMap::new();
292        for block in self.mock_chain.read().proven_blocks() {
293            let block_number = block.header().block_num();
294            // Only include blocks in range [block_from, page_end_block]
295            if block_number < block_from || block_number > page_end_block {
296                continue;
297            }
298
299            for update in block
300                .body()
301                .updated_accounts()
302                .iter()
303                .filter(|block_acc_update| block_acc_update.account_id() == account_id)
304            {
305                let AccountUpdateDetails::Public(patch) = update.details().clone() else {
306                    continue;
307                };
308
309                for (slot_name, map_patch) in patch.storage().maps() {
310                    if let Some(entries) = map_patch.entries() {
311                        map_entries
312                            .entry(slot_name.clone())
313                            .or_default()
314                            .as_map_mut()
315                            .extend(entries.as_map().clone());
316                    }
317                }
318            }
319        }
320
321        (chain_tip, page_end_block, map_entries)
322    }
323
324    pub fn get_available_notes(&self) -> Vec<MockChainNote> {
325        self.mock_chain.read().committed_notes().values().cloned().collect()
326    }
327
328    pub fn get_public_available_notes(&self) -> Vec<MockChainNote> {
329        self.mock_chain
330            .read()
331            .committed_notes()
332            .values()
333            .filter(|n| matches!(n, MockChainNote::Public(_, _)))
334            .cloned()
335            .collect()
336    }
337
338    pub fn get_private_available_notes(&self) -> Vec<MockChainNote> {
339        self.mock_chain
340            .read()
341            .committed_notes()
342            .values()
343            .filter(|n| matches!(n, MockChainNote::Private(_, _, _, _)))
344            .cloned()
345            .collect()
346    }
347
348    pub fn advance_blocks(&self, num_blocks: u32) {
349        let mut mock_chain = self.mock_chain.write();
350        let block_num = mock_chain.latest_block_header().block_num();
351        let snapshot = Arc::new(mock_chain.clone());
352        mock_chain.prove_until_block(block_num + num_blocks).unwrap();
353        self.historical_chains.write().insert(block_num, snapshot);
354    }
355}
356#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
357#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
358impl NodeRpcClient for MockRpcApi {
359    /// Always reports the commitment as unset, unlike a real client.
360    ///
361    /// A real client's RPC connection is its own, so whoever set the commitment also stored the
362    /// header. Tests share one mock across clients with separate stores, where a commitment set by
363    /// the first would stop every later client from storing genesis at all.
364    fn has_genesis_commitment(&self) -> Option<Word> {
365        None
366    }
367
368    async fn set_genesis_commitment(&self, _commitment: Word) -> Result<(), RpcError> {
369        // The mock sends no request headers, so there is nothing to pin the commitment to.
370        Ok(())
371    }
372
373    /// Returns note updates in the inclusive block range `[block_from, block_to]`.
374    /// Only notes that match the provided tags will be returned, grouped by block.
375    async fn sync_notes(
376        &self,
377        block_from: BlockNumber,
378        block_to: BlockNumber,
379        note_tags: &BTreeSet<NoteTag>,
380    ) -> Result<Vec<SyncNotesBlock>, RpcError> {
381        let mut blocks_with_notes: BTreeMap<BlockNumber, BTreeMap<NoteId, CommittedNote>> =
382            BTreeMap::new();
383        for note in self.mock_chain.read().committed_notes().values() {
384            let note_block = note.inclusion_proof().location().block_num();
385            if note_tags.contains(&note.metadata().tag())
386                && note_block >= block_from
387                && note_block <= block_to
388            {
389                let mut committed =
390                    CommittedNote::new(note.id(), *note.metadata(), note.inclusion_proof().clone());
391                // Mirror the node: a single-word attachment is sent verbatim and the record is
392                // complete. A larger one is sent as a commitment only.
393                let attachments = note.attachments();
394                if attachments.iter().all(|attachment| attachment.num_words() == 1) {
395                    committed = committed
396                        .with_attachments(attachments.clone())
397                        .expect("the note's own attachments match its commitment");
398                }
399                blocks_with_notes.entry(note_block).or_default().insert(note.id(), committed);
400            }
401        }
402
403        Ok(blocks_with_notes
404            .into_iter()
405            .map(|(bn, notes)| {
406                let block_header = self.get_block_by_num(bn);
407                let mmr_path =
408                    self.sync_notes_mmr_path_overrides.read().get(&bn).cloned().unwrap_or_else(
409                        || self.get_mmr().open(bn.as_usize()).unwrap().merkle_path().clone(),
410                    );
411                SyncNotesBlock { block_header, mmr_path, notes }
412            })
413            .collect())
414    }
415
416    async fn sync_chain_mmr(
417        &self,
418        current_block_height: BlockNumber,
419        upper_bound: SyncTarget,
420    ) -> Result<ChainMmrInfo, RpcError> {
421        let chain_tip = self.get_chain_tip_block_num();
422        // The mock chain doesn't distinguish committed vs proven tips.
423        let target_block = match upper_bound {
424            SyncTarget::CommittedChainTip | SyncTarget::ProvenChainTip => chain_tip,
425        };
426
427        let from_forest = if current_block_height == target_block {
428            target_block.as_usize()
429        } else {
430            current_block_height.as_u32() as usize + 1
431        };
432
433        let mmr_delta = self
434            .get_mmr()
435            .get_delta(
436                Forest::new(from_forest).unwrap(),
437                Forest::new(target_block.as_usize()).unwrap(),
438            )
439            .unwrap();
440
441        let block_header = self.get_block_by_num(target_block);
442
443        Ok(ChainMmrInfo {
444            block_from: current_block_height,
445            block_to: target_block,
446            mmr_delta,
447            block_header,
448        })
449    }
450
451    /// Retrieves the block header for the specified block number. If the block number is not
452    /// provided, the chain tip block header will be returned.
453    async fn get_block_header_by_number(
454        &self,
455        block_num: Option<BlockNumber>,
456        include_mmr_proof: bool,
457    ) -> Result<(BlockHeader, Option<MmrProof>), RpcError> {
458        let block = if let Some(block_num) = block_num {
459            self.mock_chain.read().block_header(block_num.as_usize())
460        } else {
461            self.mock_chain.read().latest_block_header()
462        };
463
464        let mmr_proof = if include_mmr_proof {
465            Some(self.get_mmr().open(block_num.unwrap().as_usize()).unwrap())
466        } else {
467            None
468        };
469
470        Ok((block, mmr_proof))
471    }
472
473    /// Returns the node's tracked notes that match the provided note IDs.
474    async fn get_notes_by_id(&self, note_ids: &[NoteId]) -> Result<Vec<FetchedNote>, RpcError> {
475        self.get_notes_by_id_calls.fetch_add(1, Ordering::Relaxed);
476
477        // assume all public notes for now
478        let notes = self.mock_chain.read().committed_notes().clone();
479
480        let hit_notes = note_ids.iter().filter_map(|id| notes.get(id));
481        let mut return_notes = vec![];
482        for note in hit_notes {
483            let fetched_note = match note {
484                MockChainNote::Private(note_id, note_metadata, _, note_inclusion_proof) => {
485                    let attachments = self
486                        .private_note_attachments
487                        .read()
488                        .get(note_id)
489                        .cloned()
490                        .unwrap_or_else(NoteAttachments::empty);
491                    FetchedNote::Private(
492                        *note_id,
493                        *note_metadata,
494                        attachments,
495                        note_inclusion_proof.clone(),
496                    )
497                },
498                MockChainNote::Public(note, note_inclusion_proof) => {
499                    FetchedNote::Public(note.clone(), note_inclusion_proof.clone())
500                },
501            };
502            return_notes.push(fetched_note);
503        }
504        Ok(return_notes)
505    }
506
507    /// The mock does not serve the encryption key. Verifying an attestation needs a validator
508    /// signature the mock chain cannot produce, so tests that submit transactions seed the key
509    /// directly through `Client::seed_transaction_encryption_key` instead.
510    async fn get_transaction_encryption_key(
511        &self,
512    ) -> Result<AttestedTransactionEncryptionKey, RpcError> {
513        Err(RpcError::TransactionEncryptionKeyRejected(
514            "the mock RPC client does not serve a transaction encryption key".into(),
515        ))
516    }
517
518    /// Simulates the submission of a proven transaction to the node. This will create a new block
519    /// just for the new transaction and return the block number of the newly created block.
520    async fn submit_proven_transaction(
521        &self,
522        proven_transaction: ProvenTransaction,
523        _sealed_transaction_inputs: SealedTransactionInputs, /* Unnecessary for testing client
524                                                              * itself. */
525    ) -> Result<BlockNumber, RpcError> {
526        if let Some(error) = self.take_failure(RpcEndpoint::SubmitProvenTx) {
527            return Err(error);
528        }
529
530        // Record private-note attachment content the way a real node does: attachments are
531        // stored on-chain even for private notes, so `get_notes_by_id` must be able to serve
532        // them. The mock chain itself only keeps private note headers.
533        for note in proven_transaction.output_notes().iter() {
534            if let OutputNote::Private(private_note) = note
535                && !private_note.attachments().is_empty()
536            {
537                self.private_note_attachments
538                    .write()
539                    .insert(private_note.id(), private_note.attachments().clone());
540            }
541        }
542
543        {
544            let mut mock_chain = self.mock_chain.write();
545            mock_chain.add_pending_proven_transaction(proven_transaction.clone());
546        };
547
548        let block_num = self.get_chain_tip_block_num();
549
550        Ok(block_num)
551    }
552
553    /// Simulates the submission of a proven batch to the node by adding it to the mock chain's
554    /// pending batches. The `proposed_batch` and `sealed_transaction_inputs` arguments are accepted
555    /// to match the trait signature but are unused — the mock relies on the `ProvenBatch`
556    /// alone, matching how `submit_proven_transaction` ignores its `sealed_transaction_inputs`.
557    async fn submit_proven_batch(
558        &self,
559        proven_batch: ProvenBatch,
560        _proposed_batch: ProposedBatch,
561        _sealed_transaction_inputs: Vec<SealedTransactionInputs>,
562    ) -> Result<BlockNumber, RpcError> {
563        let mut mock_chain = self.mock_chain.write();
564        mock_chain.add_pending_batch(proven_batch);
565        drop(mock_chain);
566
567        let block_num = self.get_chain_tip_block_num();
568
569        Ok(block_num)
570    }
571
572    /// Returns the account proof for the specified account. The `known_code` and `vault` fields
573    /// are ignored: full account data is returned, with truncation flags set when it exceeds
574    /// `oversize_threshold`.
575    async fn get_account(
576        &self,
577        account_id: AccountId,
578        request: GetAccountRequest,
579    ) -> Result<(BlockNumber, AccountProof), RpcError> {
580        let current_chain = self.mock_chain.read();
581        let current_block_number = current_chain.latest_block_header().block_num();
582        let block_number = match request.at {
583            AccountStateAt::Block(number) => number,
584            AccountStateAt::ChainTip => current_block_number,
585        };
586        let historical_chain = match request.at {
587            AccountStateAt::Block(_) if block_number != current_block_number => Some(
588                self.historical_chains.read().get(&block_number).cloned().ok_or_else(|| {
589                    RpcError::InvalidResponse(alloc::format!(
590                        "no mock chain snapshot at block {block_number}"
591                    ))
592                })?,
593            ),
594            AccountStateAt::ChainTip | AccountStateAt::Block(_) => None,
595        };
596        let mock_chain = historical_chain.as_deref().unwrap_or(&*current_chain);
597
598        let headers = if account_id.is_public() {
599            let account = mock_chain.committed_account(account_id).unwrap();
600
601            // `All` enumerates the account's map slots directly — the mock can introspect the
602            // account, so it simulates the (not-yet-on-the-wire) "all storage maps" request.
603            // A slot maps to the keys requested for it, empty meaning "every entry".
604            let requested_slots: Vec<(StorageSlotName, Vec<StorageMapKey>)> = match &request.storage
605            {
606                StorageMapFetch::Skip => Vec::new(),
607                StorageMapFetch::Slots(reqs) => {
608                    reqs.inner().iter().map(|(name, keys)| (name.clone(), keys.clone())).collect()
609                },
610                StorageMapFetch::All => account
611                    .storage()
612                    .to_header()
613                    .slots()
614                    .filter(|slot| slot.slot_type() == StorageSlotType::Map)
615                    .map(|slot| (slot.name().clone(), Vec::new()))
616                    .collect(),
617            };
618
619            let mut map_details = vec![];
620            for (slot_name, requested_keys) in &requested_slots {
621                if let Some(StorageSlotContent::Map(storage_map)) =
622                    account.storage().get(slot_name).map(StorageSlot::content)
623                {
624                    // Mirror the node: named keys come back as one partial SMT covering them,
625                    // and an empty key list comes back as the whole map, or as `LimitExceeded`
626                    // once it grows past the threshold.
627                    let entries = if requested_keys.is_empty() {
628                        let entries: Vec<StorageMapEntry> = storage_map
629                            .entries()
630                            .map(|(key, value)| StorageMapEntry { key: *key, value: *value })
631                            .collect();
632
633                        if entries.len() > self.oversize_threshold {
634                            StorageMapEntries::LimitExceeded
635                        } else {
636                            StorageMapEntries::AllEntries(entries)
637                        }
638                    } else {
639                        let partial_smt = PartialSmt::from_proofs(
640                            requested_keys.iter().map(|key| storage_map.open(key).into()),
641                        )
642                        .expect("proofs from one map share a root");
643
644                        StorageMapEntries::PartialMap {
645                            map_keys: requested_keys.clone(),
646                            partial_smt,
647                        }
648                    };
649
650                    map_details
651                        .push(AccountStorageMapDetails { slot_name: slot_name.clone(), entries });
652                } else {
653                    panic!("Storage slot {slot_name} is not a map");
654                }
655            }
656
657            let storage_details = AccountStorageDetails {
658                header: account.storage().to_header(),
659                map_details,
660            };
661
662            // Mirror the node: `Skip` sends no assets, and `IfChangedFrom` omits them when the
663            // account's vault root already equals the sent commitment.
664            let include_assets = match request.vault {
665                VaultFetch::Skip => false,
666                VaultFetch::Always => true,
667                VaultFetch::IfChangedFrom(root) => root != account.vault().root(),
668            };
669            let mut assets = vec![];
670            if include_assets {
671                for asset in account.vault().assets() {
672                    assets.push(asset);
673                }
674            }
675            let vault_details = AccountVaultDetails {
676                too_many_assets: assets.len() > self.oversize_threshold,
677                assets,
678            };
679
680            Some(AccountDetails {
681                header: account.into(),
682                storage_details,
683                code: account.code().clone(),
684                vault_details,
685            })
686        } else {
687            None
688        };
689
690        let witness = mock_chain.account_tree().open(account_id);
691
692        let proof = AccountProof::new(witness, headers).unwrap();
693
694        Ok((block_number, proof))
695    }
696
697    /// Returns the nullifiers created after the specified block number that match the provided
698    /// prefixes.
699    async fn sync_nullifiers(
700        &self,
701        prefixes: &[u16],
702        block_from: BlockNumber,
703        block_to: BlockNumber,
704    ) -> Result<Vec<NullifierUpdate>, RpcError> {
705        let nullifiers = self
706            .mock_chain
707            .read()
708            .nullifier_tree()
709            .entries()
710            .filter_map(|(nullifier, block_num)| {
711                let within_range = block_num >= block_from && block_num <= block_to;
712
713                if prefixes.contains(&nullifier.prefix()) && within_range {
714                    Some(NullifierUpdate { nullifier, block_num })
715                } else {
716                    None
717                }
718            })
719            .collect::<Vec<_>>();
720
721        Ok(nullifiers)
722    }
723
724    async fn get_block_by_number(
725        &self,
726        block_num: BlockNumber,
727        _include_proof: bool,
728    ) -> Result<ProvenBlock, RpcError> {
729        let block = self
730            .mock_chain
731            .read()
732            .proven_blocks()
733            .iter()
734            .find(|b| b.header().block_num() == block_num)
735            .unwrap()
736            .clone();
737
738        Ok(block)
739    }
740
741    async fn get_note_script_by_root(&self, root: Word) -> Result<Option<NoteScript>, RpcError> {
742        let script = self
743            .get_available_notes()
744            .iter()
745            .filter_map(|note| note.note())
746            .find(|n| Word::from(n.script().root()) == root)
747            .map(|n| n.script().clone());
748
749        Ok(script)
750    }
751
752    async fn sync_storage_maps(
753        &self,
754        block_from: BlockNumber,
755        block_to: BlockNumber,
756        account_id: AccountId,
757    ) -> Result<StorageMapInfo, RpcError> {
758        let mut map_entries: BTreeMap<StorageSlotName, StorageMapPatchEntries> = BTreeMap::new();
759        let mut current_block_from = block_from;
760        let chain_tip = self.get_chain_tip_block_num();
761        let target_block = block_to.min(chain_tip);
762
763        loop {
764            let (page_chain_tip, page_block_number, page_entries) =
765                self.get_sync_storage_maps_request(current_block_from, block_to, account_id);
766            for (slot_name, entries) in page_entries {
767                map_entries
768                    .entry(slot_name)
769                    .or_default()
770                    .as_map_mut()
771                    .extend(entries.into_map());
772            }
773
774            if page_block_number >= target_block {
775                return Ok(StorageMapInfo {
776                    chain_tip: page_chain_tip,
777                    block_number: page_block_number,
778                    map_entries,
779                });
780            }
781
782            current_block_from = (page_block_number.as_u32() + 1).into();
783        }
784    }
785
786    async fn sync_account_vault(
787        &self,
788        block_from: BlockNumber,
789        block_to: BlockNumber,
790        account_id: AccountId,
791    ) -> Result<AccountVaultInfo, RpcError> {
792        let mut vault_patch = AccountVaultPatch::default();
793        let mut current_block_from = block_from;
794        let chain_tip = self.get_chain_tip_block_num();
795        let target_block = block_to.min(chain_tip);
796
797        loop {
798            let (page_chain_tip, page_block_number, page_patch) =
799                self.get_sync_account_vault_request(current_block_from, block_to, account_id);
800            vault_patch.merge(page_patch);
801
802            if page_block_number >= target_block {
803                return Ok(AccountVaultInfo {
804                    chain_tip: page_chain_tip,
805                    block_number: page_block_number,
806                    vault_patch,
807                });
808            }
809
810            current_block_from = (page_block_number.as_u32() + 1).into();
811        }
812    }
813
814    async fn sync_transactions(
815        &self,
816        block_from: BlockNumber,
817        block_to: BlockNumber,
818        account_ids: Vec<AccountId>,
819    ) -> Result<Vec<TransactionRecord>, RpcError> {
820        Ok(self.get_sync_transactions_request(block_from, block_to, &account_ids))
821    }
822
823    async fn get_network_id(&self) -> Result<NetworkId, RpcError> {
824        Ok(NetworkId::Testnet)
825    }
826
827    async fn get_rpc_limits(&self) -> Result<crate::rpc::RpcLimits, RpcError> {
828        Ok(crate::rpc::RpcLimits::default())
829    }
830
831    fn has_rpc_limits(&self) -> Option<crate::rpc::RpcLimits> {
832        None
833    }
834
835    async fn set_rpc_limits(&self, _limits: crate::rpc::RpcLimits) {
836        // No-op for mock client
837    }
838
839    async fn get_status_unversioned(&self) -> Result<RpcStatusInfo, RpcError> {
840        Ok(RpcStatusInfo {
841            version: env!("CARGO_PKG_VERSION").into(),
842            genesis_commitment: None,
843            chain_tip: 0,
844            block_producer: None,
845        })
846    }
847
848    async fn get_network_note_status(
849        &self,
850        _note_id: NoteId,
851    ) -> Result<NetworkNoteStatusInfo, RpcError> {
852        todo!("We need to check if we want to implement this for the mockchain");
853    }
854}
855
856// CONVERSIONS
857// ================================================================================================
858
859impl From<MockChain> for MockRpcApi {
860    fn from(mock_chain: MockChain) -> Self {
861        MockRpcApi::new(mock_chain)
862    }
863}
864
865// HELPERS
866// ================================================================================================
867
868fn build_account_updates(
869    mock_chain: &MockChain,
870) -> BTreeMap<BlockNumber, BTreeMap<AccountId, Word>> {
871    let mut account_commitment_updates = BTreeMap::new();
872    for block in mock_chain.proven_blocks() {
873        let block_num = block.header().block_num();
874        let mut updates = BTreeMap::new();
875
876        for update in block.body().updated_accounts() {
877            updates.insert(update.account_id(), update.final_state_commitment());
878        }
879
880        if updates.is_empty() {
881            continue;
882        }
883
884        account_commitment_updates.insert(block_num, updates);
885    }
886    account_commitment_updates
887}