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#[derive(Clone)]
34pub struct MockNoteTransportNode {
35 notes: BTreeMap<NoteTag, Vec<(NoteInfo, NoteTransportCursor)>>,
36 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 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 let mut collected: Vec<(NoteInfo, NoteTransportCursor)> = vec![];
74 for tag in tags {
75 let tnotes = self
77 .notes
78 .get(tag)
79 .map(|pg_notes| {
80 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 collected.sort_by_key(|(_, c)| *c);
96
97 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#[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
181pub struct FaultyNoteTransportApi {
198 inner: MockNoteTransportApi,
199 fail_next: AtomicUsize,
200 send_attempts: AtomicUsize,
201}
202
203impl FaultyNoteTransportApi {
204 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 pub fn fail_next_n(&self, n: usize) {
217 self.fail_next.store(n, Ordering::SeqCst);
218 }
219
220 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
265impl 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}