Skip to main content

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