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