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