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#[derive(Clone)]
35pub struct MockNoteTransportNode {
36 notes: BTreeMap<NoteTag, Vec<(NoteInfo, NoteTransportCursor)>>,
37 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 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 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 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 let mut collected: Vec<(NoteInfo, NoteTransportCursor)> = vec![];
98 for tag in tags {
99 let tnotes = self
101 .notes
102 .get(tag)
103 .map(|pg_notes| {
104 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 collected.sort_by_key(|(_, c)| *c);
120
121 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#[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
224pub 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 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 pub fn fail_next_n(&self, n: usize) {
265 self.fail_next.store(n, Ordering::SeqCst);
266 }
267
268 pub fn send_attempts(&self) -> usize {
270 self.send_attempts.load(Ordering::SeqCst)
271 }
272
273 pub fn fail_next_n_fetches(&self, n: usize) {
276 self.fail_next_fetches.store(n, Ordering::SeqCst);
277 }
278
279 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
354impl 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}