linera_core/unit_tests/

test_utils.rs

// Copyright (c) Zefchain Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

#![allow(clippy::cast_possible_truncation)]

use std::{
    collections::{BTreeMap, HashMap, HashSet},
    sync::Arc,
    time::Duration,
    vec,
};

use async_trait::async_trait;
use futures::{
    future::Either,
    lock::{Mutex, MutexGuard},
    stream, Future, StreamExt as _,
};
use linera_base::{
    crypto::{
        AccountPublicKey, CryptoHash, ValidatorKeypair, ValidatorPublicKey, ValidatorSecretKey,
    },
    data_types::*,
    identifiers::{AccountOwner, BlobId, ChainId, EventId},
    ownership::ChainOwnership,
};
use linera_chain::{
    data_types::BlockProposal,
    types::{
        CertificateKind, Certified, ConfirmedBlock, ConfirmedBlockCertificate, GenericCertificate,
        LiteCertificate, Timeout, ValidatedBlock, ValidatedBlockCertificate,
    },
};
use linera_execution::{committee::Committee, ResourceControlPolicy, WasmRuntime};
use linera_storage::{Arc as CacheArc, DbStorage, ResultReadCertificates, Storage, TestClock};
#[cfg(all(not(target_arch = "wasm32"), feature = "storage-service"))]
use linera_storage_service::client::StorageServiceDatabase;
use linera_version::VersionInfo;
#[cfg(feature = "scylladb")]
use linera_views::scylla_db::ScyllaDbDatabase;
use linera_views::{
    memory::MemoryDatabase,
    random::generate_test_namespace,
    store::{KeyValueStore, TestKeyValueDatabase},
};
use tokio::sync::oneshot;
use tokio_stream::wrappers::UnboundedReceiverStream;
#[cfg(feature = "rocksdb")]
use {
    linera_views::rocks_db::RocksDbDatabase,
    tokio::sync::{Semaphore, SemaphorePermit},
};

use crate::{
    chain_worker::ChainWorkerConfig,
    client::{chain_client, Client},
    data_types::*,
    environment::{TestSigner, TestWallet},
    node::{
        CrossChainMessageDelivery, NodeError, NotificationStream, ValidatorNode,
        ValidatorNodeProvider,
    },
    notifier::ChannelNotifier,
    worker::{
        Notification, ProcessableCertificate, WorkerState, DEFAULT_BLOCK_CACHE_SIZE,
        DEFAULT_EXECUTION_STATE_CACHE_SIZE,
    },
};

/// The kind of misbehavior a test validator simulates.
#[derive(Debug, PartialEq, Clone, Copy)]
#[allow(missing_docs)]
pub enum FaultType {
    Honest,
    Offline,
    OfflineWithInfo,
    NoChains,
    DontSendConfirmVote,
    DontProcessValidated,
    DontSendValidateVote,
}

/// A validator used for testing. "Faulty" validators ignore block proposals (but not
/// certificates or info queries) and have the wrong initial balance for all chains.
///
/// All methods are executed in spawned Tokio tasks, so that canceling a client task doesn't cause
/// the validator's tasks to be canceled: In a real network, a validator also wouldn't cancel
/// tasks if the client stopped waiting for the response.
struct LocalValidator<S>
where
    S: Storage,
{
    state: WorkerState<S>,
    notifier: Arc<ChannelNotifier<Notification>>,
    /// Abort handles for the notification streams handed out by `do_subscribe`, so
    /// tests can sever them mid-flight as a validator (or proxy) restart would.
    notification_stream_aborts: Vec<stream::AbortHandle>,
    /// How many times `do_subscribe` has been called. Monotonic — never drained by
    /// `disconnect_notification_subscribers` — so a test can bound how often a client
    /// re-subscribes, which is what pins the circuit breaker's backoff.
    subscribe_calls: usize,
}

/// A client used by tests to talk to an in-process `LocalValidator`.
#[derive(Clone)]
pub struct LocalValidatorClient<S>
where
    S: Storage,
{
    public_key: ValidatorPublicKey,
    client: Arc<Mutex<LocalValidator<S>>>,
    fault_type: FaultType,
    /// Held open for reading by every chain-info query, but ONLY once armed. A test takes
    /// the write lock to stall the initial sync a probe runs after `subscribe` succeeded.
    ///
    /// The flag gate matters: `tokio::sync::RwLock` is write-preferring, so taking the read
    /// unconditionally would queue every chain-info query at every validator in every test
    /// behind a writer that only one test ever creates — and would hold the guard across
    /// the validator mutex and the whole query.
    chain_info_gate: Arc<tokio::sync::RwLock<()>>,
    chain_info_gate_armed: Arc<std::sync::atomic::AtomicBool>,
}

impl<S> ValidatorNode for LocalValidatorClient<S>
where
    S: Storage + Clone + Send + Sync + 'static,
{
    type NotificationStream = NotificationStream;

    fn address(&self) -> String {
        format!("local:{}", self.public_key)
    }

    async fn handle_block_proposal(
        &self,
        proposal: BlockProposal,
    ) -> Result<ChainInfoResponse, NodeError> {
        self.spawn_and_receive(move |validator, sender| {
            validator.do_handle_block_proposal(proposal, sender)
        })
        .await
    }

    async fn handle_lite_certificate(
        &self,
        certificate: LiteCertificate<'_>,
        _delivery: CrossChainMessageDelivery,
    ) -> Result<ChainInfoResponse, NodeError> {
        let certificate = certificate.cloned();
        self.spawn_and_receive(move |validator, sender| {
            validator.do_handle_lite_certificate(certificate, sender)
        })
        .await
    }

    async fn handle_timeout_certificate(
        &self,
        certificate: GenericCertificate<Timeout>,
    ) -> Result<ChainInfoResponse, NodeError> {
        self.spawn_and_receive(move |validator, sender| {
            validator.do_handle_certificate::<Timeout>(certificate, sender)
        })
        .await
    }

    async fn handle_validated_certificate(
        &self,
        certificate: ValidatedBlockCertificate,
    ) -> Result<ChainInfoResponse, NodeError> {
        self.spawn_and_receive(move |validator, sender| {
            validator.do_handle_certificate::<ValidatedBlock>(certificate, sender)
        })
        .await
    }

    async fn handle_confirmed_certificate(
        &self,
        certificate: CacheArc<ConfirmedBlockCertificate>,
        _delivery: CrossChainMessageDelivery,
    ) -> Result<ChainInfoResponse, NodeError> {
        self.spawn_and_receive(move |validator, sender| {
            validator.do_handle_certificate::<ConfirmedBlock>(
                CacheArc::unwrap_or_clone(certificate),
                sender,
            )
        })
        .await
    }

    async fn handle_chain_info_query(
        &self,
        query: ChainInfoQuery,
    ) -> Result<ChainInfoResponse, NodeError> {
        self.spawn_and_receive(move |validator, sender| {
            validator.do_handle_chain_info_query(query, sender)
        })
        .await
    }

    async fn subscribe(&self, chains: Vec<ChainId>) -> Result<NotificationStream, NodeError> {
        self.spawn_and_receive(move |validator, sender| validator.do_subscribe(chains, sender))
            .await
    }

    async fn get_version_info(&self) -> Result<VersionInfo, NodeError> {
        Ok(Default::default())
    }

    async fn get_network_description(&self) -> Result<NetworkDescription, NodeError> {
        Ok(self
            .client
            .lock()
            .await
            .state
            .storage_client()
            .read_network_description()
            .await
            .transpose()
            .ok_or_else(|| NodeError::ViewError {
                error: "missing NetworkDescription".to_owned(),
            })??)
    }

    async fn upload_blob(&self, content: BlobContent) -> Result<BlobId, NodeError> {
        self.spawn_and_receive(move |validator, sender| validator.do_upload_blob(content, sender))
            .await
    }

    async fn download_blob(&self, blob_id: BlobId) -> Result<BlobContent, NodeError> {
        self.spawn_and_receive(move |validator, sender| validator.do_download_blob(blob_id, sender))
            .await
    }

    async fn download_blobs(
        &self,
        blob_ids: Vec<BlobId>,
    ) -> Result<crate::node::BlobStream, NodeError> {
        let this = self.clone();
        let stream = futures::stream::unfold(blob_ids.into_iter(), move |mut iter| {
            let this = this.clone();
            async move {
                let blob_id = iter.next()?;
                let result = this.download_blob(blob_id).await;
                Some((result, iter))
            }
        });
        Ok(Box::pin(stream))
    }

    async fn download_pending_blob(
        &self,
        chain_id: ChainId,
        blob_id: BlobId,
    ) -> Result<BlobContent, NodeError> {
        self.spawn_and_receive(move |validator, sender| {
            validator.do_download_pending_blob(chain_id, blob_id, sender)
        })
        .await
    }

    async fn handle_pending_blob(
        &self,
        chain_id: ChainId,
        blob: BlobContent,
    ) -> Result<ChainInfoResponse, NodeError> {
        self.spawn_and_receive(move |validator, sender| {
            validator.do_handle_pending_blob(chain_id, blob, sender)
        })
        .await
    }

    async fn download_certificate(
        &self,
        hash: CryptoHash,
    ) -> Result<ConfirmedBlockCertificate, NodeError> {
        self.spawn_and_receive(move |validator, sender| {
            validator.do_download_certificate(hash, sender)
        })
        .await
    }

    async fn download_certificates(
        &self,
        hashes: Vec<CryptoHash>,
    ) -> Result<Vec<ConfirmedBlockCertificate>, NodeError> {
        self.spawn_and_receive(move |validator, sender| {
            validator.do_download_certificates(hashes, sender)
        })
        .await
    }

    async fn download_certificates_by_heights(
        &self,
        chain_id: ChainId,
        heights: Vec<BlockHeight>,
    ) -> Result<Vec<ConfirmedBlockCertificate>, NodeError> {
        self.spawn_and_receive(move |validator, sender| {
            validator.do_download_certificates_by_heights(chain_id, heights, sender)
        })
        .await
    }

    async fn blob_last_used_by(&self, blob_id: BlobId) -> Result<CryptoHash, NodeError> {
        self.spawn_and_receive(move |validator, sender| {
            validator.do_blob_last_used_by(blob_id, sender)
        })
        .await
    }

    async fn blob_last_used_by_certificate(
        &self,
        blob_id: BlobId,
    ) -> Result<ConfirmedBlockCertificate, NodeError> {
        self.spawn_and_receive(move |validator, sender| {
            validator.do_blob_last_used_by_certificate(blob_id, sender)
        })
        .await
    }

    async fn missing_blob_ids(&self, blob_ids: Vec<BlobId>) -> Result<Vec<BlobId>, NodeError> {
        self.spawn_and_receive(move |validator, sender| {
            validator.do_missing_blob_ids(blob_ids, sender)
        })
        .await
    }

    async fn event_block_heights(
        &self,
        event_ids: Vec<EventId>,
    ) -> Result<Vec<Option<BlockHeight>>, NodeError> {
        self.spawn_and_receive(move |validator, sender| {
            validator.do_event_block_heights(event_ids, sender)
        })
        .await
    }

    async fn get_shard_info(
        &self,
        _chain_id: ChainId,
    ) -> Result<crate::data_types::ShardInfo, NodeError> {
        // For test purposes, return a dummy shard info
        Ok(crate::data_types::ShardInfo {
            shard_id: 0,
            total_shards: 1,
        })
    }
}

impl<S> LocalValidatorClient<S>
where
    S: Storage + Clone + Send + Sync + 'static,
{
    fn new(public_key: ValidatorPublicKey, state: WorkerState<S>) -> Self {
        let client = LocalValidator {
            state,
            notifier: Arc::new(ChannelNotifier::default()),
            notification_stream_aborts: Vec::new(),
            subscribe_calls: 0,
        };
        Self {
            public_key,
            client: Arc::new(Mutex::new(client)),
            fault_type: FaultType::Honest,
            chain_info_gate: Arc::default(),
            chain_info_gate_armed: Arc::default(),
        }
    }

    /// Returns the validator's public key.
    pub fn name(&self) -> ValidatorPublicKey {
        self.public_key
    }

    /// Returns the validator's currently configured [`FaultType`].
    pub fn fault_type(&self) -> FaultType {
        self.fault_type
    }

    /// How many times this validator has been asked to `subscribe`.
    pub async fn subscribe_calls(&self) -> usize {
        self.client.lock().await.subscribe_calls
    }

    /// Ends every notification stream previously handed out by this validator — as
    /// clients experience when the validator (or its proxy) restarts.
    pub async fn disconnect_notification_subscribers(&self) {
        let mut validator = self.client.lock().await;
        for abort in validator.notification_stream_aborts.drain(..) {
            abort.abort();
        }
    }

    fn set_fault_type(&mut self, fault_type: FaultType) {
        self.fault_type = fault_type;
    }

    /// Obtains the basic `ChainInfo` data for the local validator chain, with chain manager values.
    pub async fn chain_info_with_manager_values(
        &mut self,
        chain_id: ChainId,
    ) -> Result<Box<ChainInfo>, NodeError> {
        let query = ChainInfoQuery::new(chain_id).with_manager_values();
        let response = self.handle_chain_info_query(query).await?;
        Ok(response.info)
    }

    /// Executes the future produced by `f` in a new thread in a new Tokio runtime.
    /// Returns the value that the future puts into the sender.
    async fn spawn_and_receive<F, R, T>(&self, f: F) -> T
    where
        T: Send + 'static,
        R: Future<Output = Result<(), T>> + Send,
        F: FnOnce(Self, oneshot::Sender<T>) -> R + Send + 'static,
    {
        let validator = self.clone();
        let (sender, receiver) = oneshot::channel();
        tokio::spawn(async move {
            if f(validator, sender).await.is_err() {
                tracing::debug!("result could not be sent");
            }
        });
        receiver.await.unwrap()
    }

    async fn do_handle_block_proposal(
        self,
        proposal: BlockProposal,
        sender: oneshot::Sender<Result<ChainInfoResponse, NodeError>>,
    ) -> Result<(), Result<ChainInfoResponse, NodeError>> {
        let result = match self.fault_type {
            FaultType::Offline | FaultType::OfflineWithInfo => Err(NodeError::ClientIoError {
                error: "offline".to_string(),
            }),
            FaultType::NoChains => Err(NodeError::InactiveChain(proposal.content.block.chain_id)),
            FaultType::DontSendValidateVote
            | FaultType::Honest
            | FaultType::DontSendConfirmVote
            | FaultType::DontProcessValidated => {
                let (response_result, _actions) = self
                    .client
                    .lock()
                    .await
                    .state
                    .handle_block_proposal(proposal)
                    .await;
                let result = response_result.map_err(NodeError::from);
                if self.fault_type == FaultType::DontSendValidateVote {
                    Err(NodeError::ClientIoError {
                        error: "refusing to validate".to_string(),
                    })
                } else {
                    result
                }
            }
        };
        // In a local node cross-chain messages can't get lost, so we can ignore the actions here.
        sender.send(result)
    }

    async fn do_handle_lite_certificate(
        self,
        certificate: LiteCertificate<'_>,
        sender: oneshot::Sender<Result<ChainInfoResponse, NodeError>>,
    ) -> Result<(), Result<ChainInfoResponse, NodeError>> {
        let client = self.client.clone();
        let validator = client.lock().await;
        let result = async move {
            match validator.state.full_certificate(certificate).await? {
                Either::Left(confirmed) => {
                    self.do_handle_certificate_internal::<ConfirmedBlock>(confirmed, &validator)
                        .await
                }
                Either::Right(validated) => {
                    self.do_handle_certificate_internal::<ValidatedBlock>(validated, &validator)
                        .await
                }
            }
        }
        .await;
        sender.send(result)
    }

    async fn do_handle_certificate_internal<T: ProcessableCertificate>(
        &self,
        certificate: T::Certificate,
        validator: &MutexGuard<'_, LocalValidator<S>>,
    ) -> Result<ChainInfoResponse, NodeError> {
        match self.fault_type {
            FaultType::DontProcessValidated if T::KIND == CertificateKind::Validated => {
                Err(NodeError::ClientIoError {
                    error: "refusing to process validated block".to_string(),
                })
            }
            FaultType::NoChains => Err(NodeError::InactiveChain(certificate.value().chain_id())),
            FaultType::Honest
            | FaultType::DontSendConfirmVote
            | FaultType::DontProcessValidated
            | FaultType::DontSendValidateVote => {
                let result = validator
                    .state
                    .fully_handle_certificate_with_notifications(certificate, &validator.notifier)
                    .await
                    .map_err(Into::into);
                if T::KIND == CertificateKind::Validated
                    && self.fault_type == FaultType::DontSendConfirmVote
                {
                    Err(NodeError::ClientIoError {
                        error: "refusing to confirm".to_string(),
                    })
                } else {
                    result
                }
            }
            FaultType::Offline | FaultType::OfflineWithInfo => Err(NodeError::ClientIoError {
                error: "offline".to_string(),
            }),
        }
    }

    async fn do_handle_certificate<T: ProcessableCertificate>(
        self,
        certificate: T::Certificate,
        sender: oneshot::Sender<Result<ChainInfoResponse, NodeError>>,
    ) -> Result<(), Result<ChainInfoResponse, NodeError>> {
        let validator = self.client.lock().await;
        let result = self
            .do_handle_certificate_internal::<T>(certificate, &validator)
            .await;
        sender.send(result)
    }

    async fn do_handle_chain_info_query(
        self,
        query: ChainInfoQuery,
        sender: oneshot::Sender<Result<ChainInfoResponse, NodeError>>,
    ) -> Result<(), Result<ChainInfoResponse, NodeError>> {
        let gate = self.chain_info_gate.clone();
        let _open = if self
            .chain_info_gate_armed
            .load(std::sync::atomic::Ordering::Relaxed)
        {
            Some(gate.read().await)
        } else {
            None
        };
        let validator = self.client.lock().await;
        let result = match self.fault_type {
            FaultType::Offline => Err(NodeError::ClientIoError {
                error: "offline".to_string(),
            }),
            FaultType::NoChains => Err(NodeError::InactiveChain(query.chain_id)),
            FaultType::Honest
            | FaultType::DontSendConfirmVote
            | FaultType::DontProcessValidated
            | FaultType::DontSendValidateVote
            | FaultType::OfflineWithInfo => validator
                .state
                .handle_chain_info_query(query)
                .await
                .map_err(Into::into),
        };
        sender.send(result)
    }

    async fn do_subscribe(
        self,
        chains: Vec<ChainId>,
        sender: oneshot::Sender<Result<NotificationStream, NodeError>>,
    ) -> Result<(), Result<NotificationStream, NodeError>> {
        let mut validator = self.client.lock().await;
        validator.subscribe_calls += 1;
        // Both offline variants refuse: a notification stream is not the info query that
        // `OfflineWithInfo` still answers, and a probe must be able to genuinely FAIL.
        if matches!(
            self.fault_type,
            FaultType::Offline | FaultType::OfflineWithInfo
        ) {
            return sender.send(Err(NodeError::ClientIoError {
                error: "offline".to_string(),
            }));
        }
        let rx = validator.notifier.subscribe(chains);
        let (stream, abort) = stream::abortable(UnboundedReceiverStream::new(rx));
        // Marking the handle aborted when the client merely DROPS its end is what makes
        // the prune below see the real set: a dropped stream is never aborted by anyone,
        // so without this the vec grows across every re-subscribe and
        // `disconnect_notification_subscribers` walks handles nobody holds.
        let abort_on_drop = AbortOnDrop(abort.clone());
        let stream = stream.map(move |notification| {
            let _abort_on_drop = &abort_on_drop;
            notification
        });
        validator
            .notification_stream_aborts
            .retain(|abort| !abort.is_aborted());
        validator.notification_stream_aborts.push(abort);
        let stream: NotificationStream = Box::pin(stream);
        sender.send(Ok(stream))
    }

    async fn do_upload_blob(
        self,
        content: BlobContent,
        sender: oneshot::Sender<Result<BlobId, NodeError>>,
    ) -> Result<(), Result<BlobId, NodeError>> {
        let validator = self.client.lock().await;
        let blob = Blob::new(content);
        let id = blob.id();
        let storage = validator.state.storage_client();
        let result = match storage.maybe_write_blobs(&[blob]).await {
            Ok(has_state) if has_state.first() == Some(&true) => Ok(id),
            Ok(_) => Err(NodeError::BlobsNotFound(vec![id])),
            Err(error) => Err(error.into()),
        };
        sender.send(result)
    }

    async fn do_download_blob(
        self,
        blob_id: BlobId,
        sender: oneshot::Sender<Result<BlobContent, NodeError>>,
    ) -> Result<(), Result<BlobContent, NodeError>> {
        let validator = self.client.lock().await;
        let blob = validator
            .state
            .storage_client()
            .read_blob(blob_id)
            .await
            .map_err(Into::into);
        let blob = match blob {
            Ok(blob) => blob.ok_or_else(|| NodeError::BlobsNotFound(vec![blob_id])),
            Err(error) => Err(error),
        };
        sender.send(blob.map(|blob| CacheArc::unwrap_or_clone(blob).into_content()))
    }

    async fn do_download_pending_blob(
        self,
        chain_id: ChainId,
        blob_id: BlobId,
        sender: oneshot::Sender<Result<BlobContent, NodeError>>,
    ) -> Result<(), Result<BlobContent, NodeError>> {
        let validator = self.client.lock().await;
        let result = validator
            .state
            .download_pending_blob(chain_id, blob_id)
            .await
            .map_err(Into::into);
        sender.send(result.map(|blob| blob.content().clone()))
    }

    async fn do_handle_pending_blob(
        self,
        chain_id: ChainId,
        blob: BlobContent,
        sender: oneshot::Sender<Result<ChainInfoResponse, NodeError>>,
    ) -> Result<(), Result<ChainInfoResponse, NodeError>> {
        let validator = self.client.lock().await;
        let result = validator
            .state
            .handle_pending_blob(chain_id, Blob::new(blob))
            .await
            .map_err(Into::into);
        sender.send(result)
    }

    async fn do_download_certificate(
        self,
        hash: CryptoHash,
        sender: oneshot::Sender<Result<ConfirmedBlockCertificate, NodeError>>,
    ) -> Result<(), Result<ConfirmedBlockCertificate, NodeError>> {
        let validator = self.client.lock().await;
        let certificate = validator
            .state
            .storage_client()
            .read_certificate(hash)
            .await
            .map_err(Into::into);

        let certificate = match certificate {
            Err(error) => Err(error),
            Ok(entry) => match entry {
                Some(certificate) => Ok(CacheArc::unwrap_or_clone(certificate)),
                None => {
                    panic!("Missing certificate: {hash}");
                }
            },
        };

        sender.send(certificate)
    }

    async fn do_download_certificates(
        self,
        hashes: Vec<CryptoHash>,
        sender: oneshot::Sender<Result<Vec<ConfirmedBlockCertificate>, NodeError>>,
    ) -> Result<(), Result<Vec<ConfirmedBlockCertificate>, NodeError>> {
        let validator = self.client.lock().await;
        let certificates = validator
            .state
            .storage_client()
            .read_certificates(&hashes)
            .await
            .map_err(Into::into);

        let certificates = match certificates {
            Err(error) => Err(error),
            Ok(certificates) => match ResultReadCertificates::new(certificates, hashes) {
                ResultReadCertificates::Certificates(certificates) => Ok(certificates),
                ResultReadCertificates::InvalidHashes(hashes) => {
                    panic!("Missing certificates: {hashes:?}")
                }
            },
        };

        sender.send(certificates)
    }

    async fn do_download_certificates_by_heights(
        self,
        chain_id: ChainId,
        heights: Vec<BlockHeight>,
        sender: oneshot::Sender<Result<Vec<ConfirmedBlockCertificate>, NodeError>>,
    ) -> Result<(), Result<Vec<ConfirmedBlockCertificate>, NodeError>> {
        // First, use do_handle_chain_info_query to get the certificate hashes
        let (query_sender, query_receiver) = oneshot::channel();
        let query = ChainInfoQuery::new(chain_id).with_sent_certificate_hashes_by_heights(heights);

        let self_clone = self.clone();
        self.do_handle_chain_info_query(query, query_sender)
            .await
            .expect("Failed to handle chain info query");

        // Get the response from the chain info query
        let chain_info_response = query_receiver.await.map_err(|_| {
            Err(NodeError::ClientIoError {
                error: "Failed to receive chain info response".to_string(),
            })
        })?;

        let hashes = match chain_info_response {
            Ok(response) => response.info.requested_sent_certificate_hashes,
            Err(e) => {
                return sender.send(Err(e));
            }
        };

        // Now use do_download_certificates to get the actual certificates
        let (cert_sender, cert_receiver) = oneshot::channel();
        self_clone
            .do_download_certificates(hashes, cert_sender)
            .await?;

        // Forward the result to the original sender
        let result = cert_receiver.await.map_err(|_| {
            Err(NodeError::ClientIoError {
                error: "Failed to receive certificates".to_string(),
            })
        })?;

        sender.send(result)
    }

    async fn do_blob_last_used_by(
        self,
        blob_id: BlobId,
        sender: oneshot::Sender<Result<CryptoHash, NodeError>>,
    ) -> Result<(), Result<CryptoHash, NodeError>> {
        let validator = self.client.lock().await;
        let blob_state = validator
            .state
            .storage_client()
            .read_blob_state(blob_id)
            .await
            .map_err(Into::into);
        let certificate_hash = match blob_state {
            Err(err) => Err(err),
            Ok(blob_state) => match blob_state {
                None => Err(NodeError::BlobsNotFound(vec![blob_id])),
                Some(blob_state) => blob_state
                    .last_used_by
                    .ok_or_else(|| NodeError::BlobsNotFound(vec![blob_id])),
            },
        };

        sender.send(certificate_hash)
    }

    async fn do_blob_last_used_by_certificate(
        self,
        blob_id: BlobId,
        sender: oneshot::Sender<Result<ConfirmedBlockCertificate, NodeError>>,
    ) -> Result<(), Result<ConfirmedBlockCertificate, NodeError>> {
        match self.blob_last_used_by(blob_id).await {
            Ok(cert_hash) => {
                let cert = self.download_certificate(cert_hash).await;
                sender.send(cert)
            }
            Err(err) => sender.send(Err(err)),
        }
    }

    async fn do_missing_blob_ids(
        self,
        blob_ids: Vec<BlobId>,
        sender: oneshot::Sender<Result<Vec<BlobId>, NodeError>>,
    ) -> Result<(), Result<Vec<BlobId>, NodeError>> {
        let validator = self.client.lock().await;
        let missing_blob_ids = validator
            .state
            .storage_client()
            .missing_blobs(&blob_ids)
            .await
            .map_err(Into::into);
        sender.send(missing_blob_ids)
    }

    async fn do_event_block_heights(
        self,
        event_ids: Vec<EventId>,
        sender: oneshot::Sender<Result<Vec<Option<BlockHeight>>, NodeError>>,
    ) -> Result<(), Result<Vec<Option<BlockHeight>>, NodeError>> {
        let validator = self.client.lock().await;
        let heights = validator
            .state
            .storage_client()
            .read_event_block_heights(&event_ids)
            .await
            .map_err(Into::into);
        sender.send(heights)
    }
}

/// Polls `condition` until it holds, or gives up after `timeout` of real time.
///
/// The probe-timer tests drive a simulated clock but still have to let a spawned listener
/// run; waiting on an observable counter rather than a fixed sleep keeps them from failing
/// as "never repaired" on a loaded runner, where the real work is a `subscribe` plus a
/// full sync against the storage backend under test.
pub async fn wait_until<F, Fut>(timeout: Duration, mut condition: F) -> bool
where
    F: FnMut() -> Fut,
    Fut: Future<Output = bool>,
{
    let deadline = std::time::Instant::now() + timeout;
    loop {
        if condition().await {
            return true;
        }
        if std::time::Instant::now() >= deadline {
            return false;
        }
        tokio::time::sleep(Duration::from_millis(20)).await;
    }
}

/// Aborts its handle when dropped, so a notification stream the client lets go of is
/// distinguishable from one that is still held.
struct AbortOnDrop(stream::AbortHandle);

impl Drop for AbortOnDrop {
    fn drop(&mut self) {
        self.0.abort();
    }
}

/// A [`ValidatorNodeProvider`] holding the in-process test validator clients.
#[derive(Clone)]
pub struct NodeProvider<S>
where
    S: Storage,
{
    clients: Arc<std::sync::Mutex<Vec<LocalValidatorClient<S>>>>,
    /// How many further calls to build the validator set must fail, so a test can make
    /// `update_notification_streams` return `Err` at a chosen moment.
    failures_left: Arc<std::sync::atomic::AtomicUsize>,
    /// Every attempt to build the validator set, successful or not. Monotonic, so a test
    /// can freeze the clock and assert the listener is PARKED rather than spinning.
    builds: Arc<std::sync::atomic::AtomicUsize>,
    /// Microseconds to advance the clock on each build, simulating a SLOW committee read —
    /// the one await in `update_notification_streams`. A test uses it to prove deadlines
    /// are computed from a clock sampled after that await rather than before it.
    slow_build_micros: Arc<std::sync::atomic::AtomicU64>,
    clock: Option<TestClock>,
}

impl<S> NodeProvider<S>
where
    S: Storage + Clone,
{
    fn all_nodes(&self) -> Vec<LocalValidatorClient<S>> {
        self.clients.lock().unwrap().clone()
    }

    /// The same validators, with a private failure counter. The validators' own export
    /// tasks resolve each other through a provider too, so without this they would race
    /// the chain client for injected failures and a test could pass without ever failing
    /// the update it names.
    fn without_failure_injection(&self) -> Self {
        Self {
            clients: self.clients.clone(),
            failures_left: Arc::default(),
            // A private build counter too, so export-task builds stay out of a test's count.
            builds: Arc::default(),
            slow_build_micros: Arc::default(),
            clock: self.clock.clone(),
        }
    }
}

impl<S> ValidatorNodeProvider for NodeProvider<S>
where
    S: Storage + Clone + Send + Sync + 'static,
{
    type Node = LocalValidatorClient<S>;

    fn make_node(&self, _name: &str) -> Result<Self::Node, NodeError> {
        unimplemented!()
    }

    fn make_nodes_from_list<A>(
        &self,
        validators: impl IntoIterator<Item = (ValidatorPublicKey, A)>,
    ) -> Result<impl Iterator<Item = (ValidatorPublicKey, Self::Node)>, NodeError>
    where
        A: AsRef<str>,
    {
        self.builds
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        let slow = self
            .slow_build_micros
            .load(std::sync::atomic::Ordering::Relaxed);
        if slow > 0 {
            if let Some(clock) = &self.clock {
                clock.add(TimeDelta::from_micros(slow));
            }
        }
        if self
            .failures_left
            .fetch_update(
                std::sync::atomic::Ordering::Relaxed,
                std::sync::atomic::Ordering::Relaxed,
                |n| n.checked_sub(1),
            )
            .is_ok()
        {
            return Err(NodeError::CannotResolveValidatorAddress {
                address: "injected failure".to_string(),
            });
        }
        let list = self.clients.lock().unwrap();
        Ok(validators
            .into_iter()
            .map(|(public_key, address)| {
                list.iter()
                    .find(|client| client.public_key == public_key)
                    .ok_or_else(|| NodeError::CannotResolveValidatorAddress {
                        address: address.as_ref().to_string(),
                    })
                    .map(|client| (public_key, client.clone()))
            })
            .collect::<Result<Vec<_>, _>>()?
            .into_iter())
    }
}

impl<S> FromIterator<LocalValidatorClient<S>> for NodeProvider<S>
where
    S: Storage,
{
    fn from_iter<T>(iter: T) -> Self
    where
        T: IntoIterator<Item = LocalValidatorClient<S>>,
    {
        Self {
            clients: Arc::new(std::sync::Mutex::new(iter.into_iter().collect())),
            failures_left: Arc::default(),
            builds: Arc::default(),
            slow_build_micros: Arc::default(),
            clock: None,
        }
    }
}

// NOTE:
// * To communicate with a quorum of validators, chain clients iterate over a copy of
// `validator_clients` to spawn I/O tasks.
// * When using `LocalValidatorClient`, clients communicate with an exact quorum then stop.
// * Most tests have 1 faulty validator out 4 so that there is exactly only 1 quorum to
// communicate with.
#[allow(missing_docs)]
pub struct TestBuilder<B: StorageBuilder> {
    storage_builder: B,
    pub initial_committee: Committee,
    admin_description: Option<ChainDescription>,
    network_description: Option<NetworkDescription>,
    genesis_storage_builder: GenesisStorageBuilder,
    node_provider: NodeProvider<B::Storage>,
    pub validator_storages: HashMap<ValidatorPublicKey, B::Storage>,
    pub validator_key_pairs: HashMap<ValidatorPublicKey, ValidatorSecretKey>,
    chain_client_storages: Vec<B::Storage>,
    pub chain_owners: BTreeMap<ChainId, AccountOwner>,
    pub signer: TestSigner,
}

/// Builds storage instances of a specific backend for use in tests.
#[async_trait]
pub trait StorageBuilder {
    /// The storage type produced by this builder.
    type Storage: Storage + Clone + Send + Sync + 'static;

    /// Builds a new storage instance.
    async fn build(&mut self) -> Result<Self::Storage, anyhow::Error>;

    /// Returns the test clock shared by all storages built here.
    fn clock(&self) -> &TestClock;
}

#[derive(Default)]
struct GenesisStorageBuilder {
    accounts: Vec<GenesisAccount>,
}

struct GenesisAccount {
    description: ChainDescription,
    public_key: AccountPublicKey,
}

impl GenesisStorageBuilder {
    fn add(&mut self, description: ChainDescription, public_key: AccountPublicKey) {
        self.accounts.push(GenesisAccount {
            description,
            public_key,
        })
    }

    async fn build<S>(&self, storage: S) -> S
    where
        S: Storage + Clone + Send + Sync + 'static,
    {
        for account in &self.accounts {
            storage
                .create_chain(account.description.clone())
                .await
                .unwrap();
        }
        storage
    }
}

/// A chain client wired up to the in-process test validator network.
pub type ChainClient<S> = crate::client::ChainClient<crate::environment::Impl<S, NodeProvider<S>>>;

impl<S: Storage + Clone + Send + Sync + 'static> ChainClient<S> {
    /// Reads the hashed certificate values in descending order from the given hash.
    pub async fn read_confirmed_blocks_downward(
        &self,
        from: CryptoHash,
        limit: u32,
    ) -> anyhow::Result<Vec<Arc<ConfirmedBlock>>> {
        let mut hash = Some(from);
        let mut values = Vec::new();
        for _ in 0..limit {
            let Some(next_hash) = hash else {
                break;
            };
            let value = self.read_confirmed_block(next_hash).await?;
            hash = value.block().header.previous_block_hash;
            values.push(value);
        }
        Ok(values)
    }
}

impl<B> TestBuilder<B>
where
    B: StorageBuilder,
{
    /// The simulated clock every storage here shares, so a test can drive the export queue's
    /// tick in virtual time instead of sleeping through it.
    pub fn clock(&self) -> &TestClock {
        self.storage_builder.clock()
    }

    /// Creates a test setup with `count` validators, `with_faulty_validators` of which are faulty.
    pub async fn new(
        storage_builder: B,
        count: usize,
        with_faulty_validators: usize,
        signer: TestSigner,
    ) -> Result<Self, anyhow::Error> {
        Self::build(
            storage_builder,
            count,
            with_faulty_validators,
            signer,
            None,
            None,
        )
        .await
    }

    /// Creates a test setup like [`TestBuilder::new`], in which every validator also pushes the
    /// blocks it executes to the rest of the committee.
    pub async fn new_with_block_export(
        storage_builder: B,
        count: usize,
        with_faulty_validators: usize,
        signer: TestSigner,
    ) -> Result<Self, anyhow::Error> {
        Self::build(
            storage_builder,
            count,
            with_faulty_validators,
            signer,
            Some(Self::test_block_export_config()),
            None,
        )
        .await
    }

    /// Creates a test setup like [`TestBuilder::new_with_block_export`] with an explicit chain
    /// worker TTL, for asserting that workers expire even while export is enabled.
    pub async fn new_with_block_export_and_ttl(
        storage_builder: B,
        count: usize,
        with_faulty_validators: usize,
        signer: TestSigner,
        chain_worker_ttl: Duration,
    ) -> Result<Self, anyhow::Error> {
        Self::build(
            storage_builder,
            count,
            with_faulty_validators,
            signer,
            Some(Self::test_block_export_config()),
            Some(chain_worker_ttl),
        )
        .await
    }

    /// Creates a test setup like [`TestBuilder::new_with_block_export`], with the export tuned by
    /// the caller. Used to shrink `max_catch_up_blocks` far below its default so that a backlog a
    /// test can actually produce still takes several rounds to drain.
    pub async fn new_with_block_export_config(
        storage_builder: B,
        count: usize,
        with_faulty_validators: usize,
        signer: TestSigner,
        config: crate::BlockExportConfig,
    ) -> Result<Self, anyhow::Error> {
        Self::build(
            storage_builder,
            count,
            with_faulty_validators,
            signer,
            Some(config),
            None,
        )
        .await
    }

    /// The export settings the block-export tests run with: production backoff is measured in
    /// seconds, which would make every test that exercises a failing destination wait it out.
    pub fn test_block_export_config() -> crate::BlockExportConfig {
        crate::BlockExportConfig {
            retry_delay: Duration::from_millis(20),
            max_retry_delay: Duration::from_millis(200),
            ..crate::BlockExportConfig::default()
        }
    }

    async fn build(
        mut storage_builder: B,
        count: usize,
        with_faulty_validators: usize,
        mut signer: TestSigner,
        block_export: Option<crate::BlockExportConfig>,
        chain_worker_ttl: Option<Duration>,
    ) -> Result<Self, anyhow::Error> {
        let mut validators = Vec::new();
        for _ in 0..count {
            let validator_keypair = ValidatorKeypair::generate();
            let account_public_key = signer.generate_new();
            validators.push((validator_keypair, account_public_key));
        }
        let for_committee = validators
            .iter()
            .map(|(validating, account)| (validating.public_key, *account))
            .collect::<Vec<_>>();
        let initial_committee = Committee::make_simple(for_committee);
        // Created up front and filled in below, so that each validator's export tasks can resolve
        // the others through it even though those clients do not exist yet.
        let node_provider = NodeProvider {
            clients: Arc::new(std::sync::Mutex::new(Vec::new())),
            failures_left: Arc::default(),
            builds: Arc::default(),
            slow_build_micros: Arc::default(),
            clock: Some(storage_builder.clock().clone()),
        };
        let mut validator_storages = HashMap::new();
        let mut validator_key_pairs = HashMap::new();
        let mut faulty_validators = HashSet::new();
        for (i, (validator_keypair, _account_public_key)) in validators.into_iter().enumerate() {
            let validator_public_key = validator_keypair.public_key;
            let storage = storage_builder.build().await?;
            let secret_key_copy = validator_keypair.secret_key.copy();
            let config = ChainWorkerConfig {
                nickname: format!("Node {i}"),
                // Export folds progress into the chain state when the worker next saves, so
                // give workers a lifetime instead of dropping them after every request, and fold
                // unthrottled so assertions see progress as it happens.
                ttl: chain_worker_ttl
                    .or_else(|| block_export.is_some().then(|| Duration::from_secs(60))),
                exported_heights_fold_interval: Duration::ZERO,
                ..ChainWorkerConfig::default()
            }
            .with_key_pair(Some(validator_keypair.secret_key));
            let mut state = WorkerState::new(storage.clone(), config, None);
            if let Some(export_config) = block_export.clone() {
                let handle = crate::spawn_block_export_queue(
                    storage.clone(),
                    Arc::new(node_provider.without_failure_injection()),
                    export_config,
                    Some(validator_public_key),
                );
                state = state.with_block_export(handle);
            }
            let mut validator = LocalValidatorClient::new(validator_public_key, state);
            if i < with_faulty_validators {
                faulty_validators.insert(validator_public_key);
                validator.set_fault_type(FaultType::NoChains);
            }
            node_provider.clients.lock().unwrap().push(validator);
            validator_storages.insert(validator_public_key, storage);
            validator_key_pairs.insert(validator_public_key, secret_key_copy);
        }
        tracing::info!(
            "Test will use the following faulty validators: {:?}",
            faulty_validators
        );
        Ok(Self {
            storage_builder,
            initial_committee,
            admin_description: None,
            network_description: None,
            genesis_storage_builder: GenesisStorageBuilder::default(),
            node_provider,
            validator_storages,
            validator_key_pairs,
            chain_client_storages: Vec::new(),
            chain_owners: BTreeMap::new(),
            signer,
        })
    }

    /// Replaces the initial committee's resource control policy.
    pub fn with_policy(mut self, policy: ResourceControlPolicy) -> Self {
        let validators = self.initial_committee.validators().clone();
        self.initial_committee =
            Committee::new(validators, policy).expect("committee votes should not overflow");
        self
    }

    /// Sets the cross-chain message chunk limit on every validator in the test setup.
    pub fn with_cross_chain_message_chunk_limit(self, limit: usize) -> Self {
        let validator_clients = self.node_provider.clients.lock().unwrap();
        for validator in validator_clients.iter() {
            let mut inner = validator.client.try_lock().expect("no contention at setup");
            inner.state.set_cross_chain_message_chunk_limit(limit);
        }
        drop(validator_clients);
        self
    }

    /// Makes the next `count` attempts to build the validator set fail, so a test can
    /// force `update_notification_streams` to return `Err` at a chosen moment.
    pub fn fail_next_validator_set_builds(&self, count: usize) {
        self.node_provider
            .failures_left
            .store(count, std::sync::atomic::Ordering::Relaxed);
    }

    /// How many of the injected validator-set-build failures are still pending, so a
    /// test can count how often a failing update was retried.
    pub fn validator_set_build_failures_left(&self) -> usize {
        self.node_provider
            .failures_left
            .load(std::sync::atomic::Ordering::Relaxed)
    }

    /// Stalls every chain-info query at the validator at `index` until the returned guard
    /// is dropped, which stalls the initial sync a probe runs once `subscribe` succeeds.
    pub async fn stall_chain_info_queries(
        &self,
        index: usize,
    ) -> tokio::sync::OwnedRwLockWriteGuard<()> {
        let node = &self.node_provider.all_nodes()[index];
        node.chain_info_gate_armed
            .store(true, std::sync::atomic::Ordering::Relaxed);
        node.chain_info_gate.clone().write_owned().await
    }

    /// Makes every validator-set build advance the clock, simulating a committee read slow
    /// enough that deadlines computed from a clock sampled BEFORE it are already elapsed.
    pub fn slow_validator_set_builds(&self, by: TimeDelta) {
        self.node_provider
            .slow_build_micros
            .store(by.as_micros(), std::sync::atomic::Ordering::Relaxed);
    }

    /// Every attempt the chain clients have made to build the validator set.
    pub fn validator_set_builds(&self) -> usize {
        self.node_provider
            .builds
            .load(std::sync::atomic::Ordering::Relaxed)
    }

    /// How many times the validator at `index` has been asked to `subscribe`.
    pub async fn subscribe_calls(&self, index: usize) -> usize {
        self.node_provider.all_nodes()[index]
            .subscribe_calls()
            .await
    }

    /// Severs every notification stream on every validator in the test setup, as a
    /// fleet-wide proxy restart does. Clients keep their (now dead) ends of the streams
    /// until they notice and re-subscribe.
    pub async fn disconnect_notification_subscribers(&self) {
        for validator in self.node_provider.all_nodes() {
            validator.disconnect_notification_subscribers().await;
        }
    }

    /// Returns the [`FaultType`] currently configured for the given validator, or `None`
    /// if no validator with that key is in the test setup.
    pub fn fault_type(&self, public_key: &ValidatorPublicKey) -> Option<FaultType> {
        self.node_provider
            .clients
            .lock()
            .unwrap()
            .iter()
            .find(|client| client.public_key == *public_key)
            .map(|client| client.fault_type())
    }

    /// Sets the [`FaultType`] for the validators at the given indexes.
    pub fn set_fault_type(&mut self, indexes: impl AsRef<[usize]>, fault_type: FaultType) {
        let mut faulty_validators = vec![];
        let mut validator_clients = self.node_provider.clients.lock().unwrap();
        for index in indexes.as_ref() {
            let validator = &mut validator_clients[*index];
            validator.set_fault_type(fault_type);
            faulty_validators.push(validator.public_key);
        }
        tracing::info!(
            "Making the following validators {:?}: {:?}",
            fault_type,
            faulty_validators
        );
    }

    /// Creates the root chain with the given `index`, and returns a client for it.
    ///
    /// Root chain 0 is the admin chain and needs to be initialized first, otherwise its balance
    /// is automatically set to zero.
    pub async fn add_root_chain(
        &mut self,
        index: u32,
        balance: Amount,
    ) -> anyhow::Result<ChainClient<B::Storage>> {
        self.add_root_chain_with_ownership(index, balance, ChainOwnership::single)
            .await
    }

    /// Creates the root chain with the given `index` and a genesis ownership built from its
    /// freshly generated owner key, and returns a client for it.
    ///
    /// Root chain 0 is the admin chain and needs to be initialized first, otherwise its balance
    /// is automatically set to zero.
    pub async fn add_root_chain_with_ownership(
        &mut self,
        index: u32,
        balance: Amount,
        make_ownership: impl FnOnce(AccountOwner) -> ChainOwnership,
    ) -> anyhow::Result<ChainClient<B::Storage>> {
        // Make sure the admin chain is initialized.
        if self.admin_description.is_none() && index != 0 {
            Box::pin(self.add_root_chain(0, Amount::ZERO)).await?;
        }
        let origin = ChainOrigin::Root(index);
        let public_key = self.signer.generate_new();
        let open_chain_config = InitialChainConfig {
            ownership: make_ownership(public_key.into()),
            epoch: Epoch(0),
            account: AccountOwner::CHAIN,
            balance,
            application_permissions: ApplicationPermissions::default(),
        };
        let description = ChainDescription::new(origin, open_chain_config, Timestamp::from(0));
        let committee_blob = Blob::new_committee(bcs::to_bytes(&self.initial_committee).unwrap());
        if index == 0 {
            self.admin_description = Some(description.clone());
            self.network_description = Some(NetworkDescription {
                admin_chain_id: description.id(),
                // dummy values to fill the description
                genesis_config_hash: CryptoHash::test_hash("genesis config"),
                genesis_timestamp: Timestamp::from(0),
                genesis_committee_blob_hash: committee_blob.id().hash,
                name: "test network".to_string(),
            });
        }
        // Remember what's in the genesis store for future clients to join.
        self.genesis_storage_builder
            .add(description.clone(), public_key);

        let network_description = self.network_description.as_ref().unwrap();

        for validator in self.node_provider.all_nodes() {
            let storage = self
                .validator_storages
                .get_mut(&validator.public_key)
                .unwrap();
            storage
                .write_network_description(network_description)
                .await
                .expect("writing the NetworkDescription should succeed");
            storage
                .write_blob(&committee_blob)
                .await
                .expect("writing a blob should succeed");
            storage.create_chain(description.clone()).await.unwrap();
        }
        for storage in &mut self.chain_client_storages {
            storage.create_chain(description.clone()).await.unwrap();
        }
        let chain_id = description.id();
        self.chain_owners.insert(chain_id, public_key.into());
        self.make_client(chain_id, None, BlockHeight::ZERO).await
    }

    /// Returns the public key and balance of each genesis root chain.
    pub fn genesis_chains(&self) -> Vec<(AccountPublicKey, Amount)> {
        let mut result = Vec::new();
        for (i, genesis_account) in self.genesis_storage_builder.accounts.iter().enumerate() {
            assert_eq!(
                genesis_account.description.origin(),
                ChainOrigin::Root(i as u32)
            );
            result.push((
                genesis_account.public_key,
                genesis_account.description.config().balance,
            ));
        }
        result
    }

    /// Returns the admin chain's ID, panicking if it has not been initialized.
    pub fn admin_chain_id(&self) -> ChainId {
        self.admin_description
            .as_ref()
            .expect("admin chain not initialized")
            .id()
    }

    /// Returns the admin chain description, if the admin chain has been initialized.
    pub fn admin_description(&self) -> Option<&ChainDescription> {
        self.admin_description.as_ref()
    }

    /// Returns a clone of the node provider backing this test setup.
    pub fn make_node_provider(&self) -> NodeProvider<B::Storage> {
        self.node_provider.clone()
    }

    /// Returns the storage of the validator at `index`, which holds every block that validator
    /// has processed.
    pub fn validator_storage(&mut self, index: usize) -> B::Storage {
        let public_key = self.node(index).public_key;
        self.validator_storages.get(&public_key).unwrap().clone()
    }

    /// Returns a clone of the validator client at the given index.
    pub fn node(&mut self, index: usize) -> LocalValidatorClient<B::Storage> {
        self.node_provider.clients.lock().unwrap()[index].clone()
    }

    /// Builds a fresh storage seeded with the network description and genesis chains.
    pub async fn make_storage(&mut self) -> anyhow::Result<B::Storage> {
        let storage = self.storage_builder.build().await?;
        let network_description = self.network_description.as_ref().unwrap();
        let committee_blob = Blob::new_committee(bcs::to_bytes(&self.initial_committee).unwrap());
        storage
            .write_network_description(network_description)
            .await
            .expect("writing the NetworkDescription should succeed");
        storage
            .write_blob(&committee_blob)
            .await
            .expect("writing a blob should succeed");
        Ok(self.genesis_storage_builder.build(storage).await)
    }

    /// Creates a chain client for the given chain with the given client options.
    pub async fn make_client_with_options(
        &mut self,
        chain_id: ChainId,
        block_hash: Option<CryptoHash>,
        block_height: BlockHeight,
        options: chain_client::Options,
        follow_only: bool,
    ) -> anyhow::Result<ChainClient<B::Storage>> {
        // Note that new clients are only given the genesis store: they must figure out
        // the rest by asking validators.
        let storage = self.make_storage().await?;
        self.chain_client_storages.push(storage.clone());
        let mode = if follow_only {
            crate::client::ListeningMode::FollowChain
        } else {
            crate::client::ListeningMode::FullChain
        };
        let client = Arc::new(Client::new(
            crate::environment::Impl {
                network: self.make_node_provider(),
                storage,
                signer: self.signer.clone(),
                wallet: TestWallet::default(),
            },
            self.admin_chain_id(),
            false,
            [(chain_id, mode)],
            format!("Client node for {chain_id:.8}"),
            Some(Duration::from_secs(30)),
            Some(Duration::from_secs(1)),
            1000,
            options,
            DEFAULT_BLOCK_CACHE_SIZE,
            DEFAULT_EXECUTION_STATE_CACHE_SIZE,
            &crate::client::RequestsSchedulerConfig::default(),
        ));
        Ok(client.create_chain_client(
            chain_id,
            block_hash,
            block_height,
            &None,
            self.chain_owners.get(&chain_id).copied(),
            None,
            follow_only,
        ))
    }

    /// Creates a chain client for the given chain with default test options.
    pub async fn make_client(
        &mut self,
        chain_id: ChainId,
        block_hash: Option<CryptoHash>,
        block_height: BlockHeight,
    ) -> anyhow::Result<ChainClient<B::Storage>> {
        self.make_client_with_options(
            chain_id,
            block_hash,
            block_height,
            chain_client::Options::test_default(),
            false,
        )
        .await
    }

    /// Tries to find a (confirmation) certificate for the given chain_id and block height.
    pub async fn check_that_validators_have_certificate(
        &self,
        chain_id: ChainId,
        block_height: BlockHeight,
        target_count: usize,
    ) -> Option<ConfirmedBlockCertificate> {
        let query = ChainInfoQuery::new(chain_id)
            .with_sent_certificate_hashes_by_heights(vec![block_height]);
        let mut count = 0;
        let mut certificate = None;
        for validator in self.node_provider.all_nodes() {
            if let Ok(response) = validator.handle_chain_info_query(query.clone()).await {
                if response.check(validator.public_key).is_ok() {
                    let ChainInfo {
                        mut requested_sent_certificate_hashes,
                        ..
                    } = *response.info;
                    debug_assert!(requested_sent_certificate_hashes.len() <= 1);
                    if let Some(cert_hash) = requested_sent_certificate_hashes.pop() {
                        if let Ok(cert) = validator.download_certificate(cert_hash).await {
                            if cert.inner().block().header.chain_id == chain_id
                                && cert.inner().block().header.height == block_height
                            {
                                cert.check(&self.initial_committee).unwrap();
                                count += 1;
                                certificate = Some(cert);
                            }
                        }
                    }
                }
            }
        }
        assert!(count >= target_count);
        certificate
    }

    /// Tries to find a (confirmation) certificate for the given chain_id and block height, and are
    /// in the expected round.
    pub async fn check_that_validators_are_in_round(
        &self,
        chain_id: ChainId,
        block_height: BlockHeight,
        round: Round,
        target_count: usize,
    ) {
        let query = ChainInfoQuery::new(chain_id);
        let mut count = 0;
        for validator in self.node_provider.all_nodes() {
            if let Ok(response) = validator.handle_chain_info_query(query.clone()).await {
                if response.info.manager.current_round == round
                    && response.info.next_block_height == block_height
                    && response.check(validator.public_key).is_ok()
                {
                    count += 1;
                }
            }
        }
        assert!(count >= target_count);
    }

    /// Returns how far the validator at `index` believes it has exported the given chain to each
    /// of the other validators.
    pub async fn exported_heights(
        &self,
        index: usize,
        chain_id: ChainId,
    ) -> BTreeMap<ValidatorPublicKey, BlockHeight> {
        let validator = self.node_provider.all_nodes()[index].clone();
        let guard = validator.client.lock().await;
        let chain = guard.state.chain_state_view(chain_id).await.unwrap();
        chain.exported_heights.get().clone().into()
    }

    /// Returns how many chain workers the validator at `index` currently has resident.
    pub async fn resident_chain_workers(&self, index: usize) -> usize {
        let validator = self.node_provider.all_nodes()[index].clone();
        let guard = validator.client.lock().await;
        guard.state.resident_chain_worker_count()
    }

    /// Returns the next block height the validator at `index` has for the given chain.
    pub async fn next_block_height(&self, index: usize, chain_id: ChainId) -> BlockHeight {
        let validator = self.node_provider.all_nodes()[index].clone();
        let response = validator
            .handle_chain_info_query(ChainInfoQuery::new(chain_id))
            .await
            .unwrap();
        response.info.next_block_height
    }

    /// Panics if any validator has a nonempty outbox for the given chain.
    pub async fn check_that_validators_have_empty_outboxes(&self, chain_id: ChainId) {
        for validator in self.node_provider.all_nodes() {
            let guard = validator.client.lock().await;
            let chain = guard.state.chain_state_view(chain_id).await.unwrap();
            assert_eq!(chain.outboxes.indices().await.unwrap(), []);
        }
    }
}

#[cfg(feature = "rocksdb")]
/// Limit concurrency for RocksDB tests to avoid "too many open files" errors.
static ROCKS_DB_SEMAPHORE: Semaphore = Semaphore::const_new(5);

/// State shared by every [`StorageBuilder`] in this module. The actual
/// database type varies, so `build_storage` is generic over it.
#[derive(Default)]
struct CommonStorageBuilder {
    namespace: String,
    instance_counter: usize,
    wasm_runtime: Option<WasmRuntime>,
    clock: TestClock,
}

impl CommonStorageBuilder {
    fn with_wasm_runtime(wasm_runtime: impl Into<Option<WasmRuntime>>) -> Self {
        Self {
            wasm_runtime: wasm_runtime.into(),
            ..Self::default()
        }
    }

    async fn build_storage<DB>(
        &mut self,
        config: DB::Config,
    ) -> anyhow::Result<DbStorage<DB, TestClock>>
    where
        DB: TestKeyValueDatabase + Clone + Send + Sync + 'static,
        DB::Store: KeyValueStore + Clone + Send + Sync + 'static,
        DB::Error: std::error::Error + Send + Sync + 'static,
    {
        self.instance_counter += 1;
        if self.namespace.is_empty() {
            self.namespace = generate_test_namespace();
        }
        let namespace = format!("{}_{}", self.namespace, self.instance_counter);
        Ok(
            DbStorage::new_for_testing(config, &namespace, self.wasm_runtime, self.clock.clone())
                .await?,
        )
    }
}

/// A [`StorageBuilder`] backed by in-memory storage.
#[derive(Default)]
pub struct MemoryStorageBuilder {
    inner: CommonStorageBuilder,
}

impl MemoryStorageBuilder {
    /// Creates a [`MemoryStorageBuilder`] that uses the specified [`WasmRuntime`] to run Wasm
    /// applications.
    pub fn with_wasm_runtime(wasm_runtime: impl Into<Option<WasmRuntime>>) -> Self {
        Self {
            inner: CommonStorageBuilder::with_wasm_runtime(wasm_runtime),
        }
    }
}

#[async_trait]
impl StorageBuilder for MemoryStorageBuilder {
    type Storage = DbStorage<MemoryDatabase, TestClock>;

    async fn build(&mut self) -> Result<Self::Storage, anyhow::Error> {
        let config = MemoryDatabase::new_test_config().await?;
        self.inner.build_storage::<MemoryDatabase>(config).await
    }

    fn clock(&self) -> &TestClock {
        &self.inner.clock
    }
}

#[cfg(feature = "rocksdb")]
/// A [`StorageBuilder`] backed by RocksDB storage.
pub struct RocksDbStorageBuilder {
    inner: CommonStorageBuilder,
    _permit: SemaphorePermit<'static>,
}

#[cfg(feature = "rocksdb")]
impl RocksDbStorageBuilder {
    /// Creates a [`RocksDbStorageBuilder`], acquiring a concurrency permit.
    pub async fn new() -> Self {
        Self {
            inner: CommonStorageBuilder::default(),
            _permit: ROCKS_DB_SEMAPHORE.acquire().await.unwrap(),
        }
    }

    /// Creates a [`RocksDbStorageBuilder`] that uses the specified [`WasmRuntime`] to run Wasm
    /// applications.
    #[cfg(any(feature = "wasmer", feature = "wasmtime"))]
    pub async fn with_wasm_runtime(wasm_runtime: impl Into<Option<WasmRuntime>>) -> Self {
        Self {
            inner: CommonStorageBuilder::with_wasm_runtime(wasm_runtime),
            _permit: ROCKS_DB_SEMAPHORE.acquire().await.unwrap(),
        }
    }
}

#[cfg(feature = "rocksdb")]
#[async_trait]
impl StorageBuilder for RocksDbStorageBuilder {
    type Storage = DbStorage<RocksDbDatabase, TestClock>;

    async fn build(&mut self) -> Result<Self::Storage, anyhow::Error> {
        let config = RocksDbDatabase::new_test_config().await?;
        self.inner.build_storage::<RocksDbDatabase>(config).await
    }

    fn clock(&self) -> &TestClock {
        &self.inner.clock
    }
}

#[cfg(all(not(target_arch = "wasm32"), feature = "storage-service"))]
/// A [`StorageBuilder`] backed by the storage service.
#[derive(Default)]
pub struct ServiceStorageBuilder {
    inner: CommonStorageBuilder,
}

#[cfg(all(not(target_arch = "wasm32"), feature = "storage-service"))]
impl ServiceStorageBuilder {
    /// Creates a `ServiceStorage`.
    pub fn new() -> Self {
        Self::with_wasm_runtime(None)
    }

    /// Creates a `ServiceStorage` with the given Wasm runtime.
    pub fn with_wasm_runtime(wasm_runtime: impl Into<Option<WasmRuntime>>) -> Self {
        Self {
            inner: CommonStorageBuilder::with_wasm_runtime(wasm_runtime),
        }
    }
}

#[cfg(all(not(target_arch = "wasm32"), feature = "storage-service"))]
#[async_trait]
impl StorageBuilder for ServiceStorageBuilder {
    type Storage = DbStorage<StorageServiceDatabase, TestClock>;

    async fn build(&mut self) -> anyhow::Result<Self::Storage> {
        let config = StorageServiceDatabase::new_test_config().await?;
        self.inner
            .build_storage::<StorageServiceDatabase>(config)
            .await
    }

    fn clock(&self) -> &TestClock {
        &self.inner.clock
    }
}

#[cfg(feature = "scylladb")]
/// A [`StorageBuilder`] backed by ScyllaDB storage.
#[derive(Default)]
pub struct ScyllaDbStorageBuilder {
    inner: CommonStorageBuilder,
}

#[cfg(feature = "scylladb")]
impl ScyllaDbStorageBuilder {
    /// Creates a [`ScyllaDbStorageBuilder`] that uses the specified [`WasmRuntime`] to run Wasm
    /// applications.
    pub fn with_wasm_runtime(wasm_runtime: impl Into<Option<WasmRuntime>>) -> Self {
        Self {
            inner: CommonStorageBuilder::with_wasm_runtime(wasm_runtime),
        }
    }
}

#[cfg(feature = "scylladb")]
#[async_trait]
impl StorageBuilder for ScyllaDbStorageBuilder {
    type Storage = DbStorage<ScyllaDbDatabase, TestClock>;

    async fn build(&mut self) -> Result<Self::Storage, anyhow::Error> {
        let config = ScyllaDbDatabase::new_test_config().await?;
        self.inner.build_storage::<ScyllaDbDatabase>(config).await
    }

    fn clock(&self) -> &TestClock {
        &self.inner.clock
    }
}

/// Helpers for asserting on [`ClientOutcome`] results in tests.
pub trait ClientOutcomeResultExt<T, E> {
    /// Unwraps the result and panics if it's not `Committed`.
    /// Use this when you expect the operation to succeed without conflicts.
    fn unwrap_ok_committed(self) -> T;

    /// Unwraps the result, accepting both `Committed` and `Conflict` outcomes.
    /// Returns the committed value or the conflicting certificate (boxed).
    fn unwrap_ok_or_conflict(self) -> Result<T, Box<ConfirmedBlockCertificate>>;
}

impl<T, E: std::fmt::Debug> ClientOutcomeResultExt<T, E> for Result<ClientOutcome<T>, E> {
    fn unwrap_ok_committed(self) -> T {
        match self.unwrap() {
            ClientOutcome::Committed(t) => t,
            ClientOutcome::WaitForTimeout(timeout) => {
                panic!("unexpected timeout: {timeout}")
            }
            ClientOutcome::Conflict(certificate) => {
                panic!("unexpected conflict: {}", certificate.hash())
            }
        }
    }

    fn unwrap_ok_or_conflict(self) -> Result<T, Box<ConfirmedBlockCertificate>> {
        match self.unwrap() {
            ClientOutcome::Committed(t) => Ok(t),
            ClientOutcome::Conflict(certificate) => Err(certificate),
            ClientOutcome::WaitForTimeout(timeout) => {
                panic!("unexpected timeout: {timeout}")
            }
        }
    }
}