From e64f459bf8462907d9942bc5ec3d8b8713974654 Mon Sep 17 00:00:00 2001
From: girazoki <gorka.irazoki@gmail.com>
Date: Fri, 30 Aug 2024 17:38:03 +0200
Subject: [PATCH] unfold

---
 client/consensus/src/mocks.rs | 35 ++++++++++++++++++++---------------
 1 file changed, 20 insertions(+), 15 deletions(-)

diff --git a/client/consensus/src/mocks.rs b/client/consensus/src/mocks.rs
index 860fa2808..07964e709 100644
--- a/client/consensus/src/mocks.rs
+++ b/client/consensus/src/mocks.rs
@@ -299,21 +299,26 @@ impl RelayChainInterface for RelayChain {
     async fn import_notification_stream(
         &self,
     ) -> RelayChainResult<Pin<Box<dyn Stream<Item = PHeader> + Send>>> {
-        let mut notifications = vec![];
-        for _n in 0..self.block_import_iterations {
-            // It does not matter if we push always the same header, blocks will be created anyway
-            notifications.push(PHeader {
-                parent_hash: Default::default(),
-                number: 1u32,
-                // The state trie merkle root
-                state_root: Default::default(),
-                // The merkle root of the extrinsics.
-                extrinsics_root: Default::default(),
-                // A chain-specific digest of data useful for light clients or referencing auxiliary data.
-                digest: Default::default(),
-            });
-        }
-        Ok(Box::pin(futures::stream::iter(notifications)))
+        let block_import_iterations = self.block_import_iterations;
+        // Create the stream using unfold with the counter as the state
+        let stream = stream::unfold(0, move |counter| async move {
+            if counter < block_import_iterations {
+                // Create the next PHeader
+                let header = PHeader {
+                    parent_hash: Default::default(),
+                    number: 1u32,
+                    state_root: Default::default(),
+                    extrinsics_root: Default::default(),
+                    digest: Default::default(),
+                };
+                let next_state = counter + 1;
+                Some((header, next_state))
+            } else {
+                None
+            }
+        });
+
+        Ok(Box::pin(stream))
     }
 
     async fn finality_notification_stream(