Add tests for QueuedBlocks (#1268)

* Add unit test for QueuedBlocks
* Add test for pruned blocks

zebra-state/src/service/memory_state/queued_blocks.rs

@@ -96,10 +96,27 @@ impl QueuedBlocks {
         for hash in by_height.into_iter().flat_map(|(_, hashes)| hashes) {
             let expired = self.blocks.remove(&hash).expect("block is present");
             let parent_hash = &expired.block.header.previous_block_hash;
-            self.by_parent
+
+            let parent_list = self
+                .by_parent
                 .get_mut(parent_hash)
-                .expect("parent is present")
-                .remove(&hash);
+                .expect("parent is present");
+
+            if parent_list.len() == 1 {
+                let removed = self
+                    .by_parent
+                    .remove(parent_hash)
+                    .expect("parent is present");
+                assert!(
+                    removed.contains(&hash),
+                    "hash must be present in parent hash list"
+                );
+            } else {
+                assert!(
+                    parent_list.remove(&hash),
+                    "hash must be present in parent hash list"
+                );
+            }
         }
 
         tracing::trace!(num_blocks = %self.blocks.len(), "Finished pruning blocks at or beneath the finalized tip height",);
@@ -122,3 +139,136 @@ impl QueuedBlocks {
         metrics::gauge!("state.memory.queued.block.count", self.blocks.len() as _);
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use tokio::sync::oneshot;
+    use zebra_chain::{block::Block, serialization::ZcashDeserializeInto};
+    use zebra_test::prelude::*;
+
+    use crate::tests::FakeChainHelper;
+
+    use self::assert_eq;
+    use super::*;
+
+    // Quick helper trait for making queued blocks with throw away channels
+    trait IntoQueued {
+        fn into_queued(self) -> QueuedBlock;
+    }
+
+    impl IntoQueued for Arc<Block> {
+        fn into_queued(self) -> QueuedBlock {
+            let (rsp_tx, _) = oneshot::channel();
+
+            QueuedBlock {
+                block: self,
+                rsp_tx,
+            }
+        }
+    }
+
+    #[test]
+    fn dequeue_gives_right_children() -> Result<()> {
+        zebra_test::init();
+
+        let block1: Arc<Block> =
+            zebra_test::vectors::BLOCK_MAINNET_419200_BYTES.zcash_deserialize_into()?;
+        let child1: Arc<Block> =
+            zebra_test::vectors::BLOCK_MAINNET_419201_BYTES.zcash_deserialize_into()?;
+        let child2 = block1.make_fake_child();
+
+        let parent = block1.header.previous_block_hash;
+
+        let mut queue = QueuedBlocks::default();
+        // Empty to start
+        assert_eq!(0, queue.blocks.len());
+        assert_eq!(0, queue.by_parent.len());
+        assert_eq!(0, queue.by_height.len());
+
+        // Inserting the first block gives us 1 in each table
+        queue.queue(block1.clone().into_queued());
+        assert_eq!(1, queue.blocks.len());
+        assert_eq!(1, queue.by_parent.len());
+        assert_eq!(1, queue.by_height.len());
+
+        // The second gives us one in each table because its a child of the first
+        queue.queue(child1.clone().into_queued());
+        assert_eq!(2, queue.blocks.len());
+        assert_eq!(2, queue.by_parent.len());
+        assert_eq!(2, queue.by_height.len());
+
+        // The 3rd only increments blocks, because it is also a child of the
+        // first block, so for the second and third tables it gets added to the
+        // existing HashSet value
+        queue.queue(child2.clone().into_queued());
+        assert_eq!(3, queue.blocks.len());
+        assert_eq!(2, queue.by_parent.len());
+        assert_eq!(2, queue.by_height.len());
+
+        // Dequeueing the first block removes 1 block from each list
+        let children = queue.dequeue_children(parent);
+        assert_eq!(1, children.len());
+        assert_eq!(block1, children[0].block);
+        assert_eq!(2, queue.blocks.len());
+        assert_eq!(1, queue.by_parent.len());
+        assert_eq!(1, queue.by_height.len());
+
+        // Dequeueing the children of the first block removes both of the other
+        // blocks, and empties all lists
+        let parent = children[0].block.hash();
+        let children = queue.dequeue_children(parent);
+        assert_eq!(2, children.len());
+        assert!(children
+            .iter()
+            .any(|QueuedBlock { block, .. }| block == &child1));
+        assert!(children
+            .iter()
+            .any(|QueuedBlock { block, .. }| block == &child2));
+        assert_eq!(0, queue.blocks.len());
+        assert_eq!(0, queue.by_parent.len());
+        assert_eq!(0, queue.by_height.len());
+
+        Ok(())
+    }
+
+    #[test]
+    fn prune_removes_right_children() -> Result<()> {
+        zebra_test::init();
+
+        let block1: Arc<Block> =
+            zebra_test::vectors::BLOCK_MAINNET_419200_BYTES.zcash_deserialize_into()?;
+        let child1: Arc<Block> =
+            zebra_test::vectors::BLOCK_MAINNET_419201_BYTES.zcash_deserialize_into()?;
+        let child2 = block1.make_fake_child();
+
+        let mut queue = QueuedBlocks::default();
+        queue.queue(block1.clone().into_queued());
+        queue.queue(child1.clone().into_queued());
+        queue.queue(child2.clone().into_queued());
+        assert_eq!(3, queue.blocks.len());
+        assert_eq!(2, queue.by_parent.len());
+        assert_eq!(2, queue.by_height.len());
+
+        // Pruning the first height removes only block1
+        queue.prune_by_height(block1.coinbase_height().unwrap());
+        assert_eq!(2, queue.blocks.len());
+        assert_eq!(1, queue.by_parent.len());
+        assert_eq!(1, queue.by_height.len());
+        assert!(queue.get_mut(&block1.hash()).is_none());
+        assert!(queue.get_mut(&child1.hash()).is_some());
+        assert!(queue.get_mut(&child2.hash()).is_some());
+
+        // Dequeueing the children of the first block removes both of the other
+        // blocks, and empties all lists
+        queue.prune_by_height(child1.coinbase_height().unwrap());
+        assert_eq!(0, queue.blocks.len());
+        assert_eq!(0, queue.by_parent.len());
+        assert_eq!(0, queue.by_height.len());
+        assert!(queue.get_mut(&child1.hash()).is_none());
+        assert!(queue.get_mut(&child2.hash()).is_none());
+
+        Ok(())
+    }
+}