diff --git a/zebra-state/src/service/non_finalized_state.rs b/zebra-state/src/service/non_finalized_state.rs
index adbb4348..021a1e66 100644
--- a/zebra-state/src/service/non_finalized_state.rs
+++ b/zebra-state/src/service/non_finalized_state.rs
@@ -5,6 +5,11 @@
 mod chain;
 mod queued_blocks;
 
+#[cfg(any(test, feature = "proptest-impl"))]
+mod arbitrary;
+#[cfg(test)]
+mod tests;
+
 pub use queued_blocks::QueuedBlocks;
 
 use std::{collections::BTreeSet, mem, ops::Deref, sync::Arc};
@@ -271,242 +276,3 @@ impl NonFinalizedState {
         metrics::gauge!("state.memory.best.chain.length", self.best_chain_len() as _);
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use zebra_chain::serialization::ZcashDeserializeInto;
-    use zebra_test::prelude::*;
-
-    use crate::tests::{FakeChainHelper, Prepare};
-
-    use self::assert_eq;
-    use super::*;
-
-    #[test]
-    fn best_chain_wins() -> Result<()> {
-        zebra_test::init();
-
-        best_chain_wins_for_network(Network::Mainnet)?;
-        best_chain_wins_for_network(Network::Testnet)?;
-
-        Ok(())
-    }
-
-    fn best_chain_wins_for_network(network: Network) -> Result<()> {
-        let block1: Arc<Block> = match network {
-            Network::Mainnet => {
-                zebra_test::vectors::BLOCK_MAINNET_1180900_BYTES.zcash_deserialize_into()?
-            }
-            Network::Testnet => {
-                zebra_test::vectors::BLOCK_TESTNET_1326100_BYTES.zcash_deserialize_into()?
-            }
-        };
-
-        let block2 = block1.make_fake_child().set_work(10);
-        let child = block1.make_fake_child().set_work(1);
-
-        let expected_hash = block2.hash();
-
-        let mut state = NonFinalizedState::default();
-        state.commit_new_chain(block2.prepare());
-        state.commit_new_chain(child.prepare());
-
-        let best_chain = state.best_chain().unwrap();
-        assert!(best_chain.height_by_hash.contains_key(&expected_hash));
-
-        Ok(())
-    }
-
-    #[test]
-    fn finalize_pops_from_best_chain() -> Result<()> {
-        zebra_test::init();
-
-        finalize_pops_from_best_chain_for_network(Network::Mainnet)?;
-        finalize_pops_from_best_chain_for_network(Network::Testnet)?;
-
-        Ok(())
-    }
-
-    fn finalize_pops_from_best_chain_for_network(network: Network) -> Result<()> {
-        let block1: Arc<Block> = match network {
-            Network::Mainnet => {
-                zebra_test::vectors::BLOCK_MAINNET_1180900_BYTES.zcash_deserialize_into()?
-            }
-            Network::Testnet => {
-                zebra_test::vectors::BLOCK_TESTNET_1326100_BYTES.zcash_deserialize_into()?
-            }
-        };
-
-        let block2 = block1.make_fake_child().set_work(10);
-        let child = block1.make_fake_child().set_work(1);
-
-        let mut state = NonFinalizedState::default();
-        state.commit_new_chain(block1.clone().prepare());
-        state.commit_block(block2.clone().prepare());
-        state.commit_block(child.prepare());
-
-        let finalized = state.finalize();
-        assert_eq!(block1, finalized.block);
-
-        let finalized = state.finalize();
-        assert_eq!(block2, finalized.block);
-
-        assert!(state.best_chain().is_none());
-
-        Ok(())
-    }
-
-    #[test]
-    // This test gives full coverage for `take_chain_if`
-    fn commit_block_extending_best_chain_doesnt_drop_worst_chains() -> Result<()> {
-        zebra_test::init();
-
-        commit_block_extending_best_chain_doesnt_drop_worst_chains_for_network(Network::Mainnet)?;
-        commit_block_extending_best_chain_doesnt_drop_worst_chains_for_network(Network::Testnet)?;
-
-        Ok(())
-    }
-
-    fn commit_block_extending_best_chain_doesnt_drop_worst_chains_for_network(
-        network: Network,
-    ) -> Result<()> {
-        let block1: Arc<Block> = match network {
-            Network::Mainnet => {
-                zebra_test::vectors::BLOCK_MAINNET_1180900_BYTES.zcash_deserialize_into()?
-            }
-            Network::Testnet => {
-                zebra_test::vectors::BLOCK_TESTNET_1326100_BYTES.zcash_deserialize_into()?
-            }
-        };
-
-        let block2 = block1.make_fake_child().set_work(10);
-        let child1 = block1.make_fake_child().set_work(1);
-        let child2 = block2.make_fake_child().set_work(1);
-
-        let mut state = NonFinalizedState::default();
-        assert_eq!(0, state.chain_set.len());
-        state.commit_new_chain(block1.prepare());
-        assert_eq!(1, state.chain_set.len());
-        state.commit_block(block2.prepare());
-        assert_eq!(1, state.chain_set.len());
-        state.commit_block(child1.prepare());
-        assert_eq!(2, state.chain_set.len());
-        state.commit_block(child2.prepare());
-        assert_eq!(2, state.chain_set.len());
-
-        Ok(())
-    }
-
-    #[test]
-    fn shorter_chain_can_be_best_chain() -> Result<()> {
-        zebra_test::init();
-
-        shorter_chain_can_be_best_chain_for_network(Network::Mainnet)?;
-        shorter_chain_can_be_best_chain_for_network(Network::Testnet)?;
-
-        Ok(())
-    }
-
-    fn shorter_chain_can_be_best_chain_for_network(network: Network) -> Result<()> {
-        let block1: Arc<Block> = match network {
-            Network::Mainnet => {
-                zebra_test::vectors::BLOCK_MAINNET_1180900_BYTES.zcash_deserialize_into()?
-            }
-            Network::Testnet => {
-                zebra_test::vectors::BLOCK_TESTNET_1326100_BYTES.zcash_deserialize_into()?
-            }
-        };
-
-        let long_chain_block1 = block1.make_fake_child().set_work(1);
-        let long_chain_block2 = long_chain_block1.make_fake_child().set_work(1);
-
-        let short_chain_block = block1.make_fake_child().set_work(3);
-
-        let mut state = NonFinalizedState::default();
-        state.commit_new_chain(block1.prepare());
-        state.commit_block(long_chain_block1.prepare());
-        state.commit_block(long_chain_block2.prepare());
-        state.commit_block(short_chain_block.prepare());
-        assert_eq!(2, state.chain_set.len());
-
-        assert_eq!(2, state.best_chain_len());
-
-        Ok(())
-    }
-
-    #[test]
-    fn longer_chain_with_more_work_wins() -> Result<()> {
-        zebra_test::init();
-
-        longer_chain_with_more_work_wins_for_network(Network::Mainnet)?;
-        longer_chain_with_more_work_wins_for_network(Network::Testnet)?;
-
-        Ok(())
-    }
-
-    fn longer_chain_with_more_work_wins_for_network(network: Network) -> Result<()> {
-        let block1: Arc<Block> = match network {
-            Network::Mainnet => {
-                zebra_test::vectors::BLOCK_MAINNET_1180900_BYTES.zcash_deserialize_into()?
-            }
-            Network::Testnet => {
-                zebra_test::vectors::BLOCK_TESTNET_1326100_BYTES.zcash_deserialize_into()?
-            }
-        };
-
-        let long_chain_block1 = block1.make_fake_child().set_work(1);
-        let long_chain_block2 = long_chain_block1.make_fake_child().set_work(1);
-        let long_chain_block3 = long_chain_block2.make_fake_child().set_work(1);
-        let long_chain_block4 = long_chain_block3.make_fake_child().set_work(1);
-
-        let short_chain_block = block1.make_fake_child().set_work(3);
-
-        let mut state = NonFinalizedState::default();
-        state.commit_new_chain(block1.prepare());
-        state.commit_block(long_chain_block1.prepare());
-        state.commit_block(long_chain_block2.prepare());
-        state.commit_block(long_chain_block3.prepare());
-        state.commit_block(long_chain_block4.prepare());
-        state.commit_block(short_chain_block.prepare());
-        assert_eq!(2, state.chain_set.len());
-
-        assert_eq!(5, state.best_chain_len());
-
-        Ok(())
-    }
-
-    #[test]
-    fn equal_length_goes_to_more_work() -> Result<()> {
-        zebra_test::init();
-
-        equal_length_goes_to_more_work_for_network(Network::Mainnet)?;
-        equal_length_goes_to_more_work_for_network(Network::Testnet)?;
-
-        Ok(())
-    }
-    fn equal_length_goes_to_more_work_for_network(network: Network) -> Result<()> {
-        let block1: Arc<Block> = match network {
-            Network::Mainnet => {
-                zebra_test::vectors::BLOCK_MAINNET_1180900_BYTES.zcash_deserialize_into()?
-            }
-            Network::Testnet => {
-                zebra_test::vectors::BLOCK_TESTNET_1326100_BYTES.zcash_deserialize_into()?
-            }
-        };
-
-        let less_work_child = block1.make_fake_child().set_work(1);
-        let more_work_child = block1.make_fake_child().set_work(3);
-        let expected_hash = more_work_child.hash();
-
-        let mut state = NonFinalizedState::default();
-        state.commit_new_chain(block1.prepare());
-        state.commit_block(less_work_child.prepare());
-        state.commit_block(more_work_child.prepare());
-        assert_eq!(2, state.chain_set.len());
-
-        let tip_hash = state.best_tip().unwrap().1;
-        assert_eq!(expected_hash, tip_hash);
-
-        Ok(())
-    }
-}
diff --git a/zebra-state/src/service/non_finalized_state/arbitrary.rs b/zebra-state/src/service/non_finalized_state/arbitrary.rs
new file mode 100644
index 00000000..61ae4f66
--- /dev/null
+++ b/zebra-state/src/service/non_finalized_state/arbitrary.rs
@@ -0,0 +1,64 @@
+use proptest::{
+    num::usize::BinarySearch,
+    strategy::{NewTree, ValueTree},
+    test_runner::TestRunner,
+};
+use std::sync::Arc;
+
+use zebra_chain::{block::Block, LedgerState};
+use zebra_test::prelude::*;
+
+use crate::tests::Prepare;
+
+use super::*;
+
+const MAX_PARTIAL_CHAIN_BLOCKS: usize = 100;
+
+#[derive(Debug)]
+pub struct PreparedChainTree {
+    chain: Arc<Vec<PreparedBlock>>,
+    count: BinarySearch,
+}
+
+impl ValueTree for PreparedChainTree {
+    type Value = (Arc<Vec<PreparedBlock>>, <BinarySearch as ValueTree>::Value);
+
+    fn current(&self) -> Self::Value {
+        (self.chain.clone(), self.count.current())
+    }
+
+    fn simplify(&mut self) -> bool {
+        self.count.simplify()
+    }
+
+    fn complicate(&mut self) -> bool {
+        self.count.complicate()
+    }
+}
+
+#[derive(Debug, Default)]
+pub struct PreparedChain {
+    // the proptests are threaded (not async), so we want to use a threaded mutex here
+    chain: std::sync::Mutex<Option<Arc<Vec<PreparedBlock>>>>,
+}
+
+impl Strategy for PreparedChain {
+    type Tree = PreparedChainTree;
+    type Value = <PreparedChainTree as ValueTree>::Value;
+
+    fn new_tree(&self, runner: &mut TestRunner) -> NewTree<Self> {
+        let mut chain = self.chain.lock().unwrap();
+        if chain.is_none() {
+            let blocks =
+                Block::partial_chain_strategy(LedgerState::default(), MAX_PARTIAL_CHAIN_BLOCKS)
+                    .prop_map(|vec| vec.into_iter().map(|blk| blk.prepare()).collect::<Vec<_>>())
+                    .new_tree(runner)?
+                    .current();
+            *chain = Some(Arc::new(blocks));
+        }
+
+        let chain = chain.clone().expect("should be generated");
+        let count = (1..chain.len()).new_tree(runner)?;
+        Ok(PreparedChainTree { chain, count })
+    }
+}
diff --git a/zebra-state/src/service/non_finalized_state/chain.rs b/zebra-state/src/service/non_finalized_state/chain.rs
index 132ad2c4..85f7a0da 100644
--- a/zebra-state/src/service/non_finalized_state/chain.rs
+++ b/zebra-state/src/service/non_finalized_state/chain.rs
@@ -403,203 +403,3 @@ impl Ord for Chain {
         }
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use proptest::{
-        num::usize::BinarySearch,
-        strategy::{NewTree, ValueTree},
-        test_runner::TestRunner,
-    };
-    use std::{env, sync::Arc};
-
-    use zebra_chain::{
-        block::Block,
-        parameters::{Network, NetworkUpgrade},
-        serialization::ZcashDeserializeInto,
-        LedgerState,
-    };
-    use zebra_test::prelude::*;
-
-    use crate::tests::{FakeChainHelper, Prepare};
-
-    use self::assert_eq;
-    use super::*;
-
-    const MAX_PARTIAL_CHAIN_BLOCKS: usize = 100;
-    const DEFAULT_PARTIAL_CHAIN_PROPTEST_CASES: u32 = 32;
-
-    #[test]
-    fn construct_empty() {
-        zebra_test::init();
-        let _chain = Chain::default();
-    }
-
-    #[test]
-    fn construct_single() -> Result<()> {
-        zebra_test::init();
-        let block: Arc<Block> =
-            zebra_test::vectors::BLOCK_MAINNET_434873_BYTES.zcash_deserialize_into()?;
-
-        let mut chain = Chain::default();
-        chain.push(block.prepare());
-
-        assert_eq!(1, chain.blocks.len());
-
-        Ok(())
-    }
-
-    #[test]
-    fn construct_many() -> Result<()> {
-        zebra_test::init();
-
-        let mut block: Arc<Block> =
-            zebra_test::vectors::BLOCK_MAINNET_434873_BYTES.zcash_deserialize_into()?;
-        let mut blocks = vec![];
-
-        while blocks.len() < 100 {
-            let next_block = block.make_fake_child();
-            blocks.push(block);
-            block = next_block;
-        }
-
-        let mut chain = Chain::default();
-
-        for block in blocks {
-            chain.push(block.prepare());
-        }
-
-        assert_eq!(100, chain.blocks.len());
-
-        Ok(())
-    }
-
-    #[test]
-    fn ord_matches_work() -> Result<()> {
-        zebra_test::init();
-        let less_block = zebra_test::vectors::BLOCK_MAINNET_434873_BYTES
-            .zcash_deserialize_into::<Arc<Block>>()?
-            .set_work(1);
-        let more_block = less_block.clone().set_work(10);
-
-        let mut lesser_chain = Chain::default();
-        lesser_chain.push(less_block.prepare());
-
-        let mut bigger_chain = Chain::default();
-        bigger_chain.push(more_block.prepare());
-
-        assert!(bigger_chain > lesser_chain);
-
-        Ok(())
-    }
-
-    #[derive(Debug)]
-    struct PreparedChainTree {
-        chain: Arc<Vec<PreparedBlock>>,
-        count: BinarySearch,
-    }
-
-    impl ValueTree for PreparedChainTree {
-        type Value = (Arc<Vec<PreparedBlock>>, <BinarySearch as ValueTree>::Value);
-
-        fn current(&self) -> Self::Value {
-            (self.chain.clone(), self.count.current())
-        }
-
-        fn simplify(&mut self) -> bool {
-            self.count.simplify()
-        }
-
-        fn complicate(&mut self) -> bool {
-            self.count.complicate()
-        }
-    }
-
-    #[derive(Debug, Default)]
-    struct PreparedChain {
-        // the proptests are threaded (not async), so we want to use a threaded mutex here
-        chain: std::sync::Mutex<Option<Arc<Vec<PreparedBlock>>>>,
-    }
-
-    impl Strategy for PreparedChain {
-        type Tree = PreparedChainTree;
-        type Value = <PreparedChainTree as ValueTree>::Value;
-
-        fn new_tree(&self, runner: &mut TestRunner) -> NewTree<Self> {
-            let mut chain = self.chain.lock().unwrap();
-            if chain.is_none() {
-                let tip_height = NetworkUpgrade::Canopy
-                    .activation_height(Network::Mainnet)
-                    .unwrap();
-                let ledger_state = LedgerState::new(tip_height, Network::Mainnet);
-                let blocks = Block::partial_chain_strategy(ledger_state, MAX_PARTIAL_CHAIN_BLOCKS)
-                    .prop_map(|vec| vec.into_iter().map(|blk| blk.prepare()).collect::<Vec<_>>())
-                    .new_tree(runner)?
-                    .current();
-                *chain = Some(Arc::new(blocks));
-            }
-
-            let chain = chain.clone().expect("should be generated");
-            let count = (1..chain.len()).new_tree(runner)?;
-            Ok(PreparedChainTree { chain, count })
-        }
-    }
-
-    #[test]
-    fn forked_equals_pushed() -> Result<()> {
-        zebra_test::init();
-
-        proptest!(ProptestConfig::with_cases(env::var("PROPTEST_CASES")
-                                          .ok()
-                                          .and_then(|v| v.parse().ok())
-                                          .unwrap_or(DEFAULT_PARTIAL_CHAIN_PROPTEST_CASES)),
-        |((chain, count) in PreparedChain::default())| {
-            let fork_tip_hash = chain[count - 1].hash;
-            let mut full_chain = Chain::default();
-            let mut partial_chain = Chain::default();
-
-            for block in chain.iter().take(count) {
-                partial_chain.push(block.clone());
-            }
-            for block in chain.iter() {
-                full_chain.push(block.clone());
-            }
-
-            let forked = full_chain.fork(fork_tip_hash).expect("hash is present");
-
-            prop_assert_eq!(forked.blocks.len(), partial_chain.blocks.len());
-        });
-
-        Ok(())
-    }
-
-    #[test]
-    fn finalized_equals_pushed() -> Result<()> {
-        zebra_test::init();
-
-        proptest!(ProptestConfig::with_cases(env::var("PROPTEST_CASES")
-                                          .ok()
-                                          .and_then(|v| v.parse().ok())
-                                          .unwrap_or(DEFAULT_PARTIAL_CHAIN_PROPTEST_CASES)),
-        |((chain, end_count) in PreparedChain::default())| {
-            let finalized_count = chain.len() - end_count;
-            let mut full_chain = Chain::default();
-            let mut partial_chain = Chain::default();
-
-            for block in chain.iter().skip(finalized_count) {
-                partial_chain.push(block.clone());
-            }
-            for block in chain.iter() {
-                full_chain.push(block.clone());
-            }
-
-            for _ in 0..finalized_count {
-                let _finalized = full_chain.pop_root();
-            }
-
-            prop_assert_eq!(full_chain.blocks.len(), partial_chain.blocks.len());
-        });
-
-        Ok(())
-    }
-}
diff --git a/zebra-state/src/service/non_finalized_state/tests.rs b/zebra-state/src/service/non_finalized_state/tests.rs
new file mode 100644
index 00000000..cc95d9d4
--- /dev/null
+++ b/zebra-state/src/service/non_finalized_state/tests.rs
@@ -0,0 +1,2 @@
+mod prop;
+mod vectors;
diff --git a/zebra-state/src/service/non_finalized_state/tests/prop.rs b/zebra-state/src/service/non_finalized_state/tests/prop.rs
new file mode 100644
index 00000000..d28bb68f
--- /dev/null
+++ b/zebra-state/src/service/non_finalized_state/tests/prop.rs
@@ -0,0 +1,65 @@
+use std::env;
+
+use zebra_test::prelude::*;
+
+use crate::service::non_finalized_state::{arbitrary::PreparedChain, Chain};
+
+const DEFAULT_PARTIAL_CHAIN_PROPTEST_CASES: u32 = 32;
+
+#[test]
+fn forked_equals_pushed() -> Result<()> {
+    zebra_test::init();
+
+    proptest!(ProptestConfig::with_cases(env::var("PROPTEST_CASES")
+                                          .ok()
+                                          .and_then(|v| v.parse().ok())
+                                          .unwrap_or(DEFAULT_PARTIAL_CHAIN_PROPTEST_CASES)),
+        |((chain, count) in PreparedChain::default())| {
+            let fork_tip_hash = chain[count - 1].hash;
+            let mut full_chain = Chain::default();
+            let mut partial_chain = Chain::default();
+
+            for block in chain.iter().take(count) {
+                partial_chain.push(block.clone());
+            }
+            for block in chain.iter() {
+                full_chain.push(block.clone());
+            }
+
+            let forked = full_chain.fork(fork_tip_hash).expect("hash is present");
+
+            prop_assert_eq!(forked.blocks.len(), partial_chain.blocks.len());
+        });
+
+    Ok(())
+}
+
+#[test]
+fn finalized_equals_pushed() -> Result<()> {
+    zebra_test::init();
+
+    proptest!(ProptestConfig::with_cases(env::var("PROPTEST_CASES")
+                                      .ok()
+                                      .and_then(|v| v.parse().ok())
+                                      .unwrap_or(DEFAULT_PARTIAL_CHAIN_PROPTEST_CASES)),
+    |((chain, end_count) in PreparedChain::default())| {
+        let finalized_count = chain.len() - end_count;
+        let mut full_chain = Chain::default();
+        let mut partial_chain = Chain::default();
+
+        for block in chain.iter().skip(finalized_count) {
+            partial_chain.push(block.clone());
+        }
+        for block in chain.iter() {
+            full_chain.push(block.clone());
+        }
+
+        for _ in 0..finalized_count {
+            let _finalized = full_chain.pop_root();
+        }
+
+        prop_assert_eq!(full_chain.blocks.len(), partial_chain.blocks.len());
+    });
+
+    Ok(())
+}
diff --git a/zebra-state/src/service/non_finalized_state/tests/vectors.rs b/zebra-state/src/service/non_finalized_state/tests/vectors.rs
new file mode 100644
index 00000000..c6adb460
--- /dev/null
+++ b/zebra-state/src/service/non_finalized_state/tests/vectors.rs
@@ -0,0 +1,303 @@
+use std::sync::Arc;
+
+use zebra_chain::{block::Block, parameters::Network, serialization::ZcashDeserializeInto};
+use zebra_test::prelude::*;
+
+use crate::{
+    service::non_finalized_state::{Chain, NonFinalizedState},
+    tests::{FakeChainHelper, Prepare},
+};
+
+use self::assert_eq;
+
+#[test]
+fn construct_empty() {
+    zebra_test::init();
+    let _chain = Chain::default();
+}
+
+#[test]
+fn construct_single() -> Result<()> {
+    zebra_test::init();
+    let block: Arc<Block> =
+        zebra_test::vectors::BLOCK_MAINNET_434873_BYTES.zcash_deserialize_into()?;
+
+    let mut chain = Chain::default();
+    chain.push(block.prepare());
+
+    assert_eq!(1, chain.blocks.len());
+
+    Ok(())
+}
+
+#[test]
+fn construct_many() -> Result<()> {
+    zebra_test::init();
+
+    let mut block: Arc<Block> =
+        zebra_test::vectors::BLOCK_MAINNET_434873_BYTES.zcash_deserialize_into()?;
+    let mut blocks = vec![];
+
+    while blocks.len() < 100 {
+        let next_block = block.make_fake_child();
+        blocks.push(block);
+        block = next_block;
+    }
+
+    let mut chain = Chain::default();
+
+    for block in blocks {
+        chain.push(block.prepare());
+    }
+
+    assert_eq!(100, chain.blocks.len());
+
+    Ok(())
+}
+
+#[test]
+fn ord_matches_work() -> Result<()> {
+    zebra_test::init();
+    let less_block = zebra_test::vectors::BLOCK_MAINNET_434873_BYTES
+        .zcash_deserialize_into::<Arc<Block>>()?
+        .set_work(1);
+    let more_block = less_block.clone().set_work(10);
+
+    let mut lesser_chain = Chain::default();
+    lesser_chain.push(less_block.prepare());
+
+    let mut bigger_chain = Chain::default();
+    bigger_chain.push(more_block.prepare());
+
+    assert!(bigger_chain > lesser_chain);
+
+    Ok(())
+}
+
+#[test]
+fn best_chain_wins() -> Result<()> {
+    zebra_test::init();
+
+    best_chain_wins_for_network(Network::Mainnet)?;
+    best_chain_wins_for_network(Network::Testnet)?;
+
+    Ok(())
+}
+
+fn best_chain_wins_for_network(network: Network) -> Result<()> {
+    let block1: Arc<Block> = match network {
+        Network::Mainnet => {
+            zebra_test::vectors::BLOCK_MAINNET_1180900_BYTES.zcash_deserialize_into()?
+        }
+        Network::Testnet => {
+            zebra_test::vectors::BLOCK_TESTNET_1326100_BYTES.zcash_deserialize_into()?
+        }
+    };
+
+    let block2 = block1.make_fake_child().set_work(10);
+    let child = block1.make_fake_child().set_work(1);
+
+    let expected_hash = block2.hash();
+
+    let mut state = NonFinalizedState::default();
+    state.commit_new_chain(block2.prepare());
+    state.commit_new_chain(child.prepare());
+
+    let best_chain = state.best_chain().unwrap();
+    assert!(best_chain.height_by_hash.contains_key(&expected_hash));
+
+    Ok(())
+}
+
+#[test]
+fn finalize_pops_from_best_chain() -> Result<()> {
+    zebra_test::init();
+
+    finalize_pops_from_best_chain_for_network(Network::Mainnet)?;
+    finalize_pops_from_best_chain_for_network(Network::Testnet)?;
+
+    Ok(())
+}
+
+fn finalize_pops_from_best_chain_for_network(network: Network) -> Result<()> {
+    let block1: Arc<Block> = match network {
+        Network::Mainnet => {
+            zebra_test::vectors::BLOCK_MAINNET_1180900_BYTES.zcash_deserialize_into()?
+        }
+        Network::Testnet => {
+            zebra_test::vectors::BLOCK_TESTNET_1326100_BYTES.zcash_deserialize_into()?
+        }
+    };
+
+    let block2 = block1.make_fake_child().set_work(10);
+    let child = block1.make_fake_child().set_work(1);
+
+    let mut state = NonFinalizedState::default();
+    state.commit_new_chain(block1.clone().prepare());
+    state.commit_block(block2.clone().prepare());
+    state.commit_block(child.prepare());
+
+    let finalized = state.finalize();
+    assert_eq!(block1, finalized.block);
+
+    let finalized = state.finalize();
+    assert_eq!(block2, finalized.block);
+
+    assert!(state.best_chain().is_none());
+
+    Ok(())
+}
+
+#[test]
+// This test gives full coverage for `take_chain_if`
+fn commit_block_extending_best_chain_doesnt_drop_worst_chains() -> Result<()> {
+    zebra_test::init();
+
+    commit_block_extending_best_chain_doesnt_drop_worst_chains_for_network(Network::Mainnet)?;
+    commit_block_extending_best_chain_doesnt_drop_worst_chains_for_network(Network::Testnet)?;
+
+    Ok(())
+}
+
+fn commit_block_extending_best_chain_doesnt_drop_worst_chains_for_network(
+    network: Network,
+) -> Result<()> {
+    let block1: Arc<Block> = match network {
+        Network::Mainnet => {
+            zebra_test::vectors::BLOCK_MAINNET_1180900_BYTES.zcash_deserialize_into()?
+        }
+        Network::Testnet => {
+            zebra_test::vectors::BLOCK_TESTNET_1326100_BYTES.zcash_deserialize_into()?
+        }
+    };
+
+    let block2 = block1.make_fake_child().set_work(10);
+    let child1 = block1.make_fake_child().set_work(1);
+    let child2 = block2.make_fake_child().set_work(1);
+
+    let mut state = NonFinalizedState::default();
+    assert_eq!(0, state.chain_set.len());
+    state.commit_new_chain(block1.prepare());
+    assert_eq!(1, state.chain_set.len());
+    state.commit_block(block2.prepare());
+    assert_eq!(1, state.chain_set.len());
+    state.commit_block(child1.prepare());
+    assert_eq!(2, state.chain_set.len());
+    state.commit_block(child2.prepare());
+    assert_eq!(2, state.chain_set.len());
+
+    Ok(())
+}
+
+#[test]
+fn shorter_chain_can_be_best_chain() -> Result<()> {
+    zebra_test::init();
+
+    shorter_chain_can_be_best_chain_for_network(Network::Mainnet)?;
+    shorter_chain_can_be_best_chain_for_network(Network::Testnet)?;
+
+    Ok(())
+}
+
+fn shorter_chain_can_be_best_chain_for_network(network: Network) -> Result<()> {
+    let block1: Arc<Block> = match network {
+        Network::Mainnet => {
+            zebra_test::vectors::BLOCK_MAINNET_1180900_BYTES.zcash_deserialize_into()?
+        }
+        Network::Testnet => {
+            zebra_test::vectors::BLOCK_TESTNET_1326100_BYTES.zcash_deserialize_into()?
+        }
+    };
+
+    let long_chain_block1 = block1.make_fake_child().set_work(1);
+    let long_chain_block2 = long_chain_block1.make_fake_child().set_work(1);
+
+    let short_chain_block = block1.make_fake_child().set_work(3);
+
+    let mut state = NonFinalizedState::default();
+    state.commit_new_chain(block1.prepare());
+    state.commit_block(long_chain_block1.prepare());
+    state.commit_block(long_chain_block2.prepare());
+    state.commit_block(short_chain_block.prepare());
+    assert_eq!(2, state.chain_set.len());
+
+    assert_eq!(2, state.best_chain_len());
+
+    Ok(())
+}
+
+#[test]
+fn longer_chain_with_more_work_wins() -> Result<()> {
+    zebra_test::init();
+
+    longer_chain_with_more_work_wins_for_network(Network::Mainnet)?;
+    longer_chain_with_more_work_wins_for_network(Network::Testnet)?;
+
+    Ok(())
+}
+
+fn longer_chain_with_more_work_wins_for_network(network: Network) -> Result<()> {
+    let block1: Arc<Block> = match network {
+        Network::Mainnet => {
+            zebra_test::vectors::BLOCK_MAINNET_1180900_BYTES.zcash_deserialize_into()?
+        }
+        Network::Testnet => {
+            zebra_test::vectors::BLOCK_TESTNET_1326100_BYTES.zcash_deserialize_into()?
+        }
+    };
+
+    let long_chain_block1 = block1.make_fake_child().set_work(1);
+    let long_chain_block2 = long_chain_block1.make_fake_child().set_work(1);
+    let long_chain_block3 = long_chain_block2.make_fake_child().set_work(1);
+    let long_chain_block4 = long_chain_block3.make_fake_child().set_work(1);
+
+    let short_chain_block = block1.make_fake_child().set_work(3);
+
+    let mut state = NonFinalizedState::default();
+    state.commit_new_chain(block1.prepare());
+    state.commit_block(long_chain_block1.prepare());
+    state.commit_block(long_chain_block2.prepare());
+    state.commit_block(long_chain_block3.prepare());
+    state.commit_block(long_chain_block4.prepare());
+    state.commit_block(short_chain_block.prepare());
+    assert_eq!(2, state.chain_set.len());
+
+    assert_eq!(5, state.best_chain_len());
+
+    Ok(())
+}
+
+#[test]
+fn equal_length_goes_to_more_work() -> Result<()> {
+    zebra_test::init();
+
+    equal_length_goes_to_more_work_for_network(Network::Mainnet)?;
+    equal_length_goes_to_more_work_for_network(Network::Testnet)?;
+
+    Ok(())
+}
+fn equal_length_goes_to_more_work_for_network(network: Network) -> Result<()> {
+    let block1: Arc<Block> = match network {
+        Network::Mainnet => {
+            zebra_test::vectors::BLOCK_MAINNET_1180900_BYTES.zcash_deserialize_into()?
+        }
+        Network::Testnet => {
+            zebra_test::vectors::BLOCK_TESTNET_1326100_BYTES.zcash_deserialize_into()?
+        }
+    };
+
+    let less_work_child = block1.make_fake_child().set_work(1);
+    let more_work_child = block1.make_fake_child().set_work(3);
+    let expected_hash = more_work_child.hash();
+
+    let mut state = NonFinalizedState::default();
+    state.commit_new_chain(block1.prepare());
+    state.commit_block(less_work_child.prepare());
+    state.commit_block(more_work_child.prepare());
+    assert_eq!(2, state.chain_set.len());
+
+    let tip_hash = state.best_tip().unwrap().1;
+    assert_eq!(expected_hash, tip_hash);
+
+    Ok(())
+}