Skip to main content

linera_core/unit_tests/
test_utils.rs

1// Copyright (c) Zefchain Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4#![allow(clippy::cast_possible_truncation)]
5
6use std::{
7    collections::{BTreeMap, HashMap, HashSet},
8    sync::Arc,
9    time::Duration,
10    vec,
11};
12
13use async_trait::async_trait;
14use futures::{
15    future::Either,
16    lock::{Mutex, MutexGuard},
17    Future,
18};
19use linera_base::{
20    crypto::{
21        AccountPublicKey, CryptoHash, ValidatorKeypair, ValidatorPublicKey, ValidatorSecretKey,
22    },
23    data_types::*,
24    identifiers::{AccountOwner, BlobId, ChainId, EventId},
25    ownership::ChainOwnership,
26};
27use linera_chain::{
28    data_types::BlockProposal,
29    types::{
30        CertificateKind, Certified, ConfirmedBlock, ConfirmedBlockCertificate, GenericCertificate,
31        LiteCertificate, Timeout, ValidatedBlock, ValidatedBlockCertificate,
32    },
33};
34use linera_execution::{committee::Committee, ResourceControlPolicy, WasmRuntime};
35use linera_storage::{Arc as CacheArc, DbStorage, ResultReadCertificates, Storage, TestClock};
36#[cfg(all(not(target_arch = "wasm32"), feature = "storage-service"))]
37use linera_storage_service::client::StorageServiceDatabase;
38use linera_version::VersionInfo;
39#[cfg(feature = "scylladb")]
40use linera_views::scylla_db::ScyllaDbDatabase;
41use linera_views::{
42    memory::MemoryDatabase,
43    random::generate_test_namespace,
44    store::{KeyValueStore, TestKeyValueDatabase},
45};
46use tokio::sync::oneshot;
47use tokio_stream::wrappers::UnboundedReceiverStream;
48#[cfg(feature = "rocksdb")]
49use {
50    linera_views::rocks_db::RocksDbDatabase,
51    tokio::sync::{Semaphore, SemaphorePermit},
52};
53
54use crate::{
55    chain_worker::ChainWorkerConfig,
56    client::{chain_client, Client},
57    data_types::*,
58    environment::{TestSigner, TestWallet},
59    node::{
60        CrossChainMessageDelivery, NodeError, NotificationStream, ValidatorNode,
61        ValidatorNodeProvider,
62    },
63    notifier::ChannelNotifier,
64    worker::{
65        Notification, ProcessableCertificate, WorkerState, DEFAULT_BLOCK_CACHE_SIZE,
66        DEFAULT_EXECUTION_STATE_CACHE_SIZE,
67    },
68};
69
70/// The kind of misbehavior a test validator simulates.
71#[derive(Debug, PartialEq, Clone, Copy)]
72#[allow(missing_docs)]
73pub enum FaultType {
74    Honest,
75    Offline,
76    OfflineWithInfo,
77    NoChains,
78    DontSendConfirmVote,
79    DontProcessValidated,
80    DontSendValidateVote,
81}
82
83/// A validator used for testing. "Faulty" validators ignore block proposals (but not
84/// certificates or info queries) and have the wrong initial balance for all chains.
85///
86/// All methods are executed in spawned Tokio tasks, so that canceling a client task doesn't cause
87/// the validator's tasks to be canceled: In a real network, a validator also wouldn't cancel
88/// tasks if the client stopped waiting for the response.
89struct LocalValidator<S>
90where
91    S: Storage,
92{
93    state: WorkerState<S>,
94    notifier: Arc<ChannelNotifier<Notification>>,
95}
96
97/// A client used by tests to talk to an in-process `LocalValidator`.
98#[derive(Clone)]
99pub struct LocalValidatorClient<S>
100where
101    S: Storage,
102{
103    public_key: ValidatorPublicKey,
104    client: Arc<Mutex<LocalValidator<S>>>,
105    fault_type: FaultType,
106}
107
108impl<S> ValidatorNode for LocalValidatorClient<S>
109where
110    S: Storage + Clone + Send + Sync + 'static,
111{
112    type NotificationStream = NotificationStream;
113
114    fn address(&self) -> String {
115        format!("local:{}", self.public_key)
116    }
117
118    async fn handle_block_proposal(
119        &self,
120        proposal: BlockProposal,
121    ) -> Result<ChainInfoResponse, NodeError> {
122        self.spawn_and_receive(move |validator, sender| {
123            validator.do_handle_block_proposal(proposal, sender)
124        })
125        .await
126    }
127
128    async fn handle_lite_certificate(
129        &self,
130        certificate: LiteCertificate<'_>,
131        _delivery: CrossChainMessageDelivery,
132    ) -> Result<ChainInfoResponse, NodeError> {
133        let certificate = certificate.cloned();
134        self.spawn_and_receive(move |validator, sender| {
135            validator.do_handle_lite_certificate(certificate, sender)
136        })
137        .await
138    }
139
140    async fn handle_timeout_certificate(
141        &self,
142        certificate: GenericCertificate<Timeout>,
143    ) -> Result<ChainInfoResponse, NodeError> {
144        self.spawn_and_receive(move |validator, sender| {
145            validator.do_handle_certificate::<Timeout>(certificate, sender)
146        })
147        .await
148    }
149
150    async fn handle_validated_certificate(
151        &self,
152        certificate: ValidatedBlockCertificate,
153    ) -> Result<ChainInfoResponse, NodeError> {
154        self.spawn_and_receive(move |validator, sender| {
155            validator.do_handle_certificate::<ValidatedBlock>(certificate, sender)
156        })
157        .await
158    }
159
160    async fn handle_confirmed_certificate(
161        &self,
162        certificate: CacheArc<ConfirmedBlockCertificate>,
163        _delivery: CrossChainMessageDelivery,
164    ) -> Result<ChainInfoResponse, NodeError> {
165        self.spawn_and_receive(move |validator, sender| {
166            validator.do_handle_certificate::<ConfirmedBlock>(
167                CacheArc::unwrap_or_clone(certificate),
168                sender,
169            )
170        })
171        .await
172    }
173
174    async fn handle_chain_info_query(
175        &self,
176        query: ChainInfoQuery,
177    ) -> Result<ChainInfoResponse, NodeError> {
178        self.spawn_and_receive(move |validator, sender| {
179            validator.do_handle_chain_info_query(query, sender)
180        })
181        .await
182    }
183
184    async fn subscribe(&self, chains: Vec<ChainId>) -> Result<NotificationStream, NodeError> {
185        self.spawn_and_receive(move |validator, sender| validator.do_subscribe(chains, sender))
186            .await
187    }
188
189    async fn get_version_info(&self) -> Result<VersionInfo, NodeError> {
190        Ok(Default::default())
191    }
192
193    async fn get_network_description(&self) -> Result<NetworkDescription, NodeError> {
194        Ok(self
195            .client
196            .lock()
197            .await
198            .state
199            .storage_client()
200            .read_network_description()
201            .await
202            .transpose()
203            .ok_or_else(|| NodeError::ViewError {
204                error: "missing NetworkDescription".to_owned(),
205            })??)
206    }
207
208    async fn upload_blob(&self, content: BlobContent) -> Result<BlobId, NodeError> {
209        self.spawn_and_receive(move |validator, sender| validator.do_upload_blob(content, sender))
210            .await
211    }
212
213    async fn download_blob(&self, blob_id: BlobId) -> Result<BlobContent, NodeError> {
214        self.spawn_and_receive(move |validator, sender| validator.do_download_blob(blob_id, sender))
215            .await
216    }
217
218    async fn download_blobs(
219        &self,
220        blob_ids: Vec<BlobId>,
221    ) -> Result<crate::node::BlobStream, NodeError> {
222        let this = self.clone();
223        let stream = futures::stream::unfold(blob_ids.into_iter(), move |mut iter| {
224            let this = this.clone();
225            async move {
226                let blob_id = iter.next()?;
227                let result = this.download_blob(blob_id).await;
228                Some((result, iter))
229            }
230        });
231        Ok(Box::pin(stream))
232    }
233
234    async fn download_pending_blob(
235        &self,
236        chain_id: ChainId,
237        blob_id: BlobId,
238    ) -> Result<BlobContent, NodeError> {
239        self.spawn_and_receive(move |validator, sender| {
240            validator.do_download_pending_blob(chain_id, blob_id, sender)
241        })
242        .await
243    }
244
245    async fn handle_pending_blob(
246        &self,
247        chain_id: ChainId,
248        blob: BlobContent,
249    ) -> Result<ChainInfoResponse, NodeError> {
250        self.spawn_and_receive(move |validator, sender| {
251            validator.do_handle_pending_blob(chain_id, blob, sender)
252        })
253        .await
254    }
255
256    async fn download_certificate(
257        &self,
258        hash: CryptoHash,
259    ) -> Result<ConfirmedBlockCertificate, NodeError> {
260        self.spawn_and_receive(move |validator, sender| {
261            validator.do_download_certificate(hash, sender)
262        })
263        .await
264    }
265
266    async fn download_certificates(
267        &self,
268        hashes: Vec<CryptoHash>,
269    ) -> Result<Vec<ConfirmedBlockCertificate>, NodeError> {
270        self.spawn_and_receive(move |validator, sender| {
271            validator.do_download_certificates(hashes, sender)
272        })
273        .await
274    }
275
276    async fn download_certificates_by_heights(
277        &self,
278        chain_id: ChainId,
279        heights: Vec<BlockHeight>,
280    ) -> Result<Vec<ConfirmedBlockCertificate>, NodeError> {
281        self.spawn_and_receive(move |validator, sender| {
282            validator.do_download_certificates_by_heights(chain_id, heights, sender)
283        })
284        .await
285    }
286
287    async fn blob_last_used_by(&self, blob_id: BlobId) -> Result<CryptoHash, NodeError> {
288        self.spawn_and_receive(move |validator, sender| {
289            validator.do_blob_last_used_by(blob_id, sender)
290        })
291        .await
292    }
293
294    async fn blob_last_used_by_certificate(
295        &self,
296        blob_id: BlobId,
297    ) -> Result<ConfirmedBlockCertificate, NodeError> {
298        self.spawn_and_receive(move |validator, sender| {
299            validator.do_blob_last_used_by_certificate(blob_id, sender)
300        })
301        .await
302    }
303
304    async fn missing_blob_ids(&self, blob_ids: Vec<BlobId>) -> Result<Vec<BlobId>, NodeError> {
305        self.spawn_and_receive(move |validator, sender| {
306            validator.do_missing_blob_ids(blob_ids, sender)
307        })
308        .await
309    }
310
311    async fn event_block_heights(
312        &self,
313        event_ids: Vec<EventId>,
314    ) -> Result<Vec<Option<BlockHeight>>, NodeError> {
315        self.spawn_and_receive(move |validator, sender| {
316            validator.do_event_block_heights(event_ids, sender)
317        })
318        .await
319    }
320
321    async fn get_shard_info(
322        &self,
323        _chain_id: ChainId,
324    ) -> Result<crate::data_types::ShardInfo, NodeError> {
325        // For test purposes, return a dummy shard info
326        Ok(crate::data_types::ShardInfo {
327            shard_id: 0,
328            total_shards: 1,
329        })
330    }
331}
332
333impl<S> LocalValidatorClient<S>
334where
335    S: Storage + Clone + Send + Sync + 'static,
336{
337    fn new(public_key: ValidatorPublicKey, state: WorkerState<S>) -> Self {
338        let client = LocalValidator {
339            state,
340            notifier: Arc::new(ChannelNotifier::default()),
341        };
342        Self {
343            public_key,
344            client: Arc::new(Mutex::new(client)),
345            fault_type: FaultType::Honest,
346        }
347    }
348
349    /// Returns the validator's public key.
350    pub fn name(&self) -> ValidatorPublicKey {
351        self.public_key
352    }
353
354    /// Returns the validator's currently configured [`FaultType`].
355    pub fn fault_type(&self) -> FaultType {
356        self.fault_type
357    }
358
359    fn set_fault_type(&mut self, fault_type: FaultType) {
360        self.fault_type = fault_type;
361    }
362
363    /// Obtains the basic `ChainInfo` data for the local validator chain, with chain manager values.
364    pub async fn chain_info_with_manager_values(
365        &mut self,
366        chain_id: ChainId,
367    ) -> Result<Box<ChainInfo>, NodeError> {
368        let query = ChainInfoQuery::new(chain_id).with_manager_values();
369        let response = self.handle_chain_info_query(query).await?;
370        Ok(response.info)
371    }
372
373    /// Executes the future produced by `f` in a new thread in a new Tokio runtime.
374    /// Returns the value that the future puts into the sender.
375    async fn spawn_and_receive<F, R, T>(&self, f: F) -> T
376    where
377        T: Send + 'static,
378        R: Future<Output = Result<(), T>> + Send,
379        F: FnOnce(Self, oneshot::Sender<T>) -> R + Send + 'static,
380    {
381        let validator = self.clone();
382        let (sender, receiver) = oneshot::channel();
383        tokio::spawn(async move {
384            if f(validator, sender).await.is_err() {
385                tracing::debug!("result could not be sent");
386            }
387        });
388        receiver.await.unwrap()
389    }
390
391    async fn do_handle_block_proposal(
392        self,
393        proposal: BlockProposal,
394        sender: oneshot::Sender<Result<ChainInfoResponse, NodeError>>,
395    ) -> Result<(), Result<ChainInfoResponse, NodeError>> {
396        let result = match self.fault_type {
397            FaultType::Offline | FaultType::OfflineWithInfo => Err(NodeError::ClientIoError {
398                error: "offline".to_string(),
399            }),
400            FaultType::NoChains => Err(NodeError::InactiveChain(proposal.content.block.chain_id)),
401            FaultType::DontSendValidateVote
402            | FaultType::Honest
403            | FaultType::DontSendConfirmVote
404            | FaultType::DontProcessValidated => {
405                let (response_result, _actions) = self
406                    .client
407                    .lock()
408                    .await
409                    .state
410                    .handle_block_proposal(proposal)
411                    .await;
412                let result = response_result.map_err(NodeError::from);
413                if self.fault_type == FaultType::DontSendValidateVote {
414                    Err(NodeError::ClientIoError {
415                        error: "refusing to validate".to_string(),
416                    })
417                } else {
418                    result
419                }
420            }
421        };
422        // In a local node cross-chain messages can't get lost, so we can ignore the actions here.
423        sender.send(result)
424    }
425
426    async fn do_handle_lite_certificate(
427        self,
428        certificate: LiteCertificate<'_>,
429        sender: oneshot::Sender<Result<ChainInfoResponse, NodeError>>,
430    ) -> Result<(), Result<ChainInfoResponse, NodeError>> {
431        let client = self.client.clone();
432        let validator = client.lock().await;
433        let result = async move {
434            match validator.state.full_certificate(certificate).await? {
435                Either::Left(confirmed) => {
436                    self.do_handle_certificate_internal::<ConfirmedBlock>(confirmed, &validator)
437                        .await
438                }
439                Either::Right(validated) => {
440                    self.do_handle_certificate_internal::<ValidatedBlock>(validated, &validator)
441                        .await
442                }
443            }
444        }
445        .await;
446        sender.send(result)
447    }
448
449    async fn do_handle_certificate_internal<T: ProcessableCertificate>(
450        &self,
451        certificate: T::Certificate,
452        validator: &MutexGuard<'_, LocalValidator<S>>,
453    ) -> Result<ChainInfoResponse, NodeError> {
454        match self.fault_type {
455            FaultType::DontProcessValidated if T::KIND == CertificateKind::Validated => {
456                Err(NodeError::ClientIoError {
457                    error: "refusing to process validated block".to_string(),
458                })
459            }
460            FaultType::NoChains => Err(NodeError::InactiveChain(certificate.value().chain_id())),
461            FaultType::Honest
462            | FaultType::DontSendConfirmVote
463            | FaultType::DontProcessValidated
464            | FaultType::DontSendValidateVote => {
465                let result = validator
466                    .state
467                    .fully_handle_certificate_with_notifications(certificate, &validator.notifier)
468                    .await
469                    .map_err(Into::into);
470                if T::KIND == CertificateKind::Validated
471                    && self.fault_type == FaultType::DontSendConfirmVote
472                {
473                    Err(NodeError::ClientIoError {
474                        error: "refusing to confirm".to_string(),
475                    })
476                } else {
477                    result
478                }
479            }
480            FaultType::Offline | FaultType::OfflineWithInfo => Err(NodeError::ClientIoError {
481                error: "offline".to_string(),
482            }),
483        }
484    }
485
486    async fn do_handle_certificate<T: ProcessableCertificate>(
487        self,
488        certificate: T::Certificate,
489        sender: oneshot::Sender<Result<ChainInfoResponse, NodeError>>,
490    ) -> Result<(), Result<ChainInfoResponse, NodeError>> {
491        let validator = self.client.lock().await;
492        let result = self
493            .do_handle_certificate_internal::<T>(certificate, &validator)
494            .await;
495        sender.send(result)
496    }
497
498    async fn do_handle_chain_info_query(
499        self,
500        query: ChainInfoQuery,
501        sender: oneshot::Sender<Result<ChainInfoResponse, NodeError>>,
502    ) -> Result<(), Result<ChainInfoResponse, NodeError>> {
503        let validator = self.client.lock().await;
504        let result = match self.fault_type {
505            FaultType::Offline => Err(NodeError::ClientIoError {
506                error: "offline".to_string(),
507            }),
508            FaultType::NoChains => Err(NodeError::InactiveChain(query.chain_id)),
509            FaultType::Honest
510            | FaultType::DontSendConfirmVote
511            | FaultType::DontProcessValidated
512            | FaultType::DontSendValidateVote
513            | FaultType::OfflineWithInfo => validator
514                .state
515                .handle_chain_info_query(query)
516                .await
517                .map_err(Into::into),
518        };
519        sender.send(result)
520    }
521
522    async fn do_subscribe(
523        self,
524        chains: Vec<ChainId>,
525        sender: oneshot::Sender<Result<NotificationStream, NodeError>>,
526    ) -> Result<(), Result<NotificationStream, NodeError>> {
527        let validator = self.client.lock().await;
528        let rx = validator.notifier.subscribe(chains);
529        let stream: NotificationStream = Box::pin(UnboundedReceiverStream::new(rx));
530        sender.send(Ok(stream))
531    }
532
533    async fn do_upload_blob(
534        self,
535        content: BlobContent,
536        sender: oneshot::Sender<Result<BlobId, NodeError>>,
537    ) -> Result<(), Result<BlobId, NodeError>> {
538        let validator = self.client.lock().await;
539        let blob = Blob::new(content);
540        let id = blob.id();
541        let storage = validator.state.storage_client();
542        let result = match storage.maybe_write_blobs(&[blob]).await {
543            Ok(has_state) if has_state.first() == Some(&true) => Ok(id),
544            Ok(_) => Err(NodeError::BlobsNotFound(vec![id])),
545            Err(error) => Err(error.into()),
546        };
547        sender.send(result)
548    }
549
550    async fn do_download_blob(
551        self,
552        blob_id: BlobId,
553        sender: oneshot::Sender<Result<BlobContent, NodeError>>,
554    ) -> Result<(), Result<BlobContent, NodeError>> {
555        let validator = self.client.lock().await;
556        let blob = validator
557            .state
558            .storage_client()
559            .read_blob(blob_id)
560            .await
561            .map_err(Into::into);
562        let blob = match blob {
563            Ok(blob) => blob.ok_or_else(|| NodeError::BlobsNotFound(vec![blob_id])),
564            Err(error) => Err(error),
565        };
566        sender.send(blob.map(|blob| CacheArc::unwrap_or_clone(blob).into_content()))
567    }
568
569    async fn do_download_pending_blob(
570        self,
571        chain_id: ChainId,
572        blob_id: BlobId,
573        sender: oneshot::Sender<Result<BlobContent, NodeError>>,
574    ) -> Result<(), Result<BlobContent, NodeError>> {
575        let validator = self.client.lock().await;
576        let result = validator
577            .state
578            .download_pending_blob(chain_id, blob_id)
579            .await
580            .map_err(Into::into);
581        sender.send(result.map(|blob| blob.content().clone()))
582    }
583
584    async fn do_handle_pending_blob(
585        self,
586        chain_id: ChainId,
587        blob: BlobContent,
588        sender: oneshot::Sender<Result<ChainInfoResponse, NodeError>>,
589    ) -> Result<(), Result<ChainInfoResponse, NodeError>> {
590        let validator = self.client.lock().await;
591        let result = validator
592            .state
593            .handle_pending_blob(chain_id, Blob::new(blob))
594            .await
595            .map_err(Into::into);
596        sender.send(result)
597    }
598
599    async fn do_download_certificate(
600        self,
601        hash: CryptoHash,
602        sender: oneshot::Sender<Result<ConfirmedBlockCertificate, NodeError>>,
603    ) -> Result<(), Result<ConfirmedBlockCertificate, NodeError>> {
604        let validator = self.client.lock().await;
605        let certificate = validator
606            .state
607            .storage_client()
608            .read_certificate(hash)
609            .await
610            .map_err(Into::into);
611
612        let certificate = match certificate {
613            Err(error) => Err(error),
614            Ok(entry) => match entry {
615                Some(certificate) => Ok(CacheArc::unwrap_or_clone(certificate)),
616                None => {
617                    panic!("Missing certificate: {hash}");
618                }
619            },
620        };
621
622        sender.send(certificate)
623    }
624
625    async fn do_download_certificates(
626        self,
627        hashes: Vec<CryptoHash>,
628        sender: oneshot::Sender<Result<Vec<ConfirmedBlockCertificate>, NodeError>>,
629    ) -> Result<(), Result<Vec<ConfirmedBlockCertificate>, NodeError>> {
630        let validator = self.client.lock().await;
631        let certificates = validator
632            .state
633            .storage_client()
634            .read_certificates(&hashes)
635            .await
636            .map_err(Into::into);
637
638        let certificates = match certificates {
639            Err(error) => Err(error),
640            Ok(certificates) => match ResultReadCertificates::new(certificates, hashes) {
641                ResultReadCertificates::Certificates(certificates) => Ok(certificates),
642                ResultReadCertificates::InvalidHashes(hashes) => {
643                    panic!("Missing certificates: {hashes:?}")
644                }
645            },
646        };
647
648        sender.send(certificates)
649    }
650
651    async fn do_download_certificates_by_heights(
652        self,
653        chain_id: ChainId,
654        heights: Vec<BlockHeight>,
655        sender: oneshot::Sender<Result<Vec<ConfirmedBlockCertificate>, NodeError>>,
656    ) -> Result<(), Result<Vec<ConfirmedBlockCertificate>, NodeError>> {
657        // First, use do_handle_chain_info_query to get the certificate hashes
658        let (query_sender, query_receiver) = oneshot::channel();
659        let query = ChainInfoQuery::new(chain_id).with_sent_certificate_hashes_by_heights(heights);
660
661        let self_clone = self.clone();
662        self.do_handle_chain_info_query(query, query_sender)
663            .await
664            .expect("Failed to handle chain info query");
665
666        // Get the response from the chain info query
667        let chain_info_response = query_receiver.await.map_err(|_| {
668            Err(NodeError::ClientIoError {
669                error: "Failed to receive chain info response".to_string(),
670            })
671        })?;
672
673        let hashes = match chain_info_response {
674            Ok(response) => response.info.requested_sent_certificate_hashes,
675            Err(e) => {
676                return sender.send(Err(e));
677            }
678        };
679
680        // Now use do_download_certificates to get the actual certificates
681        let (cert_sender, cert_receiver) = oneshot::channel();
682        self_clone
683            .do_download_certificates(hashes, cert_sender)
684            .await?;
685
686        // Forward the result to the original sender
687        let result = cert_receiver.await.map_err(|_| {
688            Err(NodeError::ClientIoError {
689                error: "Failed to receive certificates".to_string(),
690            })
691        })?;
692
693        sender.send(result)
694    }
695
696    async fn do_blob_last_used_by(
697        self,
698        blob_id: BlobId,
699        sender: oneshot::Sender<Result<CryptoHash, NodeError>>,
700    ) -> Result<(), Result<CryptoHash, NodeError>> {
701        let validator = self.client.lock().await;
702        let blob_state = validator
703            .state
704            .storage_client()
705            .read_blob_state(blob_id)
706            .await
707            .map_err(Into::into);
708        let certificate_hash = match blob_state {
709            Err(err) => Err(err),
710            Ok(blob_state) => match blob_state {
711                None => Err(NodeError::BlobsNotFound(vec![blob_id])),
712                Some(blob_state) => blob_state
713                    .last_used_by
714                    .ok_or_else(|| NodeError::BlobsNotFound(vec![blob_id])),
715            },
716        };
717
718        sender.send(certificate_hash)
719    }
720
721    async fn do_blob_last_used_by_certificate(
722        self,
723        blob_id: BlobId,
724        sender: oneshot::Sender<Result<ConfirmedBlockCertificate, NodeError>>,
725    ) -> Result<(), Result<ConfirmedBlockCertificate, NodeError>> {
726        match self.blob_last_used_by(blob_id).await {
727            Ok(cert_hash) => {
728                let cert = self.download_certificate(cert_hash).await;
729                sender.send(cert)
730            }
731            Err(err) => sender.send(Err(err)),
732        }
733    }
734
735    async fn do_missing_blob_ids(
736        self,
737        blob_ids: Vec<BlobId>,
738        sender: oneshot::Sender<Result<Vec<BlobId>, NodeError>>,
739    ) -> Result<(), Result<Vec<BlobId>, NodeError>> {
740        let validator = self.client.lock().await;
741        let missing_blob_ids = validator
742            .state
743            .storage_client()
744            .missing_blobs(&blob_ids)
745            .await
746            .map_err(Into::into);
747        sender.send(missing_blob_ids)
748    }
749
750    async fn do_event_block_heights(
751        self,
752        event_ids: Vec<EventId>,
753        sender: oneshot::Sender<Result<Vec<Option<BlockHeight>>, NodeError>>,
754    ) -> Result<(), Result<Vec<Option<BlockHeight>>, NodeError>> {
755        let validator = self.client.lock().await;
756        let heights = validator
757            .state
758            .storage_client()
759            .read_event_block_heights(&event_ids)
760            .await
761            .map_err(Into::into);
762        sender.send(heights)
763    }
764}
765
766/// A [`ValidatorNodeProvider`] holding the in-process test validator clients.
767#[derive(Clone)]
768pub struct NodeProvider<S>(Arc<std::sync::Mutex<Vec<LocalValidatorClient<S>>>>)
769where
770    S: Storage;
771
772impl<S> NodeProvider<S>
773where
774    S: Storage + Clone,
775{
776    fn all_nodes(&self) -> Vec<LocalValidatorClient<S>> {
777        self.0.lock().unwrap().clone()
778    }
779}
780
781impl<S> ValidatorNodeProvider for NodeProvider<S>
782where
783    S: Storage + Clone + Send + Sync + 'static,
784{
785    type Node = LocalValidatorClient<S>;
786
787    fn make_node(&self, _name: &str) -> Result<Self::Node, NodeError> {
788        unimplemented!()
789    }
790
791    fn make_nodes_from_list<A>(
792        &self,
793        validators: impl IntoIterator<Item = (ValidatorPublicKey, A)>,
794    ) -> Result<impl Iterator<Item = (ValidatorPublicKey, Self::Node)>, NodeError>
795    where
796        A: AsRef<str>,
797    {
798        let list = self.0.lock().unwrap();
799        Ok(validators
800            .into_iter()
801            .map(|(public_key, address)| {
802                list.iter()
803                    .find(|client| client.public_key == public_key)
804                    .ok_or_else(|| NodeError::CannotResolveValidatorAddress {
805                        address: address.as_ref().to_string(),
806                    })
807                    .map(|client| (public_key, client.clone()))
808            })
809            .collect::<Result<Vec<_>, _>>()?
810            .into_iter())
811    }
812}
813
814impl<S> FromIterator<LocalValidatorClient<S>> for NodeProvider<S>
815where
816    S: Storage,
817{
818    fn from_iter<T>(iter: T) -> Self
819    where
820        T: IntoIterator<Item = LocalValidatorClient<S>>,
821    {
822        Self(Arc::new(std::sync::Mutex::new(iter.into_iter().collect())))
823    }
824}
825
826// NOTE:
827// * To communicate with a quorum of validators, chain clients iterate over a copy of
828// `validator_clients` to spawn I/O tasks.
829// * When using `LocalValidatorClient`, clients communicate with an exact quorum then stop.
830// * Most tests have 1 faulty validator out 4 so that there is exactly only 1 quorum to
831// communicate with.
832#[allow(missing_docs)]
833pub struct TestBuilder<B: StorageBuilder> {
834    storage_builder: B,
835    pub initial_committee: Committee,
836    admin_description: Option<ChainDescription>,
837    network_description: Option<NetworkDescription>,
838    genesis_storage_builder: GenesisStorageBuilder,
839    node_provider: NodeProvider<B::Storage>,
840    pub validator_storages: HashMap<ValidatorPublicKey, B::Storage>,
841    pub validator_key_pairs: HashMap<ValidatorPublicKey, ValidatorSecretKey>,
842    chain_client_storages: Vec<B::Storage>,
843    pub chain_owners: BTreeMap<ChainId, AccountOwner>,
844    pub signer: TestSigner,
845}
846
847/// Builds storage instances of a specific backend for use in tests.
848#[async_trait]
849pub trait StorageBuilder {
850    /// The storage type produced by this builder.
851    type Storage: Storage + Clone + Send + Sync + 'static;
852
853    /// Builds a new storage instance.
854    async fn build(&mut self) -> Result<Self::Storage, anyhow::Error>;
855
856    /// Returns the test clock shared by all storages built here.
857    fn clock(&self) -> &TestClock;
858}
859
860#[derive(Default)]
861struct GenesisStorageBuilder {
862    accounts: Vec<GenesisAccount>,
863}
864
865struct GenesisAccount {
866    description: ChainDescription,
867    public_key: AccountPublicKey,
868}
869
870impl GenesisStorageBuilder {
871    fn add(&mut self, description: ChainDescription, public_key: AccountPublicKey) {
872        self.accounts.push(GenesisAccount {
873            description,
874            public_key,
875        })
876    }
877
878    async fn build<S>(&self, storage: S) -> S
879    where
880        S: Storage + Clone + Send + Sync + 'static,
881    {
882        for account in &self.accounts {
883            storage
884                .create_chain(account.description.clone())
885                .await
886                .unwrap();
887        }
888        storage
889    }
890}
891
892/// A chain client wired up to the in-process test validator network.
893pub type ChainClient<S> = crate::client::ChainClient<crate::environment::Impl<S, NodeProvider<S>>>;
894
895impl<S: Storage + Clone + Send + Sync + 'static> ChainClient<S> {
896    /// Reads the hashed certificate values in descending order from the given hash.
897    pub async fn read_confirmed_blocks_downward(
898        &self,
899        from: CryptoHash,
900        limit: u32,
901    ) -> anyhow::Result<Vec<Arc<ConfirmedBlock>>> {
902        let mut hash = Some(from);
903        let mut values = Vec::new();
904        for _ in 0..limit {
905            let Some(next_hash) = hash else {
906                break;
907            };
908            let value = self.read_confirmed_block(next_hash).await?;
909            hash = value.block().header.previous_block_hash;
910            values.push(value);
911        }
912        Ok(values)
913    }
914}
915
916impl<B> TestBuilder<B>
917where
918    B: StorageBuilder,
919{
920    /// Creates a test setup with `count` validators, `with_faulty_validators` of which are faulty.
921    pub async fn new(
922        mut storage_builder: B,
923        count: usize,
924        with_faulty_validators: usize,
925        mut signer: TestSigner,
926    ) -> Result<Self, anyhow::Error> {
927        let mut validators = Vec::new();
928        for _ in 0..count {
929            let validator_keypair = ValidatorKeypair::generate();
930            let account_public_key = signer.generate_new();
931            validators.push((validator_keypair, account_public_key));
932        }
933        let for_committee = validators
934            .iter()
935            .map(|(validating, account)| (validating.public_key, *account))
936            .collect::<Vec<_>>();
937        let initial_committee = Committee::make_simple(for_committee);
938        let mut validator_clients = Vec::new();
939        let mut validator_storages = HashMap::new();
940        let mut validator_key_pairs = HashMap::new();
941        let mut faulty_validators = HashSet::new();
942        for (i, (validator_keypair, _account_public_key)) in validators.into_iter().enumerate() {
943            let validator_public_key = validator_keypair.public_key;
944            let storage = storage_builder.build().await?;
945            let secret_key_copy = validator_keypair.secret_key.copy();
946            let config = ChainWorkerConfig {
947                nickname: format!("Node {i}"),
948                ..ChainWorkerConfig::default()
949            }
950            .with_key_pair(Some(validator_keypair.secret_key));
951            let state = WorkerState::new(storage.clone(), config, None);
952            let mut validator = LocalValidatorClient::new(validator_public_key, state);
953            if i < with_faulty_validators {
954                faulty_validators.insert(validator_public_key);
955                validator.set_fault_type(FaultType::NoChains);
956            }
957            validator_clients.push(validator);
958            validator_storages.insert(validator_public_key, storage);
959            validator_key_pairs.insert(validator_public_key, secret_key_copy);
960        }
961        tracing::info!(
962            "Test will use the following faulty validators: {:?}",
963            faulty_validators
964        );
965        Ok(Self {
966            storage_builder,
967            initial_committee,
968            admin_description: None,
969            network_description: None,
970            genesis_storage_builder: GenesisStorageBuilder::default(),
971            node_provider: NodeProvider::from_iter(validator_clients),
972            validator_storages,
973            validator_key_pairs,
974            chain_client_storages: Vec::new(),
975            chain_owners: BTreeMap::new(),
976            signer,
977        })
978    }
979
980    /// Replaces the initial committee's resource control policy.
981    pub fn with_policy(mut self, policy: ResourceControlPolicy) -> Self {
982        let validators = self.initial_committee.validators().clone();
983        self.initial_committee =
984            Committee::new(validators, policy).expect("committee votes should not overflow");
985        self
986    }
987
988    /// Sets the cross-chain message chunk limit on every validator in the test setup.
989    pub fn with_cross_chain_message_chunk_limit(self, limit: usize) -> Self {
990        let validator_clients = self.node_provider.0.lock().unwrap();
991        for validator in validator_clients.iter() {
992            let mut inner = validator.client.try_lock().expect("no contention at setup");
993            inner.state.set_cross_chain_message_chunk_limit(limit);
994        }
995        drop(validator_clients);
996        self
997    }
998
999    /// Returns the [`FaultType`] currently configured for the given validator, or `None`
1000    /// if no validator with that key is in the test setup.
1001    pub fn fault_type(&self, public_key: &ValidatorPublicKey) -> Option<FaultType> {
1002        self.node_provider
1003            .0
1004            .lock()
1005            .unwrap()
1006            .iter()
1007            .find(|client| client.public_key == *public_key)
1008            .map(|client| client.fault_type())
1009    }
1010
1011    /// Sets the [`FaultType`] for the validators at the given indexes.
1012    pub fn set_fault_type(&mut self, indexes: impl AsRef<[usize]>, fault_type: FaultType) {
1013        let mut faulty_validators = vec![];
1014        let mut validator_clients = self.node_provider.0.lock().unwrap();
1015        for index in indexes.as_ref() {
1016            let validator = &mut validator_clients[*index];
1017            validator.set_fault_type(fault_type);
1018            faulty_validators.push(validator.public_key);
1019        }
1020        tracing::info!(
1021            "Making the following validators {:?}: {:?}",
1022            fault_type,
1023            faulty_validators
1024        );
1025    }
1026
1027    /// Creates the root chain with the given `index`, and returns a client for it.
1028    ///
1029    /// Root chain 0 is the admin chain and needs to be initialized first, otherwise its balance
1030    /// is automatically set to zero.
1031    pub async fn add_root_chain(
1032        &mut self,
1033        index: u32,
1034        balance: Amount,
1035    ) -> anyhow::Result<ChainClient<B::Storage>> {
1036        self.add_root_chain_with_ownership(index, balance, ChainOwnership::single)
1037            .await
1038    }
1039
1040    /// Creates the root chain with the given `index` and a genesis ownership built from its
1041    /// freshly generated owner key, and returns a client for it.
1042    ///
1043    /// Root chain 0 is the admin chain and needs to be initialized first, otherwise its balance
1044    /// is automatically set to zero.
1045    pub async fn add_root_chain_with_ownership(
1046        &mut self,
1047        index: u32,
1048        balance: Amount,
1049        make_ownership: impl FnOnce(AccountOwner) -> ChainOwnership,
1050    ) -> anyhow::Result<ChainClient<B::Storage>> {
1051        // Make sure the admin chain is initialized.
1052        if self.admin_description.is_none() && index != 0 {
1053            Box::pin(self.add_root_chain(0, Amount::ZERO)).await?;
1054        }
1055        let origin = ChainOrigin::Root(index);
1056        let public_key = self.signer.generate_new();
1057        let open_chain_config = InitialChainConfig {
1058            ownership: make_ownership(public_key.into()),
1059            epoch: Epoch(0),
1060            balance,
1061            application_permissions: ApplicationPermissions::default(),
1062        };
1063        let description = ChainDescription::new(origin, open_chain_config, Timestamp::from(0));
1064        let committee_blob = Blob::new_committee(bcs::to_bytes(&self.initial_committee).unwrap());
1065        if index == 0 {
1066            self.admin_description = Some(description.clone());
1067            self.network_description = Some(NetworkDescription {
1068                admin_chain_id: description.id(),
1069                // dummy values to fill the description
1070                genesis_config_hash: CryptoHash::test_hash("genesis config"),
1071                genesis_timestamp: Timestamp::from(0),
1072                genesis_committee_blob_hash: committee_blob.id().hash,
1073                name: "test network".to_string(),
1074            });
1075        }
1076        // Remember what's in the genesis store for future clients to join.
1077        self.genesis_storage_builder
1078            .add(description.clone(), public_key);
1079
1080        let network_description = self.network_description.as_ref().unwrap();
1081
1082        for validator in self.node_provider.all_nodes() {
1083            let storage = self
1084                .validator_storages
1085                .get_mut(&validator.public_key)
1086                .unwrap();
1087            storage
1088                .write_network_description(network_description)
1089                .await
1090                .expect("writing the NetworkDescription should succeed");
1091            storage
1092                .write_blob(&committee_blob)
1093                .await
1094                .expect("writing a blob should succeed");
1095            storage.create_chain(description.clone()).await.unwrap();
1096        }
1097        for storage in &mut self.chain_client_storages {
1098            storage.create_chain(description.clone()).await.unwrap();
1099        }
1100        let chain_id = description.id();
1101        self.chain_owners.insert(chain_id, public_key.into());
1102        self.make_client(chain_id, None, BlockHeight::ZERO).await
1103    }
1104
1105    /// Returns the public key and balance of each genesis root chain.
1106    pub fn genesis_chains(&self) -> Vec<(AccountPublicKey, Amount)> {
1107        let mut result = Vec::new();
1108        for (i, genesis_account) in self.genesis_storage_builder.accounts.iter().enumerate() {
1109            assert_eq!(
1110                genesis_account.description.origin(),
1111                ChainOrigin::Root(i as u32)
1112            );
1113            result.push((
1114                genesis_account.public_key,
1115                genesis_account.description.config().balance,
1116            ));
1117        }
1118        result
1119    }
1120
1121    /// Returns the admin chain's ID, panicking if it has not been initialized.
1122    pub fn admin_chain_id(&self) -> ChainId {
1123        self.admin_description
1124            .as_ref()
1125            .expect("admin chain not initialized")
1126            .id()
1127    }
1128
1129    /// Returns the admin chain description, if the admin chain has been initialized.
1130    pub fn admin_description(&self) -> Option<&ChainDescription> {
1131        self.admin_description.as_ref()
1132    }
1133
1134    /// Returns a clone of the node provider backing this test setup.
1135    pub fn make_node_provider(&self) -> NodeProvider<B::Storage> {
1136        self.node_provider.clone()
1137    }
1138
1139    /// Returns a clone of the validator client at the given index.
1140    pub fn node(&mut self, index: usize) -> LocalValidatorClient<B::Storage> {
1141        self.node_provider.0.lock().unwrap()[index].clone()
1142    }
1143
1144    /// Builds a fresh storage seeded with the network description and genesis chains.
1145    pub async fn make_storage(&mut self) -> anyhow::Result<B::Storage> {
1146        let storage = self.storage_builder.build().await?;
1147        let network_description = self.network_description.as_ref().unwrap();
1148        let committee_blob = Blob::new_committee(bcs::to_bytes(&self.initial_committee).unwrap());
1149        storage
1150            .write_network_description(network_description)
1151            .await
1152            .expect("writing the NetworkDescription should succeed");
1153        storage
1154            .write_blob(&committee_blob)
1155            .await
1156            .expect("writing a blob should succeed");
1157        Ok(self.genesis_storage_builder.build(storage).await)
1158    }
1159
1160    /// Creates a chain client for the given chain with the given client options.
1161    pub async fn make_client_with_options(
1162        &mut self,
1163        chain_id: ChainId,
1164        block_hash: Option<CryptoHash>,
1165        block_height: BlockHeight,
1166        options: chain_client::Options,
1167        follow_only: bool,
1168    ) -> anyhow::Result<ChainClient<B::Storage>> {
1169        // Note that new clients are only given the genesis store: they must figure out
1170        // the rest by asking validators.
1171        let storage = self.make_storage().await?;
1172        self.chain_client_storages.push(storage.clone());
1173        let mode = if follow_only {
1174            crate::client::ListeningMode::FollowChain
1175        } else {
1176            crate::client::ListeningMode::FullChain
1177        };
1178        let client = Arc::new(Client::new(
1179            crate::environment::Impl {
1180                network: self.make_node_provider(),
1181                storage,
1182                signer: self.signer.clone(),
1183                wallet: TestWallet::default(),
1184            },
1185            self.admin_chain_id(),
1186            false,
1187            [(chain_id, mode)],
1188            format!("Client node for {chain_id:.8}"),
1189            Some(Duration::from_secs(30)),
1190            Some(Duration::from_secs(1)),
1191            1000,
1192            options,
1193            DEFAULT_BLOCK_CACHE_SIZE,
1194            DEFAULT_EXECUTION_STATE_CACHE_SIZE,
1195            &crate::client::RequestsSchedulerConfig::default(),
1196        ));
1197        Ok(client.create_chain_client(
1198            chain_id,
1199            block_hash,
1200            block_height,
1201            &None,
1202            self.chain_owners.get(&chain_id).copied(),
1203            None,
1204            follow_only,
1205        ))
1206    }
1207
1208    /// Creates a chain client for the given chain with default test options.
1209    pub async fn make_client(
1210        &mut self,
1211        chain_id: ChainId,
1212        block_hash: Option<CryptoHash>,
1213        block_height: BlockHeight,
1214    ) -> anyhow::Result<ChainClient<B::Storage>> {
1215        self.make_client_with_options(
1216            chain_id,
1217            block_hash,
1218            block_height,
1219            chain_client::Options::test_default(),
1220            false,
1221        )
1222        .await
1223    }
1224
1225    /// Tries to find a (confirmation) certificate for the given chain_id and block height.
1226    pub async fn check_that_validators_have_certificate(
1227        &self,
1228        chain_id: ChainId,
1229        block_height: BlockHeight,
1230        target_count: usize,
1231    ) -> Option<ConfirmedBlockCertificate> {
1232        let query = ChainInfoQuery::new(chain_id)
1233            .with_sent_certificate_hashes_by_heights(vec![block_height]);
1234        let mut count = 0;
1235        let mut certificate = None;
1236        for validator in self.node_provider.all_nodes() {
1237            if let Ok(response) = validator.handle_chain_info_query(query.clone()).await {
1238                if response.check(validator.public_key).is_ok() {
1239                    let ChainInfo {
1240                        mut requested_sent_certificate_hashes,
1241                        ..
1242                    } = *response.info;
1243                    debug_assert!(requested_sent_certificate_hashes.len() <= 1);
1244                    if let Some(cert_hash) = requested_sent_certificate_hashes.pop() {
1245                        if let Ok(cert) = validator.download_certificate(cert_hash).await {
1246                            if cert.inner().block().header.chain_id == chain_id
1247                                && cert.inner().block().header.height == block_height
1248                            {
1249                                cert.check(&self.initial_committee).unwrap();
1250                                count += 1;
1251                                certificate = Some(cert);
1252                            }
1253                        }
1254                    }
1255                }
1256            }
1257        }
1258        assert!(count >= target_count);
1259        certificate
1260    }
1261
1262    /// Tries to find a (confirmation) certificate for the given chain_id and block height, and are
1263    /// in the expected round.
1264    pub async fn check_that_validators_are_in_round(
1265        &self,
1266        chain_id: ChainId,
1267        block_height: BlockHeight,
1268        round: Round,
1269        target_count: usize,
1270    ) {
1271        let query = ChainInfoQuery::new(chain_id);
1272        let mut count = 0;
1273        for validator in self.node_provider.all_nodes() {
1274            if let Ok(response) = validator.handle_chain_info_query(query.clone()).await {
1275                if response.info.manager.current_round == round
1276                    && response.info.next_block_height == block_height
1277                    && response.check(validator.public_key).is_ok()
1278                {
1279                    count += 1;
1280                }
1281            }
1282        }
1283        assert!(count >= target_count);
1284    }
1285
1286    /// Panics if any validator has a nonempty outbox for the given chain.
1287    pub async fn check_that_validators_have_empty_outboxes(&self, chain_id: ChainId) {
1288        for validator in self.node_provider.all_nodes() {
1289            let guard = validator.client.lock().await;
1290            let chain = guard.state.chain_state_view(chain_id).await.unwrap();
1291            assert_eq!(chain.outboxes.indices().await.unwrap(), []);
1292        }
1293    }
1294}
1295
1296#[cfg(feature = "rocksdb")]
1297/// Limit concurrency for RocksDB tests to avoid "too many open files" errors.
1298static ROCKS_DB_SEMAPHORE: Semaphore = Semaphore::const_new(5);
1299
1300/// State shared by every [`StorageBuilder`] in this module. The actual
1301/// database type varies, so `build_storage` is generic over it.
1302#[derive(Default)]
1303struct CommonStorageBuilder {
1304    namespace: String,
1305    instance_counter: usize,
1306    wasm_runtime: Option<WasmRuntime>,
1307    clock: TestClock,
1308}
1309
1310impl CommonStorageBuilder {
1311    fn with_wasm_runtime(wasm_runtime: impl Into<Option<WasmRuntime>>) -> Self {
1312        Self {
1313            wasm_runtime: wasm_runtime.into(),
1314            ..Self::default()
1315        }
1316    }
1317
1318    async fn build_storage<DB>(
1319        &mut self,
1320        config: DB::Config,
1321    ) -> anyhow::Result<DbStorage<DB, TestClock>>
1322    where
1323        DB: TestKeyValueDatabase + Clone + Send + Sync + 'static,
1324        DB::Store: KeyValueStore + Clone + Send + Sync + 'static,
1325        DB::Error: std::error::Error + Send + Sync + 'static,
1326    {
1327        self.instance_counter += 1;
1328        if self.namespace.is_empty() {
1329            self.namespace = generate_test_namespace();
1330        }
1331        let namespace = format!("{}_{}", self.namespace, self.instance_counter);
1332        Ok(
1333            DbStorage::new_for_testing(config, &namespace, self.wasm_runtime, self.clock.clone())
1334                .await?,
1335        )
1336    }
1337}
1338
1339/// A [`StorageBuilder`] backed by in-memory storage.
1340#[derive(Default)]
1341pub struct MemoryStorageBuilder {
1342    inner: CommonStorageBuilder,
1343}
1344
1345impl MemoryStorageBuilder {
1346    /// Creates a [`MemoryStorageBuilder`] that uses the specified [`WasmRuntime`] to run Wasm
1347    /// applications.
1348    pub fn with_wasm_runtime(wasm_runtime: impl Into<Option<WasmRuntime>>) -> Self {
1349        Self {
1350            inner: CommonStorageBuilder::with_wasm_runtime(wasm_runtime),
1351        }
1352    }
1353}
1354
1355#[async_trait]
1356impl StorageBuilder for MemoryStorageBuilder {
1357    type Storage = DbStorage<MemoryDatabase, TestClock>;
1358
1359    async fn build(&mut self) -> Result<Self::Storage, anyhow::Error> {
1360        let config = MemoryDatabase::new_test_config().await?;
1361        self.inner.build_storage::<MemoryDatabase>(config).await
1362    }
1363
1364    fn clock(&self) -> &TestClock {
1365        &self.inner.clock
1366    }
1367}
1368
1369#[cfg(feature = "rocksdb")]
1370/// A [`StorageBuilder`] backed by RocksDB storage.
1371pub struct RocksDbStorageBuilder {
1372    inner: CommonStorageBuilder,
1373    _permit: SemaphorePermit<'static>,
1374}
1375
1376#[cfg(feature = "rocksdb")]
1377impl RocksDbStorageBuilder {
1378    /// Creates a [`RocksDbStorageBuilder`], acquiring a concurrency permit.
1379    pub async fn new() -> Self {
1380        Self {
1381            inner: CommonStorageBuilder::default(),
1382            _permit: ROCKS_DB_SEMAPHORE.acquire().await.unwrap(),
1383        }
1384    }
1385
1386    /// Creates a [`RocksDbStorageBuilder`] that uses the specified [`WasmRuntime`] to run Wasm
1387    /// applications.
1388    #[cfg(any(feature = "wasmer", feature = "wasmtime"))]
1389    pub async fn with_wasm_runtime(wasm_runtime: impl Into<Option<WasmRuntime>>) -> Self {
1390        Self {
1391            inner: CommonStorageBuilder::with_wasm_runtime(wasm_runtime),
1392            _permit: ROCKS_DB_SEMAPHORE.acquire().await.unwrap(),
1393        }
1394    }
1395}
1396
1397#[cfg(feature = "rocksdb")]
1398#[async_trait]
1399impl StorageBuilder for RocksDbStorageBuilder {
1400    type Storage = DbStorage<RocksDbDatabase, TestClock>;
1401
1402    async fn build(&mut self) -> Result<Self::Storage, anyhow::Error> {
1403        let config = RocksDbDatabase::new_test_config().await?;
1404        self.inner.build_storage::<RocksDbDatabase>(config).await
1405    }
1406
1407    fn clock(&self) -> &TestClock {
1408        &self.inner.clock
1409    }
1410}
1411
1412#[cfg(all(not(target_arch = "wasm32"), feature = "storage-service"))]
1413/// A [`StorageBuilder`] backed by the storage service.
1414#[derive(Default)]
1415pub struct ServiceStorageBuilder {
1416    inner: CommonStorageBuilder,
1417}
1418
1419#[cfg(all(not(target_arch = "wasm32"), feature = "storage-service"))]
1420impl ServiceStorageBuilder {
1421    /// Creates a `ServiceStorage`.
1422    pub fn new() -> Self {
1423        Self::with_wasm_runtime(None)
1424    }
1425
1426    /// Creates a `ServiceStorage` with the given Wasm runtime.
1427    pub fn with_wasm_runtime(wasm_runtime: impl Into<Option<WasmRuntime>>) -> Self {
1428        Self {
1429            inner: CommonStorageBuilder::with_wasm_runtime(wasm_runtime),
1430        }
1431    }
1432}
1433
1434#[cfg(all(not(target_arch = "wasm32"), feature = "storage-service"))]
1435#[async_trait]
1436impl StorageBuilder for ServiceStorageBuilder {
1437    type Storage = DbStorage<StorageServiceDatabase, TestClock>;
1438
1439    async fn build(&mut self) -> anyhow::Result<Self::Storage> {
1440        let config = StorageServiceDatabase::new_test_config().await?;
1441        self.inner
1442            .build_storage::<StorageServiceDatabase>(config)
1443            .await
1444    }
1445
1446    fn clock(&self) -> &TestClock {
1447        &self.inner.clock
1448    }
1449}
1450
1451#[cfg(feature = "scylladb")]
1452/// A [`StorageBuilder`] backed by ScyllaDB storage.
1453#[derive(Default)]
1454pub struct ScyllaDbStorageBuilder {
1455    inner: CommonStorageBuilder,
1456}
1457
1458#[cfg(feature = "scylladb")]
1459impl ScyllaDbStorageBuilder {
1460    /// Creates a [`ScyllaDbStorageBuilder`] that uses the specified [`WasmRuntime`] to run Wasm
1461    /// applications.
1462    pub fn with_wasm_runtime(wasm_runtime: impl Into<Option<WasmRuntime>>) -> Self {
1463        Self {
1464            inner: CommonStorageBuilder::with_wasm_runtime(wasm_runtime),
1465        }
1466    }
1467}
1468
1469#[cfg(feature = "scylladb")]
1470#[async_trait]
1471impl StorageBuilder for ScyllaDbStorageBuilder {
1472    type Storage = DbStorage<ScyllaDbDatabase, TestClock>;
1473
1474    async fn build(&mut self) -> Result<Self::Storage, anyhow::Error> {
1475        let config = ScyllaDbDatabase::new_test_config().await?;
1476        self.inner.build_storage::<ScyllaDbDatabase>(config).await
1477    }
1478
1479    fn clock(&self) -> &TestClock {
1480        &self.inner.clock
1481    }
1482}
1483
1484/// Helpers for asserting on [`ClientOutcome`] results in tests.
1485pub trait ClientOutcomeResultExt<T, E> {
1486    /// Unwraps the result and panics if it's not `Committed`.
1487    /// Use this when you expect the operation to succeed without conflicts.
1488    fn unwrap_ok_committed(self) -> T;
1489
1490    /// Unwraps the result, accepting both `Committed` and `Conflict` outcomes.
1491    /// Returns the committed value or the conflicting certificate (boxed).
1492    fn unwrap_ok_or_conflict(self) -> Result<T, Box<ConfirmedBlockCertificate>>;
1493}
1494
1495impl<T, E: std::fmt::Debug> ClientOutcomeResultExt<T, E> for Result<ClientOutcome<T>, E> {
1496    fn unwrap_ok_committed(self) -> T {
1497        match self.unwrap() {
1498            ClientOutcome::Committed(t) => t,
1499            ClientOutcome::WaitForTimeout(timeout) => {
1500                panic!("unexpected timeout: {timeout}")
1501            }
1502            ClientOutcome::Conflict(certificate) => {
1503                panic!("unexpected conflict: {}", certificate.hash())
1504            }
1505        }
1506    }
1507
1508    fn unwrap_ok_or_conflict(self) -> Result<T, Box<ConfirmedBlockCertificate>> {
1509        match self.unwrap() {
1510            ClientOutcome::Committed(t) => Ok(t),
1511            ClientOutcome::Conflict(certificate) => Err(certificate),
1512            ClientOutcome::WaitForTimeout(timeout) => {
1513                panic!("unexpected timeout: {timeout}")
1514            }
1515        }
1516    }
1517}