iroh_blobs/
test.rs

1use std::future::IntoFuture;
2
3use n0_future::{stream, StreamExt};
4use rand::{RngCore, SeedableRng};
5
6use crate::{
7    api::{blobs::AddBytesOptions, tags::TagInfo, RequestResult, Store},
8    hashseq::HashSeq,
9    BlobFormat,
10};
11
12pub async fn create_random_blobs<R: rand::Rng>(
13    store: &Store,
14    num_blobs: usize,
15    blob_size: impl Fn(usize, &mut R) -> usize,
16    mut rand: R,
17) -> anyhow::Result<Vec<TagInfo>> {
18    // generate sizes and seeds, non-parrallelized so it is deterministic
19    let sizes = (0..num_blobs)
20        .map(|n| (blob_size(n, &mut rand), rand.r#gen::<u64>()))
21        .collect::<Vec<_>>();
22    // generate random data and add it to the store
23    let infos = stream::iter(sizes)
24        .then(|(size, seed)| {
25            let mut rand = rand::rngs::StdRng::seed_from_u64(seed);
26            let mut data = vec![0u8; size];
27            rand.fill_bytes(&mut data);
28            store.add_bytes(data).into_future()
29        })
30        .collect::<Vec<_>>()
31        .await;
32    let infos = infos.into_iter().collect::<RequestResult<Vec<_>>>()?;
33    Ok(infos)
34}
35
36pub async fn add_hash_sequences<R: rand::Rng>(
37    store: &Store,
38    tags: &[TagInfo],
39    num_seqs: usize,
40    seq_size: impl Fn(usize, &mut R) -> usize,
41    mut rand: R,
42) -> anyhow::Result<Vec<TagInfo>> {
43    let infos = stream::iter(0..num_seqs)
44        .then(|n| {
45            let size = seq_size(n, &mut rand);
46            let hs = (0..size)
47                .map(|_| {
48                    let j = rand.gen_range(0..tags.len());
49                    tags[j].hash
50                })
51                .collect::<HashSeq>();
52            store
53                .add_bytes_with_opts(AddBytesOptions {
54                    data: hs.into(),
55                    format: BlobFormat::HashSeq,
56                })
57                .into_future()
58        })
59        .collect::<Vec<_>>()
60        .await;
61    let infos = infos.into_iter().collect::<RequestResult<Vec<_>>>()?;
62    Ok(infos)
63}