Skip to main content

miden_client/test_utils/
note_transport.rs

1use alloc::boxed::Box;
2use alloc::collections::BTreeMap;
3use alloc::sync::Arc;
4use alloc::vec::Vec;
5use core::pin::Pin;
6use core::task::{Context, Poll};
7
8use chrono::Utc;
9use futures::Stream;
10use miden_protocol::note::{NoteHeader, NoteTag};
11use miden_tx::utils::sync::RwLock;
12use miden_tx::utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
13
14use crate::note_transport::{
15    NoteInfo,
16    NoteStream,
17    NoteTransportClient,
18    NoteTransportCursor,
19    NoteTransportError,
20};
21
22/// Mock Note Transport Node
23///
24/// Simulates the functionality of the note transport node.
25#[derive(Clone)]
26pub struct MockNoteTransportNode {
27    notes: BTreeMap<NoteTag, Vec<(NoteInfo, NoteTransportCursor)>>,
28}
29
30impl MockNoteTransportNode {
31    pub fn new() -> Self {
32        Self { notes: BTreeMap::default() }
33    }
34
35    pub fn add_note(&mut self, header: NoteHeader, details_bytes: Vec<u8>) {
36        let tag = header.metadata().tag();
37        let info = NoteInfo { header, details_bytes };
38        let cursor = u64::try_from(Utc::now().timestamp_micros()).unwrap();
39        self.notes.entry(tag).or_default().push((info, cursor.into()));
40    }
41
42    pub fn get_notes(
43        &self,
44        tags: &[NoteTag],
45        cursor: NoteTransportCursor,
46    ) -> (Vec<NoteInfo>, NoteTransportCursor) {
47        let mut notes = vec![];
48        let mut rcursor = NoteTransportCursor::init();
49        for tag in tags {
50            // Assumes stored notes are ordered by cursor
51            let tnotes = self
52                .notes
53                .get(tag)
54                .map(|pg_notes| {
55                    // Find first element after cursor
56                    if let Some(pos) = pg_notes.iter().position(|(_, tcursor)| *tcursor > cursor) {
57                        &pg_notes[pos..]
58                    } else {
59                        &[]
60                    }
61                })
62                .map(Vec::from)
63                .unwrap_or_default();
64            rcursor = rcursor.max(
65                tnotes
66                    .iter()
67                    .map(|(_, cursor)| *cursor)
68                    .max()
69                    .unwrap_or(NoteTransportCursor::init()),
70            );
71            notes.extend(tnotes.into_iter().map(|(note, _)| note).collect::<Vec<_>>());
72        }
73        (notes, rcursor)
74    }
75}
76
77impl Default for MockNoteTransportNode {
78    fn default() -> Self {
79        Self::new()
80    }
81}
82
83/// Mock Note Transport API
84///
85/// Simulates communications with the note transport node.
86#[derive(Clone, Default)]
87pub struct MockNoteTransportApi {
88    pub mock_node: Arc<RwLock<MockNoteTransportNode>>,
89}
90
91impl MockNoteTransportApi {
92    pub fn new(mock_node: Arc<RwLock<MockNoteTransportNode>>) -> Self {
93        Self { mock_node }
94    }
95}
96
97impl MockNoteTransportApi {
98    pub fn send_note(&self, header: NoteHeader, details_bytes: Vec<u8>) {
99        self.mock_node.write().add_note(header, details_bytes);
100    }
101
102    pub fn fetch_notes(
103        &self,
104        tags: &[NoteTag],
105        cursor: NoteTransportCursor,
106    ) -> (Vec<NoteInfo>, NoteTransportCursor) {
107        self.mock_node.read().get_notes(tags, cursor)
108    }
109}
110
111pub struct DummyNoteStream {}
112impl Stream for DummyNoteStream {
113    type Item = Result<Vec<NoteInfo>, NoteTransportError>;
114
115    fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
116        Poll::Ready(None)
117    }
118}
119impl NoteStream for DummyNoteStream {}
120
121#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
122#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
123impl NoteTransportClient for MockNoteTransportApi {
124    async fn send_note(
125        &self,
126        header: NoteHeader,
127        details: Vec<u8>,
128    ) -> Result<(), NoteTransportError> {
129        self.send_note(header, details);
130        Ok(())
131    }
132
133    async fn fetch_notes(
134        &self,
135        tags: &[NoteTag],
136        cursor: NoteTransportCursor,
137    ) -> Result<(Vec<NoteInfo>, NoteTransportCursor), NoteTransportError> {
138        Ok(self.fetch_notes(tags, cursor))
139    }
140
141    async fn stream_notes(
142        &self,
143        _tag: NoteTag,
144        _cursor: NoteTransportCursor,
145    ) -> Result<Box<dyn NoteStream>, NoteTransportError> {
146        Ok(Box::new(DummyNoteStream {}))
147    }
148}
149
150// SERIALIZATION
151// ================================================================================================
152
153impl Serializable for MockNoteTransportNode {
154    fn write_into<W: ByteWriter>(&self, target: &mut W) {
155        self.notes.write_into(target);
156    }
157}
158
159impl Deserializable for MockNoteTransportNode {
160    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
161        let notes = BTreeMap::<NoteTag, Vec<(NoteInfo, NoteTransportCursor)>>::read_from(source)?;
162
163        Ok(Self { notes })
164    }
165}