prune: add unit tests for PruneInput get_next_tx_num_range (#12081)

crates/prune/prune/src/segments/mod.rs

@@ -23,9 +23,9 @@ pub use user::{
 
 /// A segment represents a pruning of some portion of the data.
 ///
-/// Segments are called from [Pruner](crate::Pruner) with the following lifecycle:
+/// Segments are called from [`Pruner`](crate::Pruner) with the following lifecycle:
 /// 1. Call [`Segment::prune`] with `delete_limit` of [`PruneInput`].
-/// 2. If [`Segment::prune`] returned a [Some] in `checkpoint` of [`SegmentOutput`], call
+/// 2. If [`Segment::prune`] returned a [`Some`] in `checkpoint` of [`SegmentOutput`], call
 ///    [`Segment::save_checkpoint`].
 /// 3. Subtract `pruned` of [`SegmentOutput`] from `delete_limit` of next [`PruneInput`].
 pub trait Segment<Provider>: Debug + Send + Sync {
@@ -88,7 +88,7 @@ impl PruneInput {
                 },
             })
             // No checkpoint exists, prune from genesis
-            .unwrap_or(0);
+            .unwrap_or_default();
 
         let to_tx_number = match provider.block_body_indices(self.to_block)? {
             Some(body) => {
@@ -143,3 +143,206 @@ impl PruneInput {
             .unwrap_or(0)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use alloy_primitives::B256;
+    use reth_provider::{
+        providers::BlockchainProvider2,
+        test_utils::{create_test_provider_factory, MockEthProvider},
+    };
+    use reth_testing_utils::generators::{self, random_block_range, BlockRangeParams};
+
+    #[test]
+    fn test_prune_input_get_next_tx_num_range_no_to_block() {
+        let input = PruneInput {
+            previous_checkpoint: None,
+            to_block: 10,
+            limiter: PruneLimiter::default(),
+        };
+
+        // Default provider with no block corresponding to block 10
+        let provider = MockEthProvider::default();
+
+        // No block body for block 10, expected None
+        let range = input.get_next_tx_num_range(&provider).expect("Expected range");
+        assert!(range.is_none());
+    }
+
+    #[test]
+    fn test_prune_input_get_next_tx_num_range_no_tx() {
+        let input = PruneInput {
+            previous_checkpoint: None,
+            to_block: 10,
+            limiter: PruneLimiter::default(),
+        };
+
+        let mut rng = generators::rng();
+        let factory = create_test_provider_factory();
+
+        // Generate 10 random blocks with no transactions
+        let blocks = random_block_range(
+            &mut rng,
+            0..=10,
+            BlockRangeParams { parent: Some(B256::ZERO), tx_count: 0..1, ..Default::default() },
+        );
+
+        // Insert the blocks into the database
+        let provider_rw = factory.provider_rw().expect("failed to get provider_rw");
+        for block in &blocks {
+            provider_rw
+                .insert_historical_block(
+                    block.clone().seal_with_senders().expect("failed to seal block with senders"),
+                )
+                .expect("failed to insert block");
+        }
+        provider_rw.commit().expect("failed to commit");
+
+        // Create a new provider
+        let provider = BlockchainProvider2::new(factory).unwrap();
+
+        // Since there are no transactions, expected None
+        let range = input.get_next_tx_num_range(&provider).expect("Expected range");
+        assert!(range.is_none());
+    }
+
+    #[test]
+    fn test_prune_input_get_next_tx_num_range_valid() {
+        // Create a new prune input
+        let input = PruneInput {
+            previous_checkpoint: None,
+            to_block: 10,
+            limiter: PruneLimiter::default(),
+        };
+
+        let mut rng = generators::rng();
+        let factory = create_test_provider_factory();
+
+        // Generate 10 random blocks with some transactions
+        let blocks = random_block_range(
+            &mut rng,
+            0..=10,
+            BlockRangeParams { parent: Some(B256::ZERO), tx_count: 0..5, ..Default::default() },
+        );
+
+        // Insert the blocks into the database
+        let provider_rw = factory.provider_rw().expect("failed to get provider_rw");
+        for block in &blocks {
+            provider_rw
+                .insert_historical_block(
+                    block.clone().seal_with_senders().expect("failed to seal block with senders"),
+                )
+                .expect("failed to insert block");
+        }
+        provider_rw.commit().expect("failed to commit");
+
+        // Create a new provider
+        let provider = BlockchainProvider2::new(factory).unwrap();
+
+        // Get the next tx number range
+        let range = input.get_next_tx_num_range(&provider).expect("Expected range").unwrap();
+
+        // Calculate the total number of transactions
+        let num_txs =
+            blocks.iter().map(|block| block.body.transactions().count() as u64).sum::<u64>();
+
+        assert_eq!(range, 0..=num_txs - 1);
+    }
+
+    #[test]
+    fn test_prune_input_get_next_tx_checkpoint_without_tx_number() {
+        // Create a prune input with a previous checkpoint without a tx number (unexpected)
+        let input = PruneInput {
+            previous_checkpoint: Some(PruneCheckpoint {
+                block_number: Some(5),
+                tx_number: None,
+                prune_mode: PruneMode::Full,
+            }),
+            to_block: 10,
+            limiter: PruneLimiter::default(),
+        };
+
+        let mut rng = generators::rng();
+        let factory = create_test_provider_factory();
+
+        // Generate 10 random blocks
+        let blocks = random_block_range(
+            &mut rng,
+            0..=10,
+            BlockRangeParams { parent: Some(B256::ZERO), tx_count: 0..5, ..Default::default() },
+        );
+
+        // Insert the blocks into the database
+        let provider_rw = factory.provider_rw().expect("failed to get provider_rw");
+        for block in &blocks {
+            provider_rw
+                .insert_historical_block(
+                    block.clone().seal_with_senders().expect("failed to seal block with senders"),
+                )
+                .expect("failed to insert block");
+        }
+        provider_rw.commit().expect("failed to commit");
+
+        // Create a new provider
+        let provider = BlockchainProvider2::new(factory).unwrap();
+
+        // Fetch the range and check if it is correct
+        let range = input.get_next_tx_num_range(&provider).expect("Expected range").unwrap();
+
+        // Calculate the total number of transactions
+        let num_txs =
+            blocks.iter().map(|block| block.body.transactions().count() as u64).sum::<u64>();
+
+        assert_eq!(range, 0..=num_txs - 1,);
+    }
+
+    #[test]
+    fn test_prune_input_get_next_tx_empty_range() {
+        // Create a new provider via factory
+        let mut rng = generators::rng();
+        let factory = create_test_provider_factory();
+
+        // Generate 10 random blocks
+        let blocks = random_block_range(
+            &mut rng,
+            0..=10,
+            BlockRangeParams { parent: Some(B256::ZERO), tx_count: 0..5, ..Default::default() },
+        );
+
+        // Insert the blocks into the database
+        let provider_rw = factory.provider_rw().expect("failed to get provider_rw");
+        for block in &blocks {
+            provider_rw
+                .insert_historical_block(
+                    block.clone().seal_with_senders().expect("failed to seal block with senders"),
+                )
+                .expect("failed to insert block");
+        }
+        provider_rw.commit().expect("failed to commit");
+
+        // Create a new provider
+        let provider = BlockchainProvider2::new(factory).unwrap();
+
+        // Get the last tx number
+        // Calculate the total number of transactions
+        let num_txs =
+            blocks.iter().map(|block| block.body.transactions().count() as u64).sum::<u64>();
+        let max_range = num_txs - 1;
+
+        // Create a prune input with a previous checkpoint that is the last tx number
+        let input = PruneInput {
+            previous_checkpoint: Some(PruneCheckpoint {
+                block_number: Some(5),
+                tx_number: Some(max_range),
+                prune_mode: PruneMode::Full,
+            }),
+            to_block: 10,
+            limiter: PruneLimiter::default(),
+        };
+
+        // We expect an empty range since the previous checkpoint is the last tx number
+        let range = input.get_next_tx_num_range(&provider).expect("Expected range");
+        assert!(range.is_none());
+    }
+}

crates/prune/prune/src/segments/set.rs

@@ -11,7 +11,7 @@ use reth_prune_types::PruneModes;
 
 use super::{StaticFileHeaders, StaticFileReceipts, StaticFileTransactions};
 
-/// Collection of [Segment]. Thread-safe, allocated on the heap.
+/// Collection of [`Segment`]. Thread-safe, allocated on the heap.
 #[derive(Debug)]
 pub struct SegmentSet<Provider> {
     inner: Vec<Box<dyn Segment<Provider>>>,
@@ -23,7 +23,7 @@ impl<Provider> SegmentSet<Provider> {
         Self::default()
     }
 
-    /// Adds new [Segment] to collection.
+    /// Adds new [`Segment`] to collection.
     pub fn segment<S: Segment<Provider> + 'static>(mut self, segment: S) -> Self {
         self.inner.push(Box::new(segment));
         self