reth_ress_protocol/
test_utils.rs

Help
1//! Miscellaneous test utilities.
2
3use crate::RessProtocolProvider;
4use alloy_consensus::Header;
5use alloy_primitives::{map::B256HashMap, Bytes, B256};
6use reth_ethereum_primitives::BlockBody;
7use reth_storage_errors::provider::ProviderResult;
8use std::{
9    sync::{Arc, Mutex},
10    time::Duration,
11};
12
13/// Noop implementation of [`RessProtocolProvider`].
14#[derive(Clone, Copy, Default, Debug)]
15pub struct NoopRessProtocolProvider;
16
17impl RessProtocolProvider for NoopRessProtocolProvider {
18    fn header(&self, _block_hash: B256) -> ProviderResult<Option<Header>> {
19        Ok(None)
20    }
21
22    fn block_body(&self, _block_hash: B256) -> ProviderResult<Option<BlockBody>> {
23        Ok(None)
24    }
25
26    fn bytecode(&self, _code_hash: B256) -> ProviderResult<Option<Bytes>> {
27        Ok(None)
28    }
29
30    async fn witness(&self, _block_hash: B256) -> ProviderResult<Vec<Bytes>> {
31        Ok(Vec::new())
32    }
33}
34
35/// Mock implementation of [`RessProtocolProvider`].
36#[derive(Clone, Default, Debug)]
37pub struct MockRessProtocolProvider {
38    headers: Arc<Mutex<B256HashMap<Header>>>,
39    block_bodies: Arc<Mutex<B256HashMap<BlockBody>>>,
40    bytecodes: Arc<Mutex<B256HashMap<Bytes>>>,
41    witnesses: Arc<Mutex<B256HashMap<Vec<Bytes>>>>,
42    witness_delay: Option<Duration>,
43}
44
45impl MockRessProtocolProvider {
46    /// Configure witness response delay.
47    pub const fn with_witness_delay(mut self, delay: Duration) -> Self {
48        self.witness_delay = Some(delay);
49        self
50    }
51
52    /// Insert header.
53    pub fn add_header(&self, block_hash: B256, header: Header) {
54        self.headers.lock().unwrap().insert(block_hash, header);
55    }
56
57    /// Extend headers from iterator.
58    pub fn extend_headers(&self, headers: impl IntoIterator<Item = (B256, Header)>) {
59        self.headers.lock().unwrap().extend(headers);
60    }
61
62    /// Insert block body.
63    pub fn add_block_body(&self, block_hash: B256, body: BlockBody) {
64        self.block_bodies.lock().unwrap().insert(block_hash, body);
65    }
66
67    /// Extend block bodies from iterator.
68    pub fn extend_block_bodies(&self, bodies: impl IntoIterator<Item = (B256, BlockBody)>) {
69        self.block_bodies.lock().unwrap().extend(bodies);
70    }
71
72    /// Insert bytecode.
73    pub fn add_bytecode(&self, code_hash: B256, bytecode: Bytes) {
74        self.bytecodes.lock().unwrap().insert(code_hash, bytecode);
75    }
76
77    /// Extend bytecodes from iterator.
78    pub fn extend_bytecodes(&self, bytecodes: impl IntoIterator<Item = (B256, Bytes)>) {
79        self.bytecodes.lock().unwrap().extend(bytecodes);
80    }
81
82    /// Insert witness.
83    pub fn add_witness(&self, block_hash: B256, witness: Vec<Bytes>) {
84        self.witnesses.lock().unwrap().insert(block_hash, witness);
85    }
86
87    /// Extend witnesses from iterator.
88    pub fn extend_witnesses(&self, witnesses: impl IntoIterator<Item = (B256, Vec<Bytes>)>) {
89        self.witnesses.lock().unwrap().extend(witnesses);
90    }
91}
92
93impl RessProtocolProvider for MockRessProtocolProvider {
94    fn header(&self, block_hash: B256) -> ProviderResult<Option<Header>> {
95        Ok(self.headers.lock().unwrap().get(&block_hash).cloned())
96    }
97
98    fn block_body(&self, block_hash: B256) -> ProviderResult<Option<BlockBody>> {
99        Ok(self.block_bodies.lock().unwrap().get(&block_hash).cloned())
100    }
101
102    fn bytecode(&self, code_hash: B256) -> ProviderResult<Option<Bytes>> {
103        Ok(self.bytecodes.lock().unwrap().get(&code_hash).cloned())
104    }
105
106    async fn witness(&self, block_hash: B256) -> ProviderResult<Vec<Bytes>> {
107        if let Some(delay) = self.witness_delay {
108            tokio::time::sleep(delay).await;
109        }
110        Ok(self.witnesses.lock().unwrap().get(&block_hash).cloned().unwrap_or_default())
111    }
112}
reth_ress_protocol/test_utils.rs

reth_ress_protocol/
test_utils.rs
