miden_client/test_utils/
note_transport.rs

1use alloc::boxed::Box;
2use alloc::collections::BTreeMap;
3use alloc::string::ToString;
4use alloc::sync::Arc;
5use alloc::vec::Vec;
6use core::sync::atomic::{AtomicUsize, Ordering};
7
8use miden_protocol::block::BlockNumber;
9use miden_protocol::note::{NoteHeader, NoteId, NoteInclusionProof, NoteTag};
10use miden_tx::utils::serde::{
11    ByteReader,
12    ByteWriter,
13    Deserializable,
14    DeserializationError,
15    Serializable,
16};
17use miden_tx::utils::sync::RwLock;
18
19use crate::note_transport::{
20    NoteInfo,
21    NoteTransportClient,
22    NoteTransportCursor,
23    NoteTransportError,
24    NoteTransportPage,
25    TransportNote,
26};
27
28/// Mock Note Transport Node
29///
30/// Simulates the functionality of the note transport node.
31#[derive(Clone)]
32pub struct MockNoteTransportNode {
33    notes: BTreeMap<NoteTag, Vec<(NoteInfo, NoteTransportCursor)>>,
34    nonce: u64,
35    next_sequence: u64,
36    /// Optional per-response batch cap; if `Some(n)`, `get_notes` returns at most `n` entries
37    /// (total, across all tags) in one call. Used to exercise client-side pagination drain loops.
38    /// `None` = unbounded (legacy behavior).
39    max_batch: Option<usize>,
40    /// Notes stored through the with-proof path, with the block their proof named.
41    proven_notes: BTreeMap<NoteId, BlockNumber>,
42}
43
44impl MockNoteTransportNode {
45    pub fn new() -> Self {
46        Self {
47            notes: BTreeMap::default(),
48            nonce: 1,
49            next_sequence: 1,
50            max_batch: None,
51            proven_notes: BTreeMap::default(),
52        }
53    }
54
55    /// Build a mock that caps each `get_notes` response at `max_batch` entries.
56    pub fn with_max_batch(max_batch: usize) -> Self {
57        Self {
58            notes: BTreeMap::default(),
59            nonce: 1,
60            next_sequence: 1,
61            max_batch: Some(max_batch),
62            proven_notes: BTreeMap::default(),
63        }
64    }
65
66    /// Seed a note relayed with its inclusion proof. The real service verifies the proof against
67    /// its node; the mock only records the proof's block and serves it as the commitment block.
68    ///
69    /// The real service stores a note only once. The mock also ignores a note with an id that it
70    /// already stores.
71    pub fn add_note_with_proof(
72        &mut self,
73        header: NoteHeader,
74        details_bytes: Vec<u8>,
75        inclusion_proof: &NoteInclusionProof,
76    ) {
77        if self.contains_note(&header.id()) {
78            return;
79        }
80        let block_num = inclusion_proof.location().block_num();
81        self.proven_notes.insert(header.id(), block_num);
82        self.add_note_after(header, details_bytes, Some(block_num));
83    }
84
85    /// Returns whether the mock stores a note with `note_id`.
86    fn contains_note(&self, note_id: &NoteId) -> bool {
87        self.notes.values().flatten().any(|(info, _)| info.header.id() == *note_id)
88    }
89
90    /// Returns the block named by the proof a note was stored with, or `None` when the note was not
91    /// stored through the with-proof path.
92    pub fn proven_block(&self, note_id: &NoteId) -> Option<BlockNumber> {
93        self.proven_notes.get(note_id).copied()
94    }
95
96    pub fn add_note(&mut self, header: NoteHeader, details_bytes: Vec<u8>) {
97        self.add_note_after(header, details_bytes, None);
98    }
99
100    /// Seed a note that carries a commitment block floor, as a transport that stored the note
101    /// without verifying a proof would serve it.
102    pub fn add_note_after(
103        &mut self,
104        header: NoteHeader,
105        details_bytes: Vec<u8>,
106        block_hint: Option<BlockNumber>,
107    ) {
108        let tag = header.metadata().tag();
109        let info = NoteInfo { header, details_bytes, block_hint };
110        let cursor = NoteTransportCursor::from_parts(self.nonce, self.next_sequence);
111        self.next_sequence += 1;
112        self.notes.entry(tag).or_default().push((info, cursor));
113    }
114
115    /// Seed a note under an arbitrary transport tag key, regardless of the note's own tag.
116    pub fn add_note_with_tag_key(
117        &mut self,
118        tag: NoteTag,
119        header: NoteHeader,
120        details_bytes: Vec<u8>,
121    ) {
122        let info = NoteInfo { header, details_bytes, block_hint: None };
123        let cursor = NoteTransportCursor::from_parts(self.nonce, self.next_sequence);
124        self.next_sequence += 1;
125        self.notes.entry(tag).or_default().push((info, cursor));
126    }
127
128    pub fn get_notes(
129        &self,
130        tags: &[NoteTag],
131        cursor: NoteTransportCursor,
132    ) -> (Vec<NoteInfo>, NoteTransportCursor) {
133        let cursor = if cursor.parts().is_some_and(|(nonce, _)| nonce != self.nonce) {
134            NoteTransportCursor::init()
135        } else {
136            cursor
137        };
138        // Start `rcursor` at the input — matches the real server's contract (`rcursor = max(cursor,
139        // max_seq_returned)`), so an empty batch returns the caller's own cursor rather than
140        // `init()`.
141        let mut collected: Vec<(NoteInfo, NoteTransportCursor)> = vec![];
142        for tag in tags {
143            // Assumes stored notes are ordered by cursor
144            let tnotes = self
145                .notes
146                .get(tag)
147                .map(|pg_notes| {
148                    // Find first element after cursor
149                    if let Some(pos) = pg_notes.iter().position(|(_, tcursor)| *tcursor > cursor) {
150                        &pg_notes[pos..]
151                    } else {
152                        &[]
153                    }
154                })
155                .map(Vec::from)
156                .unwrap_or_default();
157            collected.extend(tnotes);
158        }
159
160        // Deterministic ordering across tags: sort by cursor ascending so the client sees notes in
161        // per-cursor order regardless of tag iteration order, matching the real server's `ORDER BY
162        // seq ASC`.
163        collected.sort_by_key(|(_, c)| *c);
164
165        // Apply the batch cap, if configured.
166        if let Some(max) = self.max_batch {
167            collected.truncate(max);
168        }
169
170        let initial_cursor = cursor
171            .parts()
172            .map_or(NoteTransportCursor::from_parts(self.nonce, 0), |_| cursor);
173        let rcursor = collected.iter().map(|(_, c)| *c).max().unwrap_or(initial_cursor);
174        let notes = collected.into_iter().map(|(n, _)| n).collect();
175        (notes, rcursor)
176    }
177}
178
179impl Default for MockNoteTransportNode {
180    fn default() -> Self {
181        Self::new()
182    }
183}
184
185/// Mock Note Transport API
186///
187/// Simulates communications with the note transport node.
188#[derive(Clone, Default)]
189pub struct MockNoteTransportApi {
190    pub mock_node: Arc<RwLock<MockNoteTransportNode>>,
191    fetch_tag_counts: Arc<RwLock<Vec<usize>>>,
192}
193
194impl MockNoteTransportApi {
195    pub fn new(mock_node: Arc<RwLock<MockNoteTransportNode>>) -> Self {
196        Self {
197            mock_node,
198            fetch_tag_counts: Arc::default(),
199        }
200    }
201
202    /// Returns the number of tags in each fetch request.
203    pub fn fetch_tag_counts(&self) -> Vec<usize> {
204        self.fetch_tag_counts.read().clone()
205    }
206}
207
208impl MockNoteTransportApi {
209    pub fn send_note_with_proof(&self, note: TransportNote, inclusion_proof: &NoteInclusionProof) {
210        let (header, details) = note.into_parts();
211        let details_bytes = details.to_bytes();
212        self.mock_node
213            .write()
214            .add_note_with_proof(header, details_bytes, inclusion_proof);
215    }
216
217    pub fn fetch_notes(
218        &self,
219        tags: &[NoteTag],
220        cursor: NoteTransportCursor,
221    ) -> (Vec<NoteInfo>, NoteTransportCursor) {
222        self.mock_node.read().get_notes(tags, cursor)
223    }
224}
225
226#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
227#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
228impl NoteTransportClient for MockNoteTransportApi {
229    async fn send_note_with_proof(
230        &self,
231        note: TransportNote,
232        inclusion_proof: NoteInclusionProof,
233    ) -> Result<(), NoteTransportError> {
234        self.send_note_with_proof(note, &inclusion_proof);
235        Ok(())
236    }
237
238    async fn fetch_notes(
239        &self,
240        tags: &[NoteTag],
241        cursor: NoteTransportCursor,
242    ) -> Result<(Vec<NoteInfo>, NoteTransportCursor), NoteTransportError> {
243        let page = self.fetch_notes_page(tags, cursor).await?;
244        Ok((page.notes, page.cursor))
245    }
246
247    async fn fetch_notes_page(
248        &self,
249        tags: &[NoteTag],
250        cursor: NoteTransportCursor,
251    ) -> Result<NoteTransportPage, NoteTransportError> {
252        self.fetch_tag_counts.write().push(tags.len());
253        let node = self.mock_node.read();
254        let (notes, cursor) = node.get_notes(tags, cursor);
255        let has_more = !node.get_notes(tags, cursor).0.is_empty();
256        Ok(NoteTransportPage { notes, cursor, has_more })
257    }
258}
259
260// FAULTY NOTE TRANSPORT API
261// ================================================================================================
262
263/// Test-only [`NoteTransportClient`] decorator that injects controlled failures into
264/// `send_note_with_proof` calls.
265///
266/// Reproduces the failure mode where the NTL is reachable but rejects a send. Tests use it to check
267/// how [`Client::send_private_note_with_proof`](crate::Client::send_private_note_with_proof)
268/// reports the failure and how a later send by the caller delivers the note.
269///
270/// The decorator counts attempts (`send_attempts`) and lets a test specify how many of the next
271/// `send_note_with_proof` calls should fail (`fail_next`); successful calls delegate to an inner
272/// [`MockNoteTransportApi`]. `fetch_notes` failures can be injected separately via
273/// [`FaultyNoteTransportApi::fail_next_n_fetches`].
274pub struct FaultyNoteTransportApi {
275    inner: MockNoteTransportApi,
276    fail_next: AtomicUsize,
277    send_attempts: AtomicUsize,
278    fail_next_fetches: AtomicUsize,
279    fail_on_fetch: AtomicUsize,
280    fetch_attempts: AtomicUsize,
281}
282
283impl FaultyNoteTransportApi {
284    /// Create a faulty transport that fails the next `fail_next` `send_note_with_proof` calls
285    /// before delegating to the inner mock.
286    pub fn new(mock_node: Arc<RwLock<MockNoteTransportNode>>, fail_next: usize) -> Self {
287        Self {
288            inner: MockNoteTransportApi::new(mock_node),
289            fail_next: AtomicUsize::new(fail_next),
290            send_attempts: AtomicUsize::new(0),
291            fail_next_fetches: AtomicUsize::new(0),
292            fail_on_fetch: AtomicUsize::new(0),
293            fetch_attempts: AtomicUsize::new(0),
294        }
295    }
296
297    /// Total `send_note_with_proof` calls observed (success + failure).
298    pub fn send_attempts(&self) -> usize {
299        self.send_attempts.load(Ordering::SeqCst)
300    }
301
302    /// Fail the next `n` `fetch_notes` calls before delegating to the inner mock again.
303    pub fn fail_next_n_fetches(&self, n: usize) {
304        self.fail_next_fetches.store(n, Ordering::SeqCst);
305    }
306
307    /// Fails one fetch attempt. Attempt numbers start at one.
308    pub fn fail_on_fetch_attempt(&self, attempt: usize) {
309        self.fail_on_fetch.store(attempt, Ordering::SeqCst);
310    }
311
312    /// Total `fetch_notes` calls observed (success + failure).
313    pub fn fetch_attempts(&self) -> usize {
314        self.fetch_attempts.load(Ordering::SeqCst)
315    }
316
317    /// Records a send attempt and returns whether it must fail.
318    fn take_send_failure(&self) -> Option<NoteTransportError> {
319        self.send_attempts.fetch_add(1, Ordering::SeqCst);
320        self.fail_next
321            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |n| n.checked_sub(1))
322            .is_ok()
323            .then(|| {
324                NoteTransportError::Network(
325                    "FaultyNoteTransportApi: simulated send_note_with_proof failure".to_string(),
326                )
327            })
328    }
329}
330
331#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
332#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
333impl NoteTransportClient for FaultyNoteTransportApi {
334    async fn send_note_with_proof(
335        &self,
336        note: TransportNote,
337        inclusion_proof: NoteInclusionProof,
338    ) -> Result<(), NoteTransportError> {
339        if let Some(error) = self.take_send_failure() {
340            return Err(error);
341        }
342        self.inner.send_note_with_proof(note, &inclusion_proof);
343        Ok(())
344    }
345
346    async fn fetch_notes(
347        &self,
348        tags: &[NoteTag],
349        cursor: NoteTransportCursor,
350    ) -> Result<(Vec<NoteInfo>, NoteTransportCursor), NoteTransportError> {
351        let attempt = self.fetch_attempts.fetch_add(1, Ordering::SeqCst) + 1;
352        let should_fail = self
353            .fail_next_fetches
354            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |n| n.checked_sub(1))
355            .is_ok();
356        if should_fail || attempt == self.fail_on_fetch.load(Ordering::SeqCst) {
357            return Err(NoteTransportError::Network(
358                "FaultyNoteTransportApi: simulated fetch_notes failure".to_string(),
359            ));
360        }
361        Ok(self.inner.fetch_notes(tags, cursor))
362    }
363}
364
365// SERIALIZATION
366// ================================================================================================
367
368impl Serializable for MockNoteTransportNode {
369    fn write_into<W: ByteWriter>(&self, target: &mut W) {
370        self.notes.write_into(target);
371        self.nonce.write_into(target);
372        self.next_sequence.write_into(target);
373    }
374}
375
376impl Deserializable for MockNoteTransportNode {
377    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
378        let notes = BTreeMap::<NoteTag, Vec<(NoteInfo, NoteTransportCursor)>>::read_from(source)?;
379        let nonce = u64::read_from(source)?;
380        let next_sequence = u64::read_from(source)?;
381
382        Ok(Self {
383            notes,
384            nonce,
385            next_sequence,
386            max_batch: None,
387            proven_notes: BTreeMap::default(),
388        })
389    }
390}
miden_client/test_utils/note_transport.rs

miden_client/test_utils/
note_transport.rs
