Test that non-finalized block rejections reset the state (#2495)

* Document Ord for Chain and Proof of Work

* Create a NonFinalizedState::new method

And add some debug and clone derives.

* Test that block rejection restores internal non-finalized chain states

As part of this change, add `eq_internal_state` methods,
and proptests for them.

* Check that the chain's nullifiers are not modified on error

* Clarify a test description

* Refactor loop for readability

Co-authored-by: Janito Vaqueiro Ferreira Filho <janito.vff@gmail.com>

* Fix a variable name typo

Co-authored-by: Janito Vaqueiro Ferreira Filho <janito.vff@gmail.com>

zebra-state/src/service.rs

@@ -67,10 +67,7 @@ impl StateService {
     pub fn new(config: Config, network: Network) -> Self {
         let disk = FinalizedState::new(&config, network);
 
-        let mem = NonFinalizedState {
+        let mem = NonFinalizedState::new(network);
-            network,
-            ..Default::default()
-        };
         let queued_blocks = QueuedBlocks::default();
         let pending_utxos = PendingUtxos::default();
 

zebra-state/src/service/check/nullifier.rs

@@ -74,8 +74,6 @@ where
         }
     }
 
-    // TODO:  test that the chain's nullifiers are not modified on error (this PR)
-
     Ok(())
 }
 

zebra-state/src/service/check/tests/prop.rs

@@ -60,6 +60,7 @@ proptest! {
             .push(transaction.into());
 
         let (mut state, _genesis) = new_state_with_mainnet_genesis();
+        let previous_mem = state.mem.clone();
 
         // randomly choose to commit the block to the finalized or non-finalized state
         if use_finalized_state {
@@ -68,6 +69,7 @@ proptest! {
 
             prop_assert_eq!(Some((Height(1), block1.hash)), state.best_tip());
             prop_assert!(commit_result.is_ok());
+            prop_assert!(state.mem.eq_internal_state(&previous_mem));
         } else {
             let block1 = Arc::new(block1).prepare();
             let commit_result =
@@ -75,6 +77,7 @@ proptest! {
 
             prop_assert_eq!(commit_result, Ok(()));
             prop_assert_eq!(Some((Height(1), block1.hash)), state.best_tip());
+            prop_assert!(!state.mem.eq_internal_state(&previous_mem));
         }
     }
 
@@ -108,6 +111,7 @@ proptest! {
             .push(transaction.into());
 
         let (mut state, genesis) = new_state_with_mainnet_genesis();
+        let previous_mem = state.mem.clone();
 
         let block1 = Arc::new(block1).prepare();
         let commit_result =
@@ -126,6 +130,7 @@ proptest! {
         );
         // block was rejected
         prop_assert_eq!(Some((Height(0), genesis.hash)), state.best_tip());
+        prop_assert!(state.mem.eq_internal_state(&previous_mem));
     }
 
     /// Make sure duplicate sprout nullifiers are rejected by state contextual validation,
@@ -161,6 +166,7 @@ proptest! {
             .push(transaction.into());
 
         let (mut state, genesis) = new_state_with_mainnet_genesis();
+        let previous_mem = state.mem.clone();
 
         let block1 = Arc::new(block1).prepare();
         let commit_result =
@@ -176,6 +182,7 @@ proptest! {
             )
         );
         prop_assert_eq!(Some((Height(0), genesis.hash)), state.best_tip());
+        prop_assert!(state.mem.eq_internal_state(&previous_mem));
     }
 
     /// Make sure duplicate sprout nullifiers are rejected by state contextual validation,
@@ -219,6 +226,7 @@ proptest! {
             .push(transaction2.into());
 
         let (mut state, genesis) = new_state_with_mainnet_genesis();
+        let previous_mem = state.mem.clone();
 
         let block1 = Arc::new(block1).prepare();
         let commit_result =
@@ -234,6 +242,7 @@ proptest! {
             )
         );
         prop_assert_eq!(Some((Height(0), genesis.hash)), state.best_tip());
+        prop_assert!(state.mem.eq_internal_state(&previous_mem));
     }
 
     /// Make sure duplicate sprout nullifiers are rejected by state contextual validation,
@@ -282,6 +291,7 @@ proptest! {
             .push(transaction2.into());
 
         let (mut state, _genesis) = new_state_with_mainnet_genesis();
+        let mut previous_mem = state.mem.clone();
 
         let block1_hash;
         // randomly choose to commit the next block to the finalized or non-finalized state
@@ -291,6 +301,7 @@ proptest! {
 
             prop_assert_eq!(Some((Height(1), block1.hash)), state.best_tip());
             prop_assert!(commit_result.is_ok());
+            prop_assert!(state.mem.eq_internal_state(&previous_mem));
 
             block1_hash = block1.hash;
         } else {
@@ -300,8 +311,10 @@ proptest! {
 
             prop_assert_eq!(commit_result, Ok(()));
             prop_assert_eq!(Some((Height(1), block1.hash)), state.best_tip());
+            prop_assert!(!state.mem.eq_internal_state(&previous_mem));
 
             block1_hash = block1.hash;
+            previous_mem = state.mem.clone();
         }
 
         let block2 = Arc::new(block2).prepare();
@@ -318,6 +331,7 @@ proptest! {
             )
         );
         prop_assert_eq!(Some((Height(1), block1_hash)), state.best_tip());
+        prop_assert!(state.mem.eq_internal_state(&previous_mem));
     }
 }
 

zebra-state/src/service/non_finalized_state.rs

@@ -24,17 +24,51 @@ use crate::{FinalizedBlock, HashOrHeight, PreparedBlock, ValidateContextError};
 use self::chain::Chain;
 
 /// The state of the chains in memory, incuding queued blocks.
-#[derive(Default)]
+#[derive(Debug, Clone)]
 pub struct NonFinalizedState {
     /// Verified, non-finalized chains, in ascending order.
     ///
     /// The best chain is `chain_set.last()` or `chain_set.iter().next_back()`.
     pub chain_set: BTreeSet<Box<Chain>>,
-    /// The configured Zcash network
+
+    /// The configured Zcash network.
+    //
+    // Note: this field is currently unused, but it's useful for debugging.
     pub network: Network,
 }
 
 impl NonFinalizedState {
+    /// Returns a new non-finalized state for `network`.
+    pub fn new(network: Network) -> NonFinalizedState {
+        NonFinalizedState {
+            chain_set: Default::default(),
+            network,
+        }
+    }
+
+    /// Is the internal state of `self` the same as `other`?
+    ///
+    /// [`Chain`] has a custom [`Eq`] implementation based on proof of work,
+    /// which is used to select the best chain. So we can't derive [`Eq`] for [`NonFinalizedState`].
+    ///
+    /// Unlike the custom trait impl, this method returns `true` if the entire internal state
+    /// of two non-finalized states is equal.
+    ///
+    /// If the internal states are different, it returns `false`,
+    /// even if the chains and blocks are equal.
+    #[cfg(test)]
+    pub(crate) fn eq_internal_state(&self, other: &NonFinalizedState) -> bool {
+        // this method must be updated every time a field is added to NonFinalizedState
+
+        self.chain_set.len() == other.chain_set.len()
+            && self
+                .chain_set
+                .iter()
+                .zip(other.chain_set.iter())
+                .all(|(self_chain, other_chain)| self_chain.eq_internal_state(other_chain))
+            && self.network == other.network
+    }
+
     /// Finalize the lowest height block in the non-finalized portion of the best
     /// chain and update all side-chains to match.
     pub fn finalize(&mut self) -> FinalizedBlock {

zebra-state/src/service/non_finalized_state/chain.rs

@@ -13,24 +13,64 @@ use zebra_chain::{
 
 use crate::{service::check, PreparedBlock, ValidateContextError};
 
-#[derive(Default, Clone)]
+#[derive(Debug, Default, Clone)]
 pub struct Chain {
     pub blocks: BTreeMap<block::Height, PreparedBlock>,
     pub height_by_hash: HashMap<block::Hash, block::Height>,
     pub tx_by_hash: HashMap<transaction::Hash, (block::Height, usize)>,
 
     pub created_utxos: HashMap<transparent::OutPoint, transparent::Utxo>,
-    spent_utxos: HashSet<transparent::OutPoint>,
+    pub(super) spent_utxos: HashSet<transparent::OutPoint>,
-    // TODO: add sprout, sapling and orchard anchors (#1320)
+
-    sprout_anchors: HashSet<sprout::tree::Root>,
+    pub(super) sprout_anchors: HashSet<sprout::tree::Root>,
-    sapling_anchors: HashSet<sapling::tree::Root>,
+    pub(super) sapling_anchors: HashSet<sapling::tree::Root>,
-    sprout_nullifiers: HashSet<sprout::Nullifier>,
+    pub(super) orchard_anchors: HashSet<orchard::tree::Root>,
-    sapling_nullifiers: HashSet<sapling::Nullifier>,
+
-    orchard_nullifiers: HashSet<orchard::Nullifier>,
+    pub(super) sprout_nullifiers: HashSet<sprout::Nullifier>,
-    partial_cumulative_work: PartialCumulativeWork,
+    pub(super) sapling_nullifiers: HashSet<sapling::Nullifier>,
+    pub(super) orchard_nullifiers: HashSet<orchard::Nullifier>,
+
+    pub(super) partial_cumulative_work: PartialCumulativeWork,
 }
 
 impl Chain {
+    /// Is the internal state of `self` the same as `other`?
+    ///
+    /// [`Chain`] has custom [`Eq`] and [`Ord`] implementations based on proof of work,
+    /// which are used to select the best chain. So we can't derive [`Eq`] for [`Chain`].
+    ///
+    /// Unlike the custom trait impls, this method returns `true` if the entire internal state
+    /// of two chains is equal.
+    ///
+    /// If the internal states are different, it returns `false`,
+    /// even if the blocks in the two chains are equal.
+    #[cfg(test)]
+    pub(crate) fn eq_internal_state(&self, other: &Chain) -> bool {
+        // this method must be updated every time a field is added to Chain
+
+        // blocks, heights, hashes
+        self.blocks == other.blocks &&
+            self.height_by_hash == other.height_by_hash &&
+            self.tx_by_hash == other.tx_by_hash &&
+
+            // transparent UTXOs
+            self.created_utxos == other.created_utxos &&
+            self.spent_utxos == other.spent_utxos &&
+
+            // anchors
+            self.sprout_anchors == other.sprout_anchors &&
+            self.sapling_anchors == other.sapling_anchors &&
+            self.orchard_anchors == other.orchard_anchors &&
+
+            // nullifiers
+            self.sprout_nullifiers == other.sprout_nullifiers &&
+            self.sapling_nullifiers == other.sapling_nullifiers &&
+            self.orchard_nullifiers == other.orchard_nullifiers &&
+
+            // proof of work
+            self.partial_cumulative_work == other.partial_cumulative_work
+    }
+
     /// Push a contextually valid non-finalized block into this chain as the new tip.
     ///
     /// If the block is invalid, drop this chain and return an error.
@@ -448,21 +488,47 @@ impl UpdateWith<Option<orchard::ShieldedData>> for Chain {
     }
 }
 
-impl PartialEq for Chain {
-    fn eq(&self, other: &Self) -> bool {
-        self.partial_cmp(other) == Some(Ordering::Equal)
-    }
-}
-
-impl Eq for Chain {}
-
-impl PartialOrd for Chain {
-    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
-        Some(self.cmp(other))
-    }
-}
-
 impl Ord for Chain {
+    /// Chain order for the [`NonFinalizedState`]'s `chain_set`.
+    /// Chains with higher cumulative Proof of Work are [`Ordering::Greater`],
+    /// breaking ties using the tip block hash.
+    ///
+    /// Despite the consensus rules, Zebra uses the tip block hash as a tie-breaker.
+    /// Zebra blocks are downloaded in parallel, so download timestamps may not be unique.
+    /// (And Zebra currently doesn't track download times, because [`Block`]s are immutable.)
+    ///
+    /// This departure from the consensus rules may delay network convergence,
+    /// for as long as the greater hash belongs to the later mined block.
+    /// But Zebra nodes should converge as soon as the tied work is broken.
+    ///
+    /// "At a given point in time, each full validator is aware of a set of candidate blocks.
+    /// These form a tree rooted at the genesis block, where each node in the tree
+    /// refers to its parent via the hashPrevBlock block header field.
+    ///
+    /// A path from the root toward the leaves of the tree consisting of a sequence
+    /// of one or more valid blocks consistent with consensus rules,
+    /// is called a valid block chain.
+    ///
+    /// In order to choose the best valid block chain in its view of the overall block tree,
+    /// a node sums the work ... of all blocks in each valid block chain,
+    /// and considers the valid block chain with greatest total work to be best.
+    ///
+    /// To break ties between leaf blocks, a node will prefer the block that it received first.
+    ///
+    /// The consensus protocol is designed to ensure that for any given block height,
+    /// the vast majority of nodes should eventually agree on their best valid block chain
+    /// up to that height."
+    ///
+    /// https://zips.z.cash/protocol/protocol.pdf#blockchain
+    ///
+    /// # Panics
+    ///
+    /// If two chains compare equal.
+    ///
+    /// This panic enforces the `NonFinalizedState.chain_set` unique chain invariant.
+    ///
+    /// If the chain set contains duplicate chains, the non-finalized state might
+    /// handle new blocks or block finalization incorrectly.
     fn cmp(&self, other: &Self) -> Ordering {
         if self.partial_cumulative_work != other.partial_cumulative_work {
             self.partial_cumulative_work
@@ -491,3 +557,25 @@ impl Ord for Chain {
         }
     }
 }
+
+impl PartialOrd for Chain {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl PartialEq for Chain {
+    /// Chain equality for the [`NonFinalizedState`]'s `chain_set`,
+    /// using proof of work, then the tip block hash as a tie-breaker.
+    ///
+    /// # Panics
+    ///
+    /// If two chains compare equal.
+    ///
+    /// See [`Chain::cmp`] for details.
+    fn eq(&self, other: &Self) -> bool {
+        self.partial_cmp(other) == Some(Ordering::Equal)
+    }
+}
+
+impl Eq for Chain {}

zebra-state/src/service/non_finalized_state/tests/prop.rs

@@ -1,11 +1,20 @@
-use std::env;
+use std::{env, sync::Arc};
 
 use zebra_test::prelude::*;
 
-use crate::service::{arbitrary::PreparedChain, non_finalized_state::Chain};
+use zebra_chain::{block::Block, fmt::DisplayToDebug, parameters::NetworkUpgrade::*, LedgerState};
+
+use crate::{
+    service::{
+        arbitrary::PreparedChain,
+        non_finalized_state::{Chain, NonFinalizedState},
+    },
+    tests::Prepare,
+};
 
 const DEFAULT_PARTIAL_CHAIN_PROPTEST_CASES: u32 = 32;
 
+/// Check that a forked chain is the same as a chain that had the same blocks appended.
 #[test]
 fn forked_equals_pushed() -> Result<()> {
     zebra_test::init();
@@ -14,12 +23,13 @@ fn forked_equals_pushed() -> Result<()> {
                                           .ok()
                                           .and_then(|v| v.parse().ok())
                                           .unwrap_or(DEFAULT_PARTIAL_CHAIN_PROPTEST_CASES)),
-        |((chain, count, _network) in PreparedChain::default())| {
+        |((chain, fork_at_count, _network) in PreparedChain::default())| {
-            let fork_tip_hash = chain[count - 1].hash;
+            // use `fork_at_count` as the fork tip
+            let fork_tip_hash = chain[fork_at_count - 1].hash;
             let mut full_chain = Chain::default();
             let mut partial_chain = Chain::default();
 
-            for block in chain.iter().take(count) {
+            for block in chain.iter().take(fork_at_count) {
                 partial_chain = partial_chain.push(block.clone())?;
             }
             for block in chain.iter() {
@@ -28,12 +38,16 @@ fn forked_equals_pushed() -> Result<()> {
 
             let forked = full_chain.fork(fork_tip_hash).expect("fork works").expect("hash is present");
 
+            // the first check is redundant, but it's useful for debugging
             prop_assert_eq!(forked.blocks.len(), partial_chain.blocks.len());
+            prop_assert!(forked.eq_internal_state(&partial_chain));
         });
 
     Ok(())
 }
 
+/// Check that a chain with some blocks finalized is the same as
+/// a chain that never had those blocks added.
 #[test]
 fn finalized_equals_pushed() -> Result<()> {
     zebra_test::init();
@@ -43,6 +57,7 @@ fn finalized_equals_pushed() -> Result<()> {
                                       .and_then(|v| v.parse().ok())
                                       .unwrap_or(DEFAULT_PARTIAL_CHAIN_PROPTEST_CASES)),
     |((chain, end_count, _network) in PreparedChain::default())| {
+        // use `end_count` as the number of non-finalized blocks at the end of the chain
         let finalized_count = chain.len() - end_count;
         let mut full_chain = Chain::default();
         let mut partial_chain = Chain::default();
@@ -59,6 +74,162 @@ fn finalized_equals_pushed() -> Result<()> {
         }
 
         prop_assert_eq!(full_chain.blocks.len(), partial_chain.blocks.len());
+        prop_assert!(full_chain.eq_internal_state(&partial_chain));
+    });
+
+    Ok(())
+}
+
+/// Check that rejected blocks do not change the internal state of a chain
+/// in a non-finalized state.
+#[test]
+fn rejection_restores_internal_state() -> 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, valid_count, network, mut bad_block) in (PreparedChain::default(), any::<bool>(), any::<bool>())
+                .prop_flat_map(|((chain, valid_count, network), is_nu5, is_v5)| {
+                    let next_height = chain[valid_count - 1].height;
+                    (
+                        Just(chain),
+                        Just(valid_count),
+                        Just(network),
+                        // generate a Canopy or NU5 block with v4 or v5 transactions
+                        LedgerState::height_strategy(
+                            next_height,
+                            if is_nu5 { Nu5 } else { Canopy },
+                            if is_nu5 && is_v5 { 5 } else { 4 },
+                            true,
+                        )
+                            .prop_flat_map(Block::arbitrary_with)
+                            .prop_map(DisplayToDebug)
+                    )
+                }
+                ))| {
+                  let mut state = NonFinalizedState::new(network);
+
+                  // use `valid_count` as the number of valid blocks before an invalid block
+                  let valid_tip_height = chain[valid_count - 1].height;
+                  let valid_tip_hash = chain[valid_count - 1].hash;
+                  let mut chain = chain.iter().take(valid_count).cloned();
+
+                  prop_assert!(state.eq_internal_state(&state));
+
+                  if let Some(first_block) = chain.next() {
+                      state.commit_new_chain(first_block)?;
+                      prop_assert!(state.eq_internal_state(&state));
+                  }
+
+                  for block in chain {
+                      state.commit_block(block)?;
+                      prop_assert!(state.eq_internal_state(&state));
+                  }
+
+                  prop_assert_eq!(state.best_tip(), Some((valid_tip_height, valid_tip_hash)));
+
+                  let mut reject_state = state.clone();
+                  // the tip check is redundant, but it's useful for debugging
+                  prop_assert_eq!(state.best_tip(), reject_state.best_tip());
+                  prop_assert!(state.eq_internal_state(&reject_state));
+
+                  bad_block.header.previous_block_hash = valid_tip_hash;
+                  let bad_block = Arc::new(bad_block.0).prepare();
+                  let reject_result = reject_state.commit_block(bad_block);
+
+                  if reject_result.is_err() {
+                      prop_assert_eq!(state.best_tip(), reject_state.best_tip());
+                      prop_assert!(state.eq_internal_state(&reject_state));
+                  } else {
+                      // the block just happened to pass all the non-finalized checks
+                      prop_assert_ne!(state.best_tip(), reject_state.best_tip());
+                      prop_assert!(!state.eq_internal_state(&reject_state));
+                  }
+              });
+
+    Ok(())
+}
+
+/// Check that different blocks create different internal chain states,
+/// and that all the state fields are covered by `eq_internal_state`.
+#[test]
+fn different_blocks_different_chains() -> 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)),
+    |((block1, block2) in (any::<bool>(), any::<bool>())
+      .prop_flat_map(|(is_nu5, is_v5)| {
+          // generate a Canopy or NU5 block with v4 or v5 transactions
+          LedgerState::coinbase_strategy(
+              if is_nu5 { Nu5 } else { Canopy },
+              if is_nu5 && is_v5 { 5 } else { 4 },
+              true,
+          )})
+      .prop_map(Block::arbitrary_with)
+      .prop_flat_map(|block_strategy| (block_strategy.clone(), block_strategy))
+      .prop_map(|(block1, block2)| (DisplayToDebug(block1), DisplayToDebug(block2)))
+    )| {
+        let chain1 = Chain::default();
+        let chain2 = Chain::default();
+
+        let block1 = Arc::new(block1.0).prepare();
+        let block2 = Arc::new(block2.0).prepare();
+
+        let result1 = chain1.push(block1.clone());
+        let result2 = chain2.push(block2.clone());
+
+        // if there is an error, we don't get the chains back
+        if let (Ok(mut chain1), Ok(chain2)) = (result1, result2) {
+            if block1 == block2 {
+                // the blocks were equal, so the chains should be equal
+
+                // the first check is redundant, but it's useful for debugging
+                prop_assert_eq!(&chain1.height_by_hash, &chain2.height_by_hash);
+                prop_assert!(chain1.eq_internal_state(&chain2));
+            } else {
+                // the blocks were different, so the chains should be different
+
+                prop_assert_ne!(&chain1.height_by_hash, &chain2.height_by_hash);
+                prop_assert!(!chain1.eq_internal_state(&chain2));
+
+                // We can't derive eq_internal_state,
+                // so we check for missing fields here.
+
+                // blocks, heights, hashes
+                chain1.blocks = chain2.blocks.clone();
+                chain1.height_by_hash = chain2.height_by_hash.clone();
+                chain1.tx_by_hash = chain2.tx_by_hash.clone();
+
+                // transparent UTXOs
+                chain1.created_utxos = chain2.created_utxos.clone();
+                chain1.spent_utxos = chain2.spent_utxos.clone();
+
+                // anchors
+                chain1.sprout_anchors = chain2.sprout_anchors.clone();
+                chain1.sapling_anchors = chain2.sapling_anchors.clone();
+                chain1.orchard_anchors = chain2.orchard_anchors.clone();
+
+                // nullifiers
+                chain1.sprout_nullifiers = chain2.sprout_nullifiers.clone();
+                chain1.sapling_nullifiers = chain2.sapling_nullifiers.clone();
+                chain1.orchard_nullifiers = chain2.orchard_nullifiers.clone();
+
+                // proof of work
+                chain1.partial_cumulative_work = chain2.partial_cumulative_work;
+
+                // If this check fails, the `Chain` fields are out
+                // of sync with `eq_internal_state` or this test.
+                prop_assert!(
+                    chain1.eq_internal_state(&chain2),
+                    "Chain fields, eq_internal_state, and this test must be consistent"
+                );
+            }
+        }
     });
 
     Ok(())

zebra-state/src/service/non_finalized_state/tests/vectors.rs

@@ -99,7 +99,7 @@ fn best_chain_wins_for_network(network: Network) -> Result<()> {
 
     let expected_hash = block2.hash();
 
-    let mut state = NonFinalizedState::default();
+    let mut state = NonFinalizedState::new(network);
     state.commit_new_chain(block2.prepare())?;
     state.commit_new_chain(child.prepare())?;
 
@@ -132,7 +132,7 @@ fn finalize_pops_from_best_chain_for_network(network: Network) -> Result<()> {
     let block2 = block1.make_fake_child().set_work(10);
     let child = block1.make_fake_child().set_work(1);
 
-    let mut state = NonFinalizedState::default();
+    let mut state = NonFinalizedState::new(network);
     state.commit_new_chain(block1.clone().prepare())?;
     state.commit_block(block2.clone().prepare())?;
     state.commit_block(child.prepare())?;
@@ -175,7 +175,7 @@ fn commit_block_extending_best_chain_doesnt_drop_worst_chains_for_network(
     let child1 = block1.make_fake_child().set_work(1);
     let child2 = block2.make_fake_child().set_work(1);
 
-    let mut state = NonFinalizedState::default();
+    let mut state = NonFinalizedState::new(network);
     assert_eq!(0, state.chain_set.len());
     state.commit_new_chain(block1.prepare())?;
     assert_eq!(1, state.chain_set.len());
@@ -214,7 +214,7 @@ fn shorter_chain_can_be_best_chain_for_network(network: Network) -> Result<()> {
 
     let short_chain_block = block1.make_fake_child().set_work(3);
 
-    let mut state = NonFinalizedState::default();
+    let mut state = NonFinalizedState::new(network);
     state.commit_new_chain(block1.prepare())?;
     state.commit_block(long_chain_block1.prepare())?;
     state.commit_block(long_chain_block2.prepare())?;
@@ -253,7 +253,7 @@ fn longer_chain_with_more_work_wins_for_network(network: Network) -> Result<()>
 
     let short_chain_block = block1.make_fake_child().set_work(3);
 
-    let mut state = NonFinalizedState::default();
+    let mut state = NonFinalizedState::new(network);
     state.commit_new_chain(block1.prepare())?;
     state.commit_block(long_chain_block1.prepare())?;
     state.commit_block(long_chain_block2.prepare())?;
@@ -290,7 +290,7 @@ fn equal_length_goes_to_more_work_for_network(network: Network) -> Result<()> {
     let more_work_child = block1.make_fake_child().set_work(3);
     let expected_hash = more_work_child.hash();
 
-    let mut state = NonFinalizedState::default();
+    let mut state = NonFinalizedState::new(network);
     state.commit_new_chain(block1.prepare())?;
     state.commit_block(less_work_child.prepare())?;
     state.commit_block(more_work_child.prepare())?;