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

miden_client/test_utils/
note_transport.rs
