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