reth_downloaders/test_utils/
mod.rs

1//! Test helper impls.
2
3#![allow(dead_code)]
4
5#[cfg(any(test, feature = "file-client"))]
6use crate::{bodies::test_utils::create_raw_bodies, file_codec::BlockFileCodec};
7use alloy_primitives::B256;
8use futures::SinkExt;
9use reth_ethereum_primitives::BlockBody;
10use reth_testing_utils::generators::{self, random_block_range, BlockRangeParams};
11use std::{collections::HashMap, io::SeekFrom, ops::RangeInclusive};
12use tokio::{fs::File, io::AsyncSeekExt};
13use tokio_util::codec::FramedWrite;
14
15mod bodies_client;
16pub use bodies_client::TestBodiesClient;
17use reth_primitives_traits::SealedHeader;
18
19/// Metrics scope used for testing.
20pub(crate) const TEST_SCOPE: &str = "downloaders.test";
21
22/// Generate a set of bodies and their corresponding block hashes
23pub(crate) fn generate_bodies(
24    range: RangeInclusive<u64>,
25) -> (Vec<SealedHeader>, HashMap<B256, BlockBody>) {
26    let mut rng = generators::rng();
27    let blocks = random_block_range(
28        &mut rng,
29        range,
30        BlockRangeParams { parent: Some(B256::ZERO), tx_count: 0..2, ..Default::default() },
31    );
32
33    let headers = blocks.iter().map(|block| block.clone_sealed_header()).collect();
34    let bodies = blocks.into_iter().map(|block| (block.hash(), block.into_body())).collect();
35
36    (headers, bodies)
37}
38
39/// Generate a set of bodies, write them to a temporary file, and return the file along with the
40/// bodies and corresponding block hashes
41#[cfg(any(test, feature = "file-client"))]
42pub(crate) async fn generate_bodies_file(
43    range: RangeInclusive<u64>,
44) -> (tokio::fs::File, Vec<SealedHeader>, HashMap<B256, BlockBody>) {
45    let (headers, bodies) = generate_bodies(range);
46    let raw_block_bodies = create_raw_bodies(headers.iter().cloned(), &mut bodies.clone());
47
48    let file: File = tempfile::tempfile().unwrap().into();
49    let mut writer = FramedWrite::new(file, BlockFileCodec::default());
50
51    // rlp encode one after the other
52    for block in raw_block_bodies {
53        writer.feed(block).await.unwrap();
54    }
55    writer.flush().await.unwrap();
56
57    // get the file back
58    let mut file: File = writer.into_inner();
59    file.seek(SeekFrom::Start(0)).await.unwrap();
60    (file, headers, bodies)
61}