blockchain-tree: add unit tests for BlockIndices (#10675)

crates/blockchain-tree/src/block_indices.rs

@@ -371,3 +371,246 @@ impl BlockIndices {
         &self.canonical_chain
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use reth_primitives::{Header, SealedBlock, SealedHeader, B256};
+
+    #[test]
+    fn pending_block_num_hash_returns_none_if_no_fork() {
+        // Create a new canonical chain with a single block (represented by its number and hash).
+        let canonical_chain = BTreeMap::from([(0, B256::from_slice(&[1; 32]))]);
+
+        let block_indices = BlockIndices::new(0, canonical_chain);
+
+        // No fork to child blocks, so there is no pending block.
+        assert_eq!(block_indices.pending_block_num_hash(), None);
+    }
+
+    #[test]
+    fn pending_block_num_hash_works() {
+        // Create a canonical chain with multiple blocks at heights 1, 2, and 3.
+        let canonical_chain = BTreeMap::from([
+            (1, B256::from_slice(&[1; 32])),
+            (2, B256::from_slice(&[2; 32])),
+            (3, B256::from_slice(&[3; 32])),
+        ]);
+
+        let mut block_indices = BlockIndices::new(3, canonical_chain);
+
+        // Define the hash of the parent block (the block at height 3 in the canonical chain).
+        let parent_hash = B256::from_slice(&[3; 32]);
+
+        // Define the hashes of two child blocks that extend the canonical chain.
+        let child_hash_1 = B256::from_slice(&[2; 32]);
+        let child_hash_2 = B256::from_slice(&[3; 32]);
+
+        // Create a set to store both child block hashes.
+        let mut child_set = LinkedHashSet::new();
+        child_set.insert(child_hash_1);
+        child_set.insert(child_hash_2);
+
+        // Associate the parent block hash with its children in the fork_to_child mapping.
+        block_indices.fork_to_child.insert(parent_hash, child_set);
+
+        // Pending block should be the first child block.
+        assert_eq!(
+            block_indices.pending_block_num_hash(),
+            Some(BlockNumHash { number: 4, hash: child_hash_1 })
+        );
+    }
+
+    #[test]
+    fn pending_blocks_returns_empty_if_no_fork() {
+        // Create a canonical chain with a single block at height 10.
+        let canonical_chain = BTreeMap::from([(10, B256::from_slice(&[1; 32]))]);
+        let block_indices = BlockIndices::new(0, canonical_chain);
+
+        // No child blocks are associated with the canonical tip.
+        assert_eq!(block_indices.pending_blocks(), (11, Vec::new()));
+    }
+
+    #[test]
+    fn pending_blocks_returns_multiple_children() {
+        // Define the hash of the parent block (the block at height 5 in the canonical chain).
+        let parent_hash = B256::from_slice(&[3; 32]);
+
+        // Create a canonical chain with a block at height 5.
+        let canonical_chain = BTreeMap::from([(5, parent_hash)]);
+        let mut block_indices = BlockIndices::new(0, canonical_chain);
+
+        // Define the hashes of two child blocks.
+        let child_hash_1 = B256::from_slice(&[4; 32]);
+        let child_hash_2 = B256::from_slice(&[5; 32]);
+
+        // Create a set to store both child block hashes.
+        let mut child_set = LinkedHashSet::new();
+        child_set.insert(child_hash_1);
+        child_set.insert(child_hash_2);
+
+        // Associate the parent block hash with its children.
+        block_indices.fork_to_child.insert(parent_hash, child_set);
+
+        // Pending blocks should be the two child blocks.
+        assert_eq!(block_indices.pending_blocks(), (6, vec![child_hash_1, child_hash_2]));
+    }
+
+    #[test]
+    fn pending_blocks_with_multiple_forked_chains() {
+        // Define hashes for parent blocks and child blocks.
+        let parent_hash_1 = B256::from_slice(&[6; 32]);
+        let parent_hash_2 = B256::from_slice(&[7; 32]);
+
+        // Create a canonical chain with blocks at heights 1 and 2.
+        let canonical_chain = BTreeMap::from([(1, parent_hash_1), (2, parent_hash_2)]);
+
+        let mut block_indices = BlockIndices::new(2, canonical_chain);
+
+        // Define hashes for child blocks.
+        let child_hash_1 = B256::from_slice(&[8; 32]);
+        let child_hash_2 = B256::from_slice(&[9; 32]);
+
+        // Create sets to store child blocks for each parent block.
+        let mut child_set_1 = LinkedHashSet::new();
+        let mut child_set_2 = LinkedHashSet::new();
+        child_set_1.insert(child_hash_1);
+        child_set_2.insert(child_hash_2);
+
+        // Associate parent block hashes with their child blocks.
+        block_indices.fork_to_child.insert(parent_hash_1, child_set_1);
+        block_indices.fork_to_child.insert(parent_hash_2, child_set_2);
+
+        // Check that the pending blocks are only those extending the canonical tip.
+        assert_eq!(block_indices.pending_blocks(), (3, vec![child_hash_2]));
+    }
+
+    #[test]
+    fn insert_non_fork_block_adds_block_correctly() {
+        // Create a new BlockIndices instance with an empty state.
+        let mut block_indices = BlockIndices::new(0, BTreeMap::new());
+
+        // Define test parameters.
+        let block_number = 1;
+        let block_hash = B256::from_slice(&[1; 32]);
+        let chain_id = SidechainId::from(42);
+
+        // Insert the block into the BlockIndices instance.
+        block_indices.insert_non_fork_block(block_number, block_hash, chain_id);
+
+        // Check that the block number to block hashes mapping includes the new block hash.
+        assert_eq!(
+            block_indices.block_number_to_block_hashes.get(&block_number),
+            Some(&HashSet::from([block_hash]))
+        );
+
+        // Check that the block hash to chain ID mapping includes the new entry.
+        assert_eq!(block_indices.blocks_to_chain.get(&block_hash), Some(&chain_id));
+    }
+
+    #[test]
+    fn insert_non_fork_block_combined_tests() {
+        // Create a new BlockIndices instance with an empty state.
+        let mut block_indices = BlockIndices::new(0, BTreeMap::new());
+
+        // Define test parameters.
+        let block_number_1 = 2;
+        let block_hash_1 = B256::from_slice(&[1; 32]);
+        let block_hash_2 = B256::from_slice(&[2; 32]);
+        let chain_id_1 = SidechainId::from(84);
+
+        let block_number_2 = 4;
+        let block_hash_3 = B256::from_slice(&[3; 32]);
+        let chain_id_2 = SidechainId::from(200);
+
+        // Insert multiple hashes for the same block number.
+        block_indices.insert_non_fork_block(block_number_1, block_hash_1, chain_id_1);
+        block_indices.insert_non_fork_block(block_number_1, block_hash_2, chain_id_1);
+
+        // Insert blocks with different numbers.
+        block_indices.insert_non_fork_block(block_number_2, block_hash_3, chain_id_2);
+
+        // Block number 1 should have two block hashes associated with it.
+        let mut expected_hashes_for_block_1 = HashSet::new();
+        expected_hashes_for_block_1.insert(block_hash_1);
+        expected_hashes_for_block_1.insert(block_hash_2);
+        assert_eq!(
+            block_indices.block_number_to_block_hashes.get(&block_number_1),
+            Some(&expected_hashes_for_block_1)
+        );
+
+        // Check that the block hashes for block_number_1 are associated with the correct chain ID.
+        assert_eq!(block_indices.blocks_to_chain.get(&block_hash_1), Some(&chain_id_1));
+        assert_eq!(block_indices.blocks_to_chain.get(&block_hash_2), Some(&chain_id_1));
+
+        // Block number 2 should have a single block hash associated with it.
+        assert_eq!(
+            block_indices.block_number_to_block_hashes.get(&block_number_2),
+            Some(&HashSet::from([block_hash_3]))
+        );
+
+        // Block hash 3 should be associated with the correct chain ID.
+        assert_eq!(block_indices.blocks_to_chain.get(&block_hash_3), Some(&chain_id_2));
+    }
+
+    #[test]
+    fn insert_chain_validates_insertion() {
+        // Create a new BlockIndices instance with an empty state.
+        let mut block_indices = BlockIndices::new(0, BTreeMap::new());
+
+        // Define test parameters.
+        let chain_id = SidechainId::from(42);
+
+        // Define some example blocks and their hashes.
+        let block_hash_1 = B256::from_slice(&[1; 32]);
+        let block_hash_2 = B256::from_slice(&[2; 32]);
+        let parent_hash = B256::from_slice(&[0; 32]);
+
+        // Define blocks with their numbers and parent hashes.
+        let block_1 = SealedBlockWithSenders {
+            block: SealedBlock {
+                header: SealedHeader::new(
+                    Header { parent_hash, number: 1, ..Default::default() },
+                    block_hash_1,
+                ),
+                ..Default::default()
+            },
+            ..Default::default()
+        };
+        let block_2 = SealedBlockWithSenders {
+            block: SealedBlock {
+                header: SealedHeader::new(
+                    Header { parent_hash: block_hash_1, number: 2, ..Default::default() },
+                    block_hash_2,
+                ),
+                ..Default::default()
+            },
+            ..Default::default()
+        };
+
+        // Define a chain containing the blocks.
+        let chain = Chain::new(vec![block_1, block_2], Default::default(), Default::default());
+
+        // Insert the chain into the BlockIndices.
+        block_indices.insert_chain(chain_id, &chain);
+
+        // Check that the blocks are correctly mapped to the chain ID.
+        assert_eq!(block_indices.blocks_to_chain.get(&block_hash_1), Some(&chain_id));
+        assert_eq!(block_indices.blocks_to_chain.get(&block_hash_2), Some(&chain_id));
+
+        // Check that block numbers map to their respective hashes.
+        let mut expected_hashes_1 = HashSet::new();
+        expected_hashes_1.insert(block_hash_1);
+        assert_eq!(block_indices.block_number_to_block_hashes.get(&1), Some(&expected_hashes_1));
+
+        let mut expected_hashes_2 = HashSet::new();
+        expected_hashes_2.insert(block_hash_2);
+        assert_eq!(block_indices.block_number_to_block_hashes.get(&2), Some(&expected_hashes_2));
+
+        // Check that the fork_to_child mapping contains the correct parent-child relationship.
+        // We take the first block of the chain.
+        let mut expected_children = LinkedHashSet::new();
+        expected_children.insert(block_hash_1);
+        assert_eq!(block_indices.fork_to_child.get(&parent_hash), Some(&expected_children));
+    }
+}