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