1use std::collections::BTreeMap;
2use std::collections::BTreeSet;
3use std::fmt::Write as _;
4use std::fs;
5use std::fs::File;
6use std::path::Path;
7
8use serde_json;
9
10pub type ArtifactId = String;
11pub type FileContent = String;
12
13#[derive(Debug, Default)]
16pub struct FakeAggregatorData {
17 epoch_settings: FileContent,
18
19 certificates_list: FileContent,
20 individual_certificates: BTreeMap<ArtifactId, FileContent>,
21
22 snapshots_list: FileContent,
23 individual_snapshots: BTreeMap<ArtifactId, FileContent>,
24
25 mithril_stake_distributions_list: FileContent,
26 individual_mithril_stake_distributions: BTreeMap<ArtifactId, FileContent>,
27
28 cardano_transaction_snapshots_list: FileContent,
29 individual_cardano_transaction_snapshots: BTreeMap<ArtifactId, FileContent>,
30 cardano_transaction_proofs: BTreeMap<ArtifactId, FileContent>,
31
32 cardano_stake_distributions_list: FileContent,
33 individual_cardano_stake_distributions: BTreeMap<ArtifactId, FileContent>,
34
35 cardano_database_snapshots_list: FileContent,
36 individual_cardano_database_snapshots: BTreeMap<ArtifactId, FileContent>,
37}
38
39impl FakeAggregatorData {
40 pub fn load_from_folder(folder: &Path) -> Self {
41 let mut data = FakeAggregatorData::default();
42
43 for entry in list_json_files_in_folder(folder) {
44 let filename = entry.file_name().to_string_lossy().to_string();
45 let file_content = fs::read_to_string(entry.path()).unwrap_or_else(|_| {
46 panic!(
47 "Could not read file content, file_path: {}",
48 entry.path().display()
49 )
50 });
51
52 match filename.as_str() {
53 "epoch-settings.json" => {
54 data.epoch_settings = file_content;
55 }
56 "mithril-stake-distributions-list.json" => {
57 data.mithril_stake_distributions_list = file_content;
58 }
59 "snapshots-list.json" => {
60 data.snapshots_list = file_content;
61 }
62 "cardano-stake-distributions-list.json" => {
63 data.cardano_stake_distributions_list = file_content;
64 }
65 "cardano-databases-list.json" => {
66 data.cardano_database_snapshots_list = file_content;
67 }
68 "certificates-list.json" => {
69 data.certificates_list = file_content;
70 }
71 "ctx-snapshots-list.json" => {
72 data.cardano_transaction_snapshots_list = file_content;
73 }
74 "mithril-stake-distributions.json" => {
75 data.individual_mithril_stake_distributions =
76 Self::read_artifacts_json_file(&entry.path());
77 }
78 "snapshots.json" => {
79 data.individual_snapshots = Self::read_artifacts_json_file(&entry.path());
80 }
81 "cardano-stake-distributions.json" => {
82 data.individual_cardano_stake_distributions =
83 Self::read_artifacts_json_file(&entry.path());
84 }
85 "cardano-databases.json" => {
86 data.individual_cardano_database_snapshots =
87 Self::read_artifacts_json_file(&entry.path());
88 }
89 "certificates.json" => {
90 data.individual_certificates = Self::read_artifacts_json_file(&entry.path());
91 }
92 "ctx-snapshots.json" => {
93 data.individual_cardano_transaction_snapshots =
94 Self::read_artifacts_json_file(&entry.path());
95 }
96 "ctx-proofs.json" => {
97 data.cardano_transaction_proofs = Self::read_artifacts_json_file(&entry.path());
98 }
99 _ => {}
101 }
102 }
103
104 data
105 }
106
107 pub fn generate_code_for_ids(self) -> String {
108 Self::assemble_code(
109 &[
110 generate_ids_array(
111 "snapshot_digests",
112 BTreeSet::from_iter(self.individual_snapshots.keys().cloned()),
113 ),
114 generate_ids_array(
115 "mithril_stake_distribution_hashes",
116 BTreeSet::from_iter(
117 self.individual_mithril_stake_distributions.keys().cloned(),
118 ),
119 ),
120 generate_ids_array(
121 "cardano_stake_distribution_hashes",
122 BTreeSet::from_iter(
123 self.individual_cardano_stake_distributions.keys().cloned(),
124 ),
125 ),
126 generate_ids_array(
127 "cardano_stake_distribution_epochs",
128 BTreeSet::from_iter(extract_cardano_stake_distribution_epochs(
129 &self.individual_cardano_stake_distributions,
130 )),
131 ),
132 generate_ids_array(
133 "cardano_database_snapshot_hashes",
134 BTreeSet::from_iter(self.individual_cardano_database_snapshots.keys().cloned()),
135 ),
136 generate_ids_array(
137 "certificate_hashes",
138 BTreeSet::from_iter(self.individual_certificates.keys().cloned()),
139 ),
140 generate_ids_array(
141 "cardano_transaction_snapshot_hashes",
142 BTreeSet::from_iter(
143 self.individual_cardano_transaction_snapshots.keys().cloned(),
144 ),
145 ),
146 generate_ids_array(
147 "proof_transaction_hashes",
148 BTreeSet::from_iter(self.cardano_transaction_proofs.keys().cloned()),
149 ),
150 ],
151 false,
152 )
153 }
154
155 pub fn generate_code_for_all_data(self) -> String {
156 Self::assemble_code(
157 &[
158 generate_list_getter("epoch_settings", self.epoch_settings),
159 generate_ids_array(
160 "snapshot_digests",
161 BTreeSet::from_iter(self.individual_snapshots.keys().cloned()),
162 ),
163 generate_artifact_getter("snapshots", self.individual_snapshots),
164 generate_list_getter("snapshot_list", self.snapshots_list),
165 generate_ids_array(
166 "mithril_stake_distribution_hashes",
167 BTreeSet::from_iter(
168 self.individual_mithril_stake_distributions.keys().cloned(),
169 ),
170 ),
171 generate_artifact_getter(
172 "mithril_stake_distributions",
173 self.individual_mithril_stake_distributions,
174 ),
175 generate_list_getter(
176 "mithril_stake_distribution_list",
177 self.mithril_stake_distributions_list,
178 ),
179 generate_ids_array(
180 "cardano_stake_distribution_hashes",
181 BTreeSet::from_iter(
182 self.individual_cardano_stake_distributions.keys().cloned(),
183 ),
184 ),
185 generate_ids_array(
186 "cardano_stake_distribution_epochs",
187 BTreeSet::from_iter(extract_cardano_stake_distribution_epochs(
188 &self.individual_cardano_stake_distributions,
189 )),
190 ),
191 generate_artifact_getter(
192 "cardano_stake_distributions",
193 self.individual_cardano_stake_distributions,
194 ),
195 generate_list_getter(
196 "cardano_stake_distribution_list",
197 self.cardano_stake_distributions_list,
198 ),
199 generate_ids_array(
200 "certificate_hashes",
201 BTreeSet::from_iter(self.individual_certificates.keys().cloned()),
202 ),
203 generate_ids_array(
204 "cardano_database_snapshot_hashes",
205 BTreeSet::from_iter(self.individual_cardano_database_snapshots.keys().cloned()),
206 ),
207 generate_artifact_getter(
208 "cardano_database_snapshots",
209 self.individual_cardano_database_snapshots,
210 ),
211 generate_list_getter(
212 "cardano_database_snapshot_list",
213 self.cardano_database_snapshots_list,
214 ),
215 generate_artifact_getter("certificates", self.individual_certificates),
216 generate_list_getter("certificate_list", self.certificates_list),
217 generate_ids_array(
218 "cardano_transaction_snapshot_hashes",
219 BTreeSet::from_iter(
220 self.individual_cardano_transaction_snapshots.keys().cloned(),
221 ),
222 ),
223 generate_artifact_getter(
224 "cardano_transaction_snapshots",
225 self.individual_cardano_transaction_snapshots,
226 ),
227 generate_list_getter(
228 "cardano_transaction_snapshots_list",
229 self.cardano_transaction_snapshots_list,
230 ),
231 generate_ids_array(
232 "proof_transaction_hashes",
233 BTreeSet::from_iter(self.cardano_transaction_proofs.keys().cloned()),
234 ),
235 generate_artifact_getter(
236 "cardano_transaction_proofs",
237 self.cardano_transaction_proofs,
238 ),
239 ],
240 true,
241 )
242 }
243
244 fn assemble_code(functions_code: &[String], include_use_btree_map: bool) -> String {
245 format!(
246 "{}{}
247",
248 if include_use_btree_map {
249 "use std::collections::BTreeMap;
250
251"
252 } else {
253 ""
254 },
255 functions_code.join(
256 "
257
258"
259 )
260 )
261 }
262
263 fn read_artifacts_json_file(json_file: &Path) -> BTreeMap<ArtifactId, FileContent> {
264 let file = File::open(json_file).unwrap();
265 let parsed_json: serde_json::Value = serde_json::from_reader(&file).unwrap();
266
267 let json_object = parsed_json.as_object().unwrap();
268 let res: Result<Vec<_>, _> = json_object
269 .iter()
270 .map(|(key, value)| extract_artifact_id_and_content(key, value))
271 .collect();
272
273 BTreeMap::from_iter(res.unwrap())
274 }
275}
276
277pub fn extract_cardano_stake_distribution_epochs(
278 individual_csds: &BTreeMap<ArtifactId, FileContent>,
279) -> Vec<String> {
280 individual_csds
281 .values()
282 .map(|content| {
283 let json_value: serde_json::Value =
284 serde_json::from_str(content).unwrap_or_else(|err| {
285 panic!("Failed to parse JSON in csd content: {content}\nError: {err}");
286 });
287
288 json_value
289 .get("epoch")
290 .and_then(|epoch| epoch.as_u64().map(|s| s.to_string()))
291 .unwrap_or_else(|| {
292 panic!("Epoch not found or invalid in csd content: {content}");
293 })
294 })
295 .collect()
296}
297
298fn extract_artifact_id_and_content(
299 key: &String,
300 value: &serde_json::Value,
301) -> Result<(ArtifactId, FileContent), String> {
302 let json_content = serde_json::to_string_pretty(value).map_err(|e| e.to_string())?;
303 Ok((key.to_owned(), json_content))
304}
305
306pub fn list_json_files_in_folder(folder: &Path) -> impl Iterator<Item = fs::DirEntry> + '_ {
307 crate::list_files_in_folder(folder)
308 .filter(|e| e.file_name().to_string_lossy().ends_with(".json"))
309}
310
311pub fn generate_artifact_getter(
313 fun_name: &str,
314 source_jsons: BTreeMap<ArtifactId, FileContent>,
315) -> String {
316 let mut artifacts_list = String::new();
317
318 for (artifact_id, file_content) in source_jsons {
319 write!(
320 artifacts_list,
321 r###"
322 (
323 "{artifact_id}",
324 r#"{file_content}"#
325 ),"###
326 )
327 .unwrap();
328 }
329
330 format!(
331 r###"pub(crate) fn {fun_name}() -> BTreeMap<String, String> {{
332 [{artifacts_list}
333 ]
334 .into_iter()
335 .map(|(k, v)| (k.to_owned(), v.to_owned()))
336 .collect()
337}}"###
338 )
339}
340
341pub fn generate_list_getter(fun_name: &str, source_json: FileContent) -> String {
343 format!(
344 r###"pub(crate) fn {fun_name}() -> &'static str {{
345 r#"{source_json}"#
346}}"###
347 )
348}
349
350pub fn generate_ids_array(array_name: &str, ids: BTreeSet<ArtifactId>) -> String {
352 let mut ids_list = String::new();
353
354 for id in &ids {
355 write!(
356 ids_list,
357 r#"
358 "{id}","#
359 )
360 .unwrap();
361 }
362
363 format!(
364 r###"pub(crate) const fn {}<'a>() -> [&'a str; {}] {{
365 [{}
366 ]
367}}"###,
368 array_name,
369 ids.len(),
370 ids_list,
371 )
372}
373
374#[cfg(test)]
375mod tests {
376 use crate::get_temp_dir;
377
378 use super::*;
379
380 #[test]
381 fn generate_ids_array_with_empty_data() {
382 assert_eq!(
383 generate_ids_array("snapshots_digests", BTreeSet::new()),
384 "pub(crate) const fn snapshots_digests<'a>() -> [&'a str; 0] {
385 [
386 ]
387}"
388 );
389 }
390
391 #[test]
392 fn generate_ids_array_with_non_empty_data() {
393 assert_eq!(
394 generate_ids_array(
395 "snapshots_digests",
396 BTreeSet::from_iter(["abc".to_string(), "def".to_string(), "hij".to_string()])
397 ),
398 r#"pub(crate) const fn snapshots_digests<'a>() -> [&'a str; 3] {
399 [
400 "abc",
401 "def",
402 "hij",
403 ]
404}"#
405 );
406 }
407
408 #[test]
409 fn assemble_code_with_btree_use() {
410 assert_eq!(
411 "use std::collections::BTreeMap;
412
413fn a() {}
414
415fn b() {}
416",
417 FakeAggregatorData::assemble_code(
418 &["fn a() {}".to_string(), "fn b() {}".to_string()],
419 true
420 )
421 )
422 }
423
424 #[test]
425 fn assemble_code_without_btree_use() {
426 assert_eq!(
427 "fn a() {}
428
429fn b() {}
430",
431 FakeAggregatorData::assemble_code(
432 &["fn a() {}".to_string(), "fn b() {}".to_string()],
433 false
434 )
435 )
436 }
437
438 #[test]
439 fn parse_artifacts_json_into_btree_of_key_and_pretty_sub_json() {
440 let dir = get_temp_dir("read_artifacts_json_file");
441 let file = dir.join("test.json");
442 fs::write(
443 &file,
444 r#"{
445 "hash1": { "name": "artifact1" },
446 "hash2": { "name": "artifact2" }
447}"#,
448 )
449 .unwrap();
450
451 let id_per_json = FakeAggregatorData::read_artifacts_json_file(&file);
452
453 let expected = BTreeMap::from([
454 (
455 "hash1".to_string(),
456 r#"{
457 "name": "artifact1"
458}"#
459 .to_string(),
460 ),
461 (
462 "hash2".to_string(),
463 r#"{
464 "name": "artifact2"
465}"#
466 .to_string(),
467 ),
468 ]);
469 assert_eq!(expected, id_per_json);
470 }
471
472 #[test]
473 fn extract_csd_epochs_with_valid_data() {
474 let mut csds = BTreeMap::new();
475 csds.insert(
476 "csd-123".to_string(),
477 r#"{"hash": "csd-123", "epoch": 123}"#.to_string(),
478 );
479 csds.insert(
480 "csd-456".to_string(),
481 r#"{"hash": "csd-456", "epoch": 456}"#.to_string(),
482 );
483
484 let epochs = extract_cardano_stake_distribution_epochs(&csds);
485
486 assert_eq!(epochs, vec![123.to_string(), 456.to_string()]);
487 }
488
489 #[test]
490 #[should_panic(expected = "Failed to parse JSON in csd content")]
491 fn extract_csd_epochs_with_invalid_json() {
492 let mut csds = BTreeMap::new();
493 csds.insert(
494 "csd-123".to_string(),
495 r#""hash": "csd-123", "epoch": "123"#.to_string(),
496 );
497
498 extract_cardano_stake_distribution_epochs(&csds);
499 }
500
501 #[test]
502 #[should_panic(expected = "Epoch not found or invalid in csd content")]
503 fn test_extract_csd_epochs_with_missing_epoch() {
504 let mut csds = BTreeMap::new();
505 csds.insert("csd-123".to_string(), r#"{"hash": "csd-123"}"#.to_string());
506
507 extract_cardano_stake_distribution_epochs(&csds);
508 }
509
510 #[test]
511 fn test_extract_csd_epochs_with_empty_map() {
512 let csds = BTreeMap::new();
513
514 let epochs = extract_cardano_stake_distribution_epochs(&csds);
515
516 assert!(epochs.is_empty());
517 }
518}