chore(derive): data source unit tests

crates/derive/src/sources/ethereum.rs

@@ -78,3 +78,71 @@ where
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use alloy_consensus::TxEnvelope;
+    use alloy_eips::eip2718::Decodable2718;
+    use alloy_primitives::{address, Address};
+    use kona_primitives::{BlockInfo, RollupConfig};
+
+    use crate::{
+        sources::{EthereumDataSource, EthereumDataSourceVariant},
+        traits::{
+            test_utils::{TestBlobProvider, TestChainProvider},
+            AsyncIterator, DataAvailabilityProvider,
+        },
+    };
+
+    #[tokio::test]
+    async fn test_validate_ethereum_data_source() {
+        let chain = TestChainProvider::default();
+        let blob = TestBlobProvider::default();
+        let block_ref = BlockInfo::default();
+        let batcher_address = Address::default();
+
+        // If the ecotone_timestamp is not set, a Calldata source should be returned.
+        let cfg = RollupConfig { ecotone_time: None, ..Default::default() };
+        let data_source = EthereumDataSource::new(chain.clone(), blob.clone(), &cfg);
+        let data_iter = data_source.open_data(&block_ref, batcher_address).await.unwrap();
+        assert!(matches!(data_iter, EthereumDataSourceVariant::Calldata(_)));
+
+        // If the ecotone_timestamp is set, and the block_ref timestamp is prior to the
+        // ecotone_timestamp, a calldata source is created.
+        let cfg = RollupConfig { ecotone_time: Some(100), ..Default::default() };
+        let data_source = EthereumDataSource::new(chain, blob, &cfg);
+        let data_iter = data_source.open_data(&block_ref, batcher_address).await.unwrap();
+        assert!(matches!(data_iter, EthereumDataSourceVariant::Calldata(_)));
+
+        // If the ecotone_timestamp is set, and the block_ref timestamp is greater than
+        // or equal to the ecotone_timestamp, a Blob source is created.
+        let block_ref = BlockInfo { timestamp: 101, ..Default::default() };
+        let data_iter = data_source.open_data(&block_ref, batcher_address).await.unwrap();
+        assert!(matches!(data_iter, EthereumDataSourceVariant::Blob(_)));
+    }
+
+    #[tokio::test]
+    async fn test_open_ethereum_calldata_source_pre_ecotone() {
+        let mut chain = TestChainProvider::default();
+        let blob = TestBlobProvider::default();
+        let batcher_address = address!("6887246668a3b87F54DeB3b94Ba47a6f63F32985");
+        let batch_inbox = address!("FF00000000000000000000000000000000000010");
+        let block_ref = BlockInfo { number: 10, ..Default::default() };
+
+        let mut cfg = RollupConfig::default();
+        cfg.genesis.system_config.batcher_addr = batcher_address;
+
+        // load a test batcher transaction
+        let raw_batcher_tx = include_bytes!("../../testdata/raw_batcher_tx.hex");
+        let tx = TxEnvelope::decode_2718(&mut raw_batcher_tx.as_ref()).unwrap();
+        chain.insert_block_with_transactions(10, block_ref, alloc::vec![tx]);
+
+        let data_source = EthereumDataSource::new(chain, blob, &cfg);
+        let mut data_iter = data_source.open_data(&block_ref, batch_inbox).await.unwrap();
+        assert!(matches!(data_iter, EthereumDataSourceVariant::Calldata(_)));
+
+        // Should successfully retrieve a calldata batch from the block
+        let calldata_batch = data_iter.next().await.unwrap().unwrap();
+        assert_eq!(calldata_batch.len(), 119823);
+    }
+}

crates/derive/src/traits/test_utils.rs

@@ -1,8 +1,10 @@
 //! Test Utilities for derive traits
 
 use crate::{
-    traits::{AsyncIterator, ChainProvider, DataAvailabilityProvider, L2ChainProvider},
-    types::{StageError, StageResult},
+    traits::{
+        AsyncIterator, BlobProvider, ChainProvider, DataAvailabilityProvider, L2ChainProvider,
+    },
+    types::{Blob, BlobProviderError, IndexedBlobHash, StageError, StageResult},
 };
 use alloc::{boxed::Box, sync::Arc, vec, vec::Vec};
 use alloy_consensus::{Header, Receipt, TxEnvelope};
@@ -72,6 +74,8 @@ pub struct TestChainProvider {
     pub headers: Vec<(B256, Header)>,
     /// Maps block hashes to receipts using a tuple list.
     pub receipts: Vec<(B256, Vec<Receipt>)>,
+    /// Maps block hashes to transactions using a tuple list.
+    pub transactions: Vec<(B256, Vec<TxEnvelope>)>,
 }
 
 impl TestChainProvider {
@@ -80,6 +84,17 @@ impl TestChainProvider {
         self.blocks.push((number, block));
     }
 
+    /// Insert a block with transactions into the mock chain provider.
+    pub fn insert_block_with_transactions(
+        &mut self,
+        number: u64,
+        block: BlockInfo,
+        txs: Vec<TxEnvelope>,
+    ) {
+        self.blocks.push((number, block));
+        self.transactions.push((block.hash, txs));
+    }
+
     /// Insert receipts into the mock chain provider.
     pub fn insert_receipts(&mut self, hash: B256, receipts: Vec<Receipt>) {
         self.receipts.push((hash, receipts));
@@ -149,7 +164,49 @@ impl ChainProvider for TestChainProvider {
             .find(|(_, b)| b.hash == hash)
             .map(|(_, b)| *b)
             .ok_or_else(|| anyhow::anyhow!("Block not found"))?;
-        Ok((block, Vec::new()))
+        let txs = self
+            .transactions
+            .iter()
+            .find(|(h, _)| *h == hash)
+            .map(|(_, txs)| txs.clone())
+            .unwrap_or_default();
+        Ok((block, txs))
+    }
+}
+
+/// A mock blob provider for testing.
+#[derive(Debug, Clone, Default)]
+pub struct TestBlobProvider {
+    /// Maps block hashes to blob data.
+    pub blobs: HashMap<B256, Blob>,
+}
+
+impl TestBlobProvider {
+    /// Insert a blob into the mock blob provider.
+    pub fn insert_blob(&mut self, hash: B256, blob: Blob) {
+        self.blobs.insert(hash, blob);
+    }
+
+    /// Clears blobs from the mock blob provider.
+    pub fn clear(&mut self) {
+        self.blobs.clear();
+    }
+}
+
+#[async_trait]
+impl BlobProvider for TestBlobProvider {
+    async fn get_blobs(
+        &mut self,
+        _block_ref: &BlockInfo,
+        blob_hashes: &[IndexedBlobHash],
+    ) -> Result<Vec<Blob>, BlobProviderError> {
+        let mut blobs = Vec::new();
+        for blob_hash in blob_hashes {
+            if let Some(data) = self.blobs.get(&blob_hash.hash) {
+                blobs.push(*data);
+            }
+        }
+        Ok(blobs)
     }
 }