reth_network_p2p/test_utils/
bodies.rs

1use crate::{
2    bodies::client::{BodiesClient, BodiesFut},
3    download::DownloadClient,
4    error::PeerRequestResult,
5    priority::Priority,
6};
7use alloy_primitives::B256;
8use futures::FutureExt;
9use reth_ethereum_primitives::BlockBody;
10use reth_network_peers::PeerId;
11use std::{
12    fmt::{Debug, Formatter},
13    ops::RangeInclusive,
14};
15use tokio::sync::oneshot;
16
17/// A test client for fetching bodies
18pub struct TestBodiesClient<F> {
19    /// The function that is called on each body request.
20    pub responder: F,
21}
22
23impl<F> Debug for TestBodiesClient<F> {
24    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
25        f.debug_struct("TestBodiesClient").finish_non_exhaustive()
26    }
27}
28
29impl<F: Sync + Send> DownloadClient for TestBodiesClient<F> {
30    fn report_bad_message(&self, _peer_id: PeerId) {
31        // noop
32    }
33
34    fn num_connected_peers(&self) -> usize {
35        0
36    }
37}
38
39impl<F> BodiesClient for TestBodiesClient<F>
40where
41    F: Fn(Vec<B256>) -> PeerRequestResult<Vec<BlockBody>> + Send + Sync,
42{
43    type Body = BlockBody;
44    type Output = BodiesFut;
45
46    fn get_block_bodies_with_priority_and_range_hint(
47        &self,
48        hashes: Vec<B256>,
49        _priority: Priority,
50        _range_hint: Option<RangeInclusive<u64>>,
51    ) -> Self::Output {
52        let (tx, rx) = oneshot::channel();
53        let _ = tx.send((self.responder)(hashes));
54        Box::pin(rx.map(|x| match x {
55            Ok(value) => value,
56            Err(err) => Err(err.into()),
57        }))
58    }
59}