fix(plasma): Unit Test the Plasma Source

crates/derive/src/traits/test_utils.rs

@@ -67,7 +67,7 @@ impl DataAvailabilityProvider for TestDAP {
 #[derive(Debug, Clone, Default)]
 pub struct TestChainProvider {
     /// Maps block numbers to block information using a tuple list.
-    pub blocks: Vec<(u64, BlockInfo)>,
+    pub blocks: Vec<(u64, BlockInfo, Vec<TxEnvelope>)>,
     /// Maps block hashes to header information using a tuple list.
     pub headers: Vec<(B256, Header)>,
     /// Maps block hashes to receipts using a tuple list.
@@ -77,7 +77,7 @@ pub struct TestChainProvider {
 impl TestChainProvider {
     /// Insert a block into the mock chain provider.
     pub fn insert_block(&mut self, number: u64, block: BlockInfo) {
-        self.blocks.push((number, block));
+        self.blocks.push((number, block, Vec::new()));
     }
 
     /// Insert receipts into the mock chain provider.
@@ -124,7 +124,7 @@ impl ChainProvider for TestChainProvider {
     }
 
     async fn block_info_by_number(&mut self, _number: u64) -> Result<BlockInfo> {
-        if let Some((_, block)) = self.blocks.iter().find(|(n, _)| *n == _number) {
+        if let Some((_, block, _)) = self.blocks.iter().find(|(n, _, _)| *n == _number) {
             Ok(*block)
         } else {
             Err(anyhow::anyhow!("Block not found"))
@@ -143,13 +143,13 @@ impl ChainProvider for TestChainProvider {
         &mut self,
         hash: B256,
     ) -> Result<(BlockInfo, Vec<TxEnvelope>)> {
-        let block = self
+        let (block, txs) = self
             .blocks
             .iter()
-            .find(|(_, b)| b.hash == hash)
-            .map(|(_, b)| *b)
+            .find(|(_, b, _)| b.hash == hash)
+            .map(|(_, b, v)| (*b, v.clone()))
             .ok_or_else(|| anyhow::anyhow!("Block not found"))?;
-        Ok((block, Vec::new()))
+        Ok((block, txs))
     }
 }
 

crates/derive/src/types/errors.rs

@@ -62,6 +62,12 @@ impl PartialEq<StageError> for StageError {
         if let (StageError::Reset(a), StageError::Reset(b)) = (self, other) {
             return a == b;
         }
+        if let (StageError::AttributesBuild(a), StageError::AttributesBuild(b)) = (self, other) {
+            return a == b;
+        }
+        if let (StageError::BlockFetch(a), StageError::BlockFetch(b)) = (self, other) {
+            return a == b;
+        }
         matches!(
             (self, other),
             (StageError::Eof, StageError::Eof) |
@@ -72,7 +78,6 @@ impl PartialEq<StageError> for StageError {
                 (StageError::NoChannel, StageError::NoChannel) |
                 (StageError::ChannelNotFound, StageError::ChannelNotFound) |
                 (StageError::MissingOrigin, StageError::MissingOrigin) |
-                (StageError::AttributesBuild(_), StageError::AttributesBuild(_)) |
                 (StageError::ReceiptFetch(_), StageError::ReceiptFetch(_)) |
                 (StageError::BlockInfoFetch(_), StageError::BlockInfoFetch(_)) |
                 (StageError::SystemConfigUpdate(_), StageError::SystemConfigUpdate(_)) |

crates/plasma/src/plasma.rs

@@ -10,7 +10,7 @@ use alloy_primitives::{Address, Bytes};
 use anyhow::Result;
 use async_trait::async_trait;
 use kona_derive::{
-    traits::{ChainProvider, DataAvailabilityProvider},
+    traits::{AsyncIterator, ChainProvider, DataAvailabilityProvider},
     types::{BlockInfo, RollupConfig},
 };
 use kona_primitives::BlockID;
@@ -21,7 +21,7 @@ pub struct PlasmaDataSource<C, PIF, I>
 where
     C: ChainProvider + Send + Clone,
     PIF: PlasmaInputFetcher<C> + Clone,
-    I: Iterator<Item = Bytes> + Send + Clone,
+    I: AsyncIterator<Item = Bytes> + Send + Clone,
 {
     /// The chain provider to use for the factory.
     pub chain_provider: C,
@@ -37,7 +37,7 @@ impl<C, PIF, I> PlasmaDataSource<C, PIF, I>
 where
     C: ChainProvider + Send + Clone + Debug,
     PIF: PlasmaInputFetcher<C> + Clone,
-    I: Iterator<Item = Bytes> + Send + Clone,
+    I: AsyncIterator<Item = Bytes> + Send + Clone,
 {
     /// Creates a new factory.
     pub fn new(provider: C, pif: PIF, s: I, cfg: &RollupConfig) -> Self {
@@ -55,7 +55,7 @@ impl<C, PIF, I> DataAvailabilityProvider for PlasmaDataSource<C, PIF, I>
 where
     C: ChainProvider + Send + Clone + Debug + Sync,
     PIF: PlasmaInputFetcher<C> + Clone + Debug + Send + Sync,
-    I: Iterator<Item = Bytes> + Send + Clone + Debug + Sync,
+    I: AsyncIterator<Item = Bytes> + Send + Clone + Debug + Sync,
 {
     type Item = Bytes;
     type DataIter = PlasmaSource<C, PIF, I>;