xet_client/cas_client/simulation/
client_testing_utils.rs

1use std::collections::HashMap;
2
3use bytes::Bytes;
4use rand::prelude::*;
5use xet_core_structures::MerkleHashMap;
6use xet_core_structures::merklehash::{MerkleHash, compute_data_hash, file_hash_with_salt};
7use xet_core_structures::metadata_shard::file_structs::{FileDataSequenceEntry, FileDataSequenceHeader, MDBFileInfo};
8use xet_core_structures::metadata_shard::shard_in_memory::MDBInMemoryShard;
9use xet_core_structures::xorb_object::{Chunk, RawXorbData, SerializedXorbObject};
10
11use super::super::interface::Client;
12use crate::error::Result;
13
14/// Information about a term (segment) in the file, referencing an XORB and chunk range.
15#[derive(Clone, Debug)]
16pub struct FileTermReference {
17    /// The XORB hash this term references.
18    pub xorb_hash: MerkleHash,
19    /// Start chunk index (inclusive) within the XORB.
20    pub chunk_start: u32,
21    /// End chunk index (exclusive) within the XORB.
22    pub chunk_end: u32,
23    /// The data for this term (concatenated chunk data).
24    pub data: Bytes,
25    /// The chunk hashes for this term.
26    pub chunk_hashes: Vec<MerkleHash>,
27}
28
29/// Complete information about a randomly generated file for testing purposes.
30///
31/// Contains all the metadata needed to verify that reconstruction and fetching
32/// operations return correct data.
33#[derive(Clone, Debug)]
34pub struct RandomFileContents {
35    /// The file hash (used for reconstruction queries).
36    pub file_hash: MerkleHash,
37    /// The complete file data.
38    pub data: Bytes,
39    /// The RawXorbData for each XORB that was created, keyed by XORB hash.
40    pub xorbs: MerkleHashMap<RawXorbData>,
41    /// Information about each term in file order.
42    pub terms: Vec<FileTermReference>,
43}
44
45impl RandomFileContents {
46    /// Verifies that the given data matches the expected data for a specific term.
47    ///
48    /// This checks that the hash of the provided data matches the expected XORB
49    /// data for the term at the given index.
50    ///
51    /// # Arguments
52    /// * `term_index` - The index of the term (0-based) in the terms list
53    /// * `data` - The data to verify against the expected term data
54    ///
55    /// # Returns
56    /// `true` if the data matches the expected term data, `false` otherwise.
57    pub fn term_matches(&self, term_index: usize, data: &[u8]) -> bool {
58        if term_index >= self.terms.len() {
59            return false;
60        }
61        let term = &self.terms[term_index];
62        term.data == data
63    }
64
65    /// Returns the expected data for a specific term.
66    pub fn term_data(&self, term_index: usize) -> Option<&Bytes> {
67        self.terms.get(term_index).map(|t| &t.data)
68    }
69
70    /// Returns the XORB hash for a specific term.
71    pub fn term_xorb_hash(&self, term_index: usize) -> Option<MerkleHash> {
72        self.terms.get(term_index).map(|t| t.xorb_hash)
73    }
74
75    /// Returns the chunk range for a specific term.
76    pub fn term_chunk_range(&self, term_index: usize) -> Option<(u32, u32)> {
77        self.terms.get(term_index).map(|t| (t.chunk_start, t.chunk_end))
78    }
79}
80
81/// A trait that adds testing utility functions to the Client interface.
82#[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
83#[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
84pub trait ClientTestingUtils: Client + Send + Sync {
85    /// Insert a random file into the local CAS.
86    ///
87    /// This function generates a random file with the given term specification.
88    /// Each term is defined as `(xorb_seed, (chunk_start, chunk_end))` where:
89    /// - `xorb_seed` determines the random data for that XORB
90    /// - `chunk_start` and `chunk_end` define the range of chunks to include
91    ///
92    /// Returns a `RandomFileContents` struct containing all the metadata needed
93    /// to verify reconstruction and fetching operations.
94    async fn upload_random_file(
95        &self,
96        term_spec: &[(u64, (u64, u64))],
97        chunk_size: usize,
98    ) -> Result<RandomFileContents> {
99        let mut xorb_num_chunks = HashMap::<u64, u64>::new();
100
101        for &(xorb_seed, (_chunk_idx_start, chunk_idx_end)) in term_spec {
102            let c: &mut u64 = xorb_num_chunks.entry(xorb_seed).or_default();
103            *c = (*c).max(chunk_idx_end);
104        }
105
106        let mut shard = MDBInMemoryShard::default();
107        let mut xorb_data = HashMap::<u64, RawXorbData>::new();
108
109        for (&xorb_seed, n_chunks) in xorb_num_chunks.iter() {
110            let mut rng = SmallRng::seed_from_u64(xorb_seed);
111            let n_chunks = *n_chunks as usize;
112            let mut chunks = Vec::with_capacity(n_chunks);
113
114            for _idx in 0..n_chunks {
115                let n = rng.random_range((chunk_size / 2 + 1)..chunk_size);
116                let n_left = chunk_size - n;
117
118                let mut rng_data = vec![0u8; n];
119                rng.fill_bytes(&mut rng_data);
120
121                let mut buf = vec![0u8; chunk_size];
122                buf[..n].copy_from_slice(&rng_data[..n]);
123                buf[n..].copy_from_slice(&rng_data[..n_left]);
124
125                let hash = compute_data_hash(&buf);
126                chunks.push(Chunk {
127                    hash,
128                    data: Bytes::from(buf),
129                });
130            }
131
132            let raw_xorb = RawXorbData::from_chunks(&chunks, vec![0]);
133
134            shard.add_xorb_block(raw_xorb.xorb_info.clone())?;
135
136            let cfg = xet_runtime::config::XetConfig::new();
137            let serialized_xorb = SerializedXorbObject::from_xorb(
138                raw_xorb.clone(),
139                true,
140                cfg.xorb.compression_policy.as_str(),
141                cfg.xorb.compression_scheme_retest_interval,
142            )?;
143
144            let upload_permit = self.acquire_upload_permit().await?;
145            self.upload_xorb("default", serialized_xorb, None, upload_permit).await?;
146
147            xorb_data.insert(xorb_seed, raw_xorb);
148        }
149
150        // Build the file info and file data from RawXorbData.
151        let mut file_segments = Vec::new();
152        let mut file_data = Vec::new();
153        let mut chunk_file_hashes = Vec::new();
154        let mut term_infos = Vec::new();
155
156        for &(xorb_seed, (chunk_idx_start, chunk_idx_end)) in term_spec {
157            let raw_xorb = xorb_data.get(&xorb_seed).unwrap();
158            let xorb_h = raw_xorb.hash();
159
160            let (c_lb, c_ub) = (chunk_idx_start as usize, chunk_idx_end as usize);
161
162            let mut n_bytes = 0;
163            let mut term_data = Vec::new();
164            let mut term_chunk_hashes = Vec::new();
165
166            for i in c_lb..c_ub {
167                let chunk_bytes = &raw_xorb.data[i];
168                let chunk_hash = raw_xorb.xorb_info.chunks[i].chunk_hash;
169
170                file_data.extend_from_slice(chunk_bytes);
171                term_data.extend_from_slice(chunk_bytes);
172                n_bytes += chunk_bytes.len();
173                chunk_file_hashes.push((chunk_hash, chunk_bytes.len() as u64));
174                term_chunk_hashes.push(chunk_hash);
175            }
176
177            file_segments.push(FileDataSequenceEntry::new(
178                xorb_h,
179                n_bytes,
180                chunk_idx_start as usize,
181                chunk_idx_end as usize,
182            ));
183
184            term_infos.push(FileTermReference {
185                xorb_hash: xorb_h,
186                chunk_start: chunk_idx_start as u32,
187                chunk_end: chunk_idx_end as u32,
188                data: Bytes::from(term_data),
189                chunk_hashes: term_chunk_hashes,
190            });
191        }
192
193        let file_hash = file_hash_with_salt(&chunk_file_hashes, &[0; 32]);
194
195        shard.add_file_reconstruction_info(MDBFileInfo {
196            metadata: FileDataSequenceHeader::new(file_hash, file_segments.len(), false, false),
197            segments: file_segments,
198            verification: vec![],
199            metadata_ext: None,
200        })?;
201
202        let upload_permit = self.acquire_upload_permit().await?;
203        self.upload_shard(shard.to_bytes()?.into(), upload_permit).await?;
204
205        // Convert xorb_data from seed-keyed to hash-keyed
206        let xorbs = xorb_data.into_values().map(|x| (x.hash(), x)).collect();
207
208        Ok(RandomFileContents {
209            file_hash,
210            data: Bytes::from(file_data),
211            xorbs,
212            terms: term_infos,
213        })
214    }
215}
216
217impl<T: ?Sized + Client + Send + Sync> ClientTestingUtils for T {}
xet_client/cas_client/simulation/client_testing_utils.rs

xet_client/cas_client/simulation/
client_testing_utils.rs
