Add BlockStreamerTestExtensions trait for tests

Author: dlachaume

mithril-common/src/cardano_block_scanner/block_scanner.rs

@@ -60,6 +60,7 @@ impl BlockScanner for CardanoBlockScanner {
 
 #[cfg(test)]
 mod tests {
+    use crate::cardano_block_scanner::BlockStreamerTestExtensions;
     use crate::test_utils::{TempDir, TestLogger};
 
     use super::*;

mithril-common/src/cardano_block_scanner/dumb_block_scanner.rs

@@ -68,6 +68,8 @@ impl BlockStreamer for DumbBlockStreamer {
 
 #[cfg(test)]
 mod tests {
+    use crate::cardano_block_scanner::BlockStreamerTestExtensions;
+
     use super::*;
 
     #[tokio::test]

mithril-common/src/cardano_block_scanner/immutable_block_streamer.rs

@@ -124,6 +124,7 @@ impl ImmutableBlockStreamer {
 
 #[cfg(test)]
 mod tests {
+    use crate::cardano_block_scanner::BlockStreamerTestExtensions;
     use crate::test_utils::{TempDir, TestLogger};
 
     use super::*;

mithril-common/src/cardano_block_scanner/interface.rs

@@ -70,20 +70,32 @@ pub enum ChainScannedBlocks {
 pub trait BlockStreamer: Sync + Send {
     /// Stream the next available blocks
     async fn poll_next(&mut self) -> StdResult<Option<ChainScannedBlocks>>;
+}
+
+cfg_test_tools! {
+    /// Tests extensions methods for the [BlockStreamer] trait.
+    #[async_trait]
+    pub trait BlockStreamerTestExtensions{
+        /// Stream all the available blocks, may be very memory intensive
+        async fn poll_all(&mut self) -> StdResult<Vec<ScannedBlock>>;
+    }
 
-    /// Stream all the available blocks, may be very memory intensive
-    async fn poll_all(&mut self) -> StdResult<Vec<ScannedBlock>> {
-        let mut all_blocks = Vec::new();
-        while let Some(next_blocks) = self.poll_next().await? {
-            match next_blocks {
-                ChainScannedBlocks::RollForwards(mut forward_blocks) => {
-                    all_blocks.append(&mut forward_blocks);
-                }
-                ChainScannedBlocks::RollBackward(_) => {
-                    return Err(anyhow!("poll_all: RollBackward not supported"));
-                }
-            };
+    #[async_trait]
+    impl <S: BlockStreamer + ?Sized> BlockStreamerTestExtensions for S {
+        async fn poll_all(&mut self) -> StdResult<Vec<ScannedBlock>> {
+            let mut all_blocks = Vec::new();
+            while let Some(next_blocks) = self.poll_next().await? {
+                match next_blocks {
+                    ChainScannedBlocks::RollForwards(mut forward_blocks) => {
+                        all_blocks.append(&mut forward_blocks);
+                    }
+                    ChainScannedBlocks::RollBackward(_) => {
+                        return Err(anyhow!("poll_all: RollBackward not supported"));
+                    }
+                };
+            }
+            Ok(all_blocks)
         }
-        Ok(all_blocks)
     }
+
 }