mocker: add functions to get and take outbuffer

bastibl · Nov 13, 2024 · 7ba136b · 7ba136b
1 parent 6a92dbb
commit 7ba136b
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 23 deletions.
diff --git a/src/runtime/mocker.rs b/src/runtime/mocker.rs
@@ -60,7 +60,7 @@ impl<K: Kernel + 'static> Mocker<K> {
     /// Initialize output buffer with given size
     pub fn init_output<T>(&mut self, id: usize, size: usize)
     where
-        T: Debug + Send + 'static,
+        T: Clone + Debug + Send + 'static,
     {
         self.block
             .sio
@@ -69,32 +69,26 @@ impl<K: Kernel + 'static> Mocker<K> {
     }
 
     /// Get data from output buffer
-    pub fn output<T>(&mut self, id: usize) -> Vec<T>
+    pub fn output<T>(&mut self, id: usize) -> (Vec<T>, Vec<ItemTag>)
     where
-        T: Debug + Send + 'static,
+        T: Clone + Debug + Send + 'static,
     {
         let w = self.block.sio.output(id).writer_mut();
         if let BufferWriter::Host(w) = w {
-            w.as_any().downcast_mut::<MockWriter<T>>().unwrap().get()
+            w.as_any().downcast_ref::<MockWriter<T>>().unwrap().get()
         } else {
             panic!("mocker: wrong output buffer (expected CPU, got Custom)");
         }
     }
 
-    /// Get the list of tags from the output buffer at `id`.
-    ///
-    /// Type parameter `T` should be the same as the type of the output buffer.
-    pub fn output_tags<T>(&mut self, id: usize) -> Vec<ItemTag>
+    /// Taking data from output buffer, freeing up the buffer
+    pub fn take_output<T>(&mut self, id: usize) -> (Vec<T>, Vec<ItemTag>)
     where
-        T: Debug + Send + 'static,
+        T: Clone + Debug + Send + 'static,
     {
         let w = self.block.sio.output(id).writer_mut();
         if let BufferWriter::Host(w) = w {
-            w.as_any()
-                .downcast_mut::<MockWriter<T>>()
-                .unwrap()
-                .tags()
-                .to_vec()
+            w.as_any().downcast_mut::<MockWriter<T>>().unwrap().take()
         } else {
             panic!("mocker: wrong output buffer (expected CPU, got Custom)");
         }
@@ -207,30 +201,33 @@ impl<T: Debug + Send + 'static> BufferReaderHost for MockReader<T> {
 }
 
 #[derive(Debug)]
-struct MockWriter<T: Debug + Send + 'static> {
+struct MockWriter<T: Clone + Debug + Send + 'static> {
     data: Vec<T>,
     tags: Vec<ItemTag>,
 }
 
-impl<T: Debug + Send + 'static> MockWriter<T> {
+impl<T: Clone + Debug + Send + 'static> MockWriter<T> {
     pub fn new(size: usize) -> Self {
         MockWriter::<T> {
             data: Vec::with_capacity(size),
             tags: Vec::new(),
         }
     }
 
-    pub fn get(&mut self) -> Vec<T> {
-        std::mem::take(&mut self.data)
+    pub fn get(&self) -> (Vec<T>, Vec<ItemTag>) {
+        (self.data.clone(), self.tags.clone())
     }
 
-    pub fn tags(&self) -> &[ItemTag] {
-        &self.tags
+    pub fn take(&mut self) -> (Vec<T>, Vec<ItemTag>) {
+        let (data, tags) = self.get();
+        self.data.clear();
+        self.tags = Vec::new();
+        (data, tags)
     }
 }
 
 #[async_trait]
-impl<T: Debug + Send + 'static> BufferWriterHost for MockWriter<T> {
+impl<T: Clone + Debug + Send + 'static> BufferWriterHost for MockWriter<T> {
     fn add_reader(
         &mut self,
         _reader_inbox: Sender<BlockMessage>,

diff --git a/tests/tag.rs b/tests/tag.rs
@@ -30,12 +30,17 @@ fn tags_through_mock() -> Result<()> {
     mock.init_output::<f32>(0, input.len() * 2);
     mock.input(0, input.clone());
     mock.run();
+
+    let (out_buffer, out_tags) = mock.output::<f32>(0);
+    assert_eq!(out_buffer.len(), 1024);
+    assert_eq!(out_tags.len(), 0);
+
     mock.input_with_tags(0, input, tags.clone());
     mock.run();
     mock.deinit();
 
-    let out_tags = mock.output_tags::<f32>(0);
-
+    let (out_buffer, out_tags) = mock.output::<f32>(0);
+    assert_eq!(out_buffer.len(), 2048);
     assert_eq!(out_tags.len(), 3);
 
     for (i, tag) in tags.iter().enumerate() {
@@ -46,5 +51,13 @@ fn tags_through_mock() -> Result<()> {
         assert!(matches!(out_tags[i].tag, Tag::Id(t) if t == tag_id));
     }
 
+    let (out_buffer, out_tags) = mock.take_output::<f32>(0);
+    assert_eq!(out_buffer.len(), 2048);
+    assert_eq!(out_tags.len(), 3);
+
+    let (out_buffer, out_tags) = mock.output::<f32>(0);
+    assert_eq!(out_buffer.len(), 0);
+    assert_eq!(out_tags.len(), 0);
+
     Ok(())
 }