Add `zebra_state::init_test` helper function for tests (#2539)

* Create a `zebra_state::init_test` helper function

This function will be used as a replacement for `zebra_state::init`
inside tests. It's a simpler alternative because it can ignore any
details that aren't relevant for tests.

* Use `init_test` inside `zebra-state` tests

Update usages of `init` to use `init_test` instead, which simplifies
most cases.

* Use `zebra_state::init_test` in `zebra-consensus`

Replace usages of `zebra_state::init` with the new helper function. This
simplifies the code a bit.

zebra-consensus/Cargo.toml

@@ -42,4 +42,5 @@ tracing-error = "0.1.2"
 tracing-subscriber = "0.2.19"
 
 zebra-chain = { path = "../zebra-chain", features = ["proptest-impl"] }
+zebra-state = { path = "../zebra-state", features = ["proptest-impl"] }
 zebra-test = { path = "../zebra-test/" }

zebra-consensus/src/block/tests.rs

@@ -117,10 +117,7 @@ async fn check_transcripts() -> Result<(), Report> {
     zebra_test::init();
 
     let network = Network::Mainnet;
-    let state_service = Buffer::new(
+    let state_service = zebra_state::init_test(network);
-        zebra_state::init(zebra_state::Config::ephemeral(), network),
-        1,
-    );
 
     let block_verifier = Buffer::new(BlockVerifier::new(network, state_service.clone()), 1);
 

zebra-consensus/src/chain/tests.rs

@@ -4,7 +4,7 @@ use std::{sync::Arc, time::Duration};
 
 use color_eyre::eyre::Report;
 use once_cell::sync::Lazy;
-use tower::{layer::Layer, timeout::TimeoutLayer, Service, ServiceBuilder};
+use tower::{layer::Layer, timeout::TimeoutLayer, Service};
 
 use zebra_chain::{
     block::{self, Block},
@@ -63,9 +63,7 @@ async fn verifiers_from_network(
         + Clone
         + 'static,
 ) {
-    let state_service = ServiceBuilder::new()
+    let state_service = zs::init_test(network);
-        .buffer(1)
-        .service(zs::init(zs::Config::ephemeral(), network));
     let chain_verifier =
         crate::chain::init(Config::default(), network, state_service.clone()).await;
 
@@ -155,14 +153,7 @@ async fn verify_checkpoint(config: Config) -> Result<(), Report> {
 
     // Test that the chain::init function works. Most of the other tests use
     // init_from_verifiers.
-    let chain_verifier = super::init(
+    let chain_verifier = super::init(config.clone(), network, zs::init_test(network)).await;
-        config.clone(),
-        network,
-        ServiceBuilder::new()
-            .buffer(1)
-            .service(zs::init(zs::Config::ephemeral(), network)),
-    )
-    .await;
 
     // Add a timeout layer
     let chain_verifier =

zebra-consensus/src/checkpoint/tests.rs

@@ -9,7 +9,7 @@ use color_eyre::eyre::{eyre, Report};
 use futures::{future::TryFutureExt, stream::FuturesUnordered};
 use std::{cmp::min, convert::TryInto, mem::drop, time::Duration};
 use tokio::{stream::StreamExt, time::timeout};
-use tower::{Service, ServiceBuilder, ServiceExt};
+use tower::{Service, ServiceExt};
 use tracing_futures::Instrument;
 
 use zebra_chain::parameters::Network::*;
@@ -45,9 +45,7 @@ async fn single_item_checkpoint_list() -> Result<(), Report> {
             .cloned()
             .collect();
 
-    let state_service = ServiceBuilder::new()
+    let state_service = zebra_state::init_test(Mainnet);
-        .buffer(1)
-        .service(zebra_state::init(zebra_state::Config::ephemeral(), Mainnet));
     let mut checkpoint_verifier =
         CheckpointVerifier::from_list(genesis_checkpoint_list, Mainnet, None, state_service)
             .map_err(|e| eyre!(e))?;
@@ -129,9 +127,7 @@ async fn multi_item_checkpoint_list() -> Result<(), Report> {
         .map(|(_block, height, hash)| (*height, *hash))
         .collect();
 
-    let state_service = ServiceBuilder::new()
+    let state_service = zebra_state::init_test(Mainnet);
-        .buffer(1)
-        .service(zebra_state::init(zebra_state::Config::ephemeral(), Mainnet));
     let mut checkpoint_verifier =
         CheckpointVerifier::from_list(checkpoint_list, Mainnet, None, state_service)
             .map_err(|e| eyre!(e))?;
@@ -277,9 +273,7 @@ async fn continuous_blockchain(
         let initial_tip = restart_height.map(|block::Height(height)| {
             (blockchain[height as usize].1, blockchain[height as usize].2)
         });
-        let state_service = ServiceBuilder::new()
+        let state_service = zebra_state::init_test(Mainnet);
-            .buffer(1)
-            .service(zebra_state::init(zebra_state::Config::ephemeral(), Mainnet));
         let mut checkpoint_verifier = CheckpointVerifier::from_list(
             checkpoint_list,
             network,
@@ -457,9 +451,7 @@ async fn block_higher_than_max_checkpoint_fail() -> Result<(), Report> {
             .cloned()
             .collect();
 
-    let state_service = ServiceBuilder::new()
+    let state_service = zebra_state::init_test(Mainnet);
-        .buffer(1)
-        .service(zebra_state::init(zebra_state::Config::ephemeral(), Mainnet));
     let mut checkpoint_verifier =
         CheckpointVerifier::from_list(genesis_checkpoint_list, Mainnet, None, state_service)
             .map_err(|e| eyre!(e))?;
@@ -536,9 +528,7 @@ async fn wrong_checkpoint_hash_fail() -> Result<(), Report> {
             .cloned()
             .collect();
 
-    let state_service = ServiceBuilder::new()
+    let state_service = zebra_state::init_test(Mainnet);
-        .buffer(1)
-        .service(zebra_state::init(zebra_state::Config::ephemeral(), Mainnet));
     let mut checkpoint_verifier =
         CheckpointVerifier::from_list(genesis_checkpoint_list, Mainnet, None, state_service)
             .map_err(|e| eyre!(e))?;
@@ -720,9 +710,7 @@ async fn checkpoint_drop_cancel() -> Result<(), Report> {
         .map(|(_block, height, hash)| (*height, *hash))
         .collect();
 
-    let state_service = ServiceBuilder::new()
+    let state_service = zebra_state::init_test(Mainnet);
-        .buffer(1)
-        .service(zebra_state::init(zebra_state::Config::ephemeral(), Mainnet));
     let mut checkpoint_verifier =
         CheckpointVerifier::from_list(checkpoint_list, Mainnet, None, state_service)
             .map_err(|e| eyre!(e))?;
@@ -808,9 +796,7 @@ async fn hard_coded_mainnet() -> Result<(), Report> {
         Arc::<Block>::zcash_deserialize(&zebra_test::vectors::BLOCK_MAINNET_GENESIS_BYTES[..])?;
     let hash0 = block0.hash();
 
-    let state_service = ServiceBuilder::new()
+    let state_service = zebra_state::init_test(Mainnet);
-        .buffer(1)
-        .service(zebra_state::init(zebra_state::Config::ephemeral(), Mainnet));
     // Use the hard-coded checkpoint list
     let mut checkpoint_verifier = CheckpointVerifier::new(Network::Mainnet, None, state_service);
 

zebra-state/src/lib.rs

@@ -37,5 +37,7 @@ pub use error::{BoxError, CloneError, CommitBlockError, ValidateContextError};
 pub use request::{FinalizedBlock, HashOrHeight, PreparedBlock, Request};
 pub use response::Response;
 pub use service::init;
+#[cfg(any(test, feature = "proptest-impl"))]
+pub use service::init_test;
 
 pub(crate) use request::ContextuallyValidBlock;

zebra-state/src/service.rs

@@ -10,6 +10,8 @@ use check::difficulty::POW_MEDIAN_BLOCK_SPAN;
 use futures::future::FutureExt;
 use non_finalized_state::{NonFinalizedState, QueuedBlocks};
 use tokio::sync::oneshot;
+#[cfg(any(test, feature = "proptest-impl"))]
+use tower::buffer::Buffer;
 use tower::{util::BoxService, Service};
 use tracing::instrument;
 use zebra_chain::{
@@ -746,6 +748,16 @@ pub fn init(config: Config, network: Network) -> BoxService<Request, Response, B
     BoxService::new(StateService::new(config, network))
 }
 
+/// Initialize a state service with an ephemeral [`Config`] and a buffer with a single slot.
+///
+/// This can be used to create a state service for testing. See also [`init`].
+#[cfg(any(test, feature = "proptest-impl"))]
+pub fn init_test(network: Network) -> Buffer<BoxService<Request, Response, BoxError>, Request> {
+    let state_service = StateService::new(Config::ephemeral(), network);
+
+    Buffer::new(BoxService::new(state_service), 1)
+}
+
 /// Check if zebra is following a legacy chain and return an error if so.
 fn legacy_chain_check<I>(
     nu5_activation_height: block::Height,

zebra-state/src/service/tests.rs

@@ -1,7 +1,7 @@
 use std::{env, sync::Arc};
 
 use futures::stream::FuturesUnordered;
-use tower::{util::BoxService, Service, ServiceExt};
+use tower::{buffer::Buffer, util::BoxService, Service, ServiceExt};
 use zebra_chain::{
     block::Block,
     parameters::{Network, NetworkUpgrade},
@@ -10,20 +10,19 @@ use zebra_chain::{
 };
 use zebra_test::{prelude::*, transcript::Transcript};
 
-use crate::{init, tests::setup::partial_nu5_chain_strategy, BoxError, Config, Request, Response};
+use crate::{init_test, tests::setup::partial_nu5_chain_strategy, BoxError, Request, Response};
 
 const LAST_BLOCK_HEIGHT: u32 = 10;
 
 async fn populated_state(
     blocks: impl IntoIterator<Item = Arc<Block>>,
-) -> BoxService<Request, Response, BoxError> {
+) -> Buffer<BoxService<Request, Response, BoxError>, Request> {
     let requests = blocks
         .into_iter()
         .map(|block| Request::CommitFinalizedBlock(block.into()));
 
-    let config = Config::ephemeral();
     let network = Network::Mainnet;
-    let mut state = init(config, network);
+    let mut state = init_test(network);
 
     let mut responses = FuturesUnordered::new();
 
@@ -41,7 +40,7 @@ async fn populated_state(
 }
 
 async fn test_populated_state_responds_correctly(
-    mut state: BoxService<Request, Response, BoxError>,
+    mut state: Buffer<BoxService<Request, Response, BoxError>, Request>,
 ) -> Result<()> {
     let blocks = zebra_test::vectors::MAINNET_BLOCKS
         .range(0..=LAST_BLOCK_HEIGHT)
@@ -152,9 +151,8 @@ async fn empty_state_still_responds_to_requests() -> Result<()> {
     .into_iter();
     let transcript = Transcript::from(iter);
 
-    let config = Config::ephemeral();
     let network = Network::Mainnet;
-    let state = init(config, network);
+    let state = init_test(network);
 
     transcript.check(state).await?;