diff --git a/protocol-testing/Cargo.lock b/protocol-testing/Cargo.lock index 343020e..be60f27 100644 --- a/protocol-testing/Cargo.lock +++ b/protocol-testing/Cargo.lock @@ -2720,6 +2720,7 @@ dependencies = [ "postgres", "serde", "serde_json", + "serde_yaml", "tokio", "tracing", "tracing-subscriber", diff --git a/protocol-testing/Cargo.toml b/protocol-testing/Cargo.toml index c5fb3d9..bd76d81 100644 --- a/protocol-testing/Cargo.toml +++ b/protocol-testing/Cargo.toml @@ -19,3 +19,4 @@ serde = { version = "1.0.218", features = ["derive"] } hex = "0.4.3" tracing-subscriber = "0.3.19" postgres = "0.19.10" +serde_yaml = "0.9.34" diff --git a/protocol-testing/src/config.rs b/protocol-testing/src/config.rs index 2580f3d..2ebdff4 100644 --- a/protocol-testing/src/config.rs +++ b/protocol-testing/src/config.rs @@ -1,6 +1,5 @@ use std::collections::HashMap; -use figment::providers::Format; use hex::FromHex; use serde::{Deserialize, Serialize}; @@ -18,9 +17,9 @@ impl From for HexBytes { } } -impl Into for HexBytes { - fn into(self) -> String { - format!("0x{}", hex::encode(self.0)) +impl From for String { + fn from(val: HexBytes) -> Self { + format!("0x{}", hex::encode(val.0)) } } diff --git a/protocol-testing/src/utils.rs b/protocol-testing/src/utils.rs new file mode 100644 index 0000000..6175473 --- /dev/null +++ b/protocol-testing/src/utils.rs @@ -0,0 +1,134 @@ +use std::{error::Error, fs, path::Path, process::Command}; + +use figment::{ + providers::{Format, Yaml}, + value::Value, + Figment, +}; + +/// Build a Substreams package with modifications to the YAML file. +pub fn build_spkg(yaml_file_path: &str, modify_func: F) -> Result> +where + F: FnOnce(&mut Value) -> Result<(), Box>, +{ + // Create a backup file of the unmodified Substreams protocol YAML config file. + let backup_file_path = format!("{}.backup", yaml_file_path); + fs::copy(yaml_file_path, &backup_file_path)?; + + let figment = Figment::new().merge(Yaml::file(yaml_file_path)); + let mut data: Value = figment.extract()?; + + // Apply the modification function to update the YAML files + modify_func(&mut data).expect("Failed to modify the YAML config file."); + + let parent_dir = Path::new(yaml_file_path) + .parent() + .unwrap_or_else(|| Path::new("")) + .to_str() + .unwrap_or(""); + + let package_name = data + .clone() + .find("package") + .expect("Package not found on YAML") + .find("name") + .expect("Name not found on YAML") + .as_str() + .expect("Failed to convert name to string.") + .replace("_", "-"); + + let binding = data + .clone() + .find("package") + .expect("Package not found on YAML") + .find("version") + .expect("Version not found on YAML"); + + let package_version = binding.as_str().unwrap_or(""); + + let spkg_name = format!("{}/{}-{}.spkg", parent_dir, package_name, package_version); + + // Write the modified YAML back to the file + let yaml_string = serde_yaml::to_string(&data)?; + fs::write(yaml_file_path, yaml_string)?; + + // Run the substreams pack command to create the spkg + match Command::new("substreams") + .arg("pack") + .arg(yaml_file_path) + .output() + { + Ok(output) => { + if !output.status.success() { + println!( + "Substreams pack command failed: {}", + String::from_utf8_lossy(&output.stderr) + ); + } + } + Err(e) => { + println!("Error running substreams pack command: {}", e); + } + } + + // Restore the original YAML from backup + fs::copy(&backup_file_path, yaml_file_path)?; + fs::remove_file(&backup_file_path)?; + + Ok(spkg_name) +} + +/// Update the initial block for all modules in the configuration data. +pub fn modify_initial_block(data: &mut Value, start_block: usize) { + if let Value::Dict(_, ref mut dict) = data { + if let Some(Value::Array(_, modules)) = dict.get_mut("modules") { + for module in modules.iter_mut() { + if let Value::Dict(_, ref mut module_dict) = module { + module_dict.insert("initialBlock".to_string(), Value::from(start_block)); + } + } + } + } +} + +#[cfg(test)] +mod tests { + use figment::value::Value; + + use super::*; + + fn create_test_data() -> Value { + let file_path = Path::new("src/assets/substreams_example.yaml"); + let figment = Figment::new().merge(Yaml::file(file_path)); + + figment + .extract() + .expect("Failed to parse YAML file") + } + + #[test] + fn test_modify_initial_block_normal_case() { + let mut data = create_test_data(); + + // Apply modification + let new_block = 12345; + modify_initial_block(&mut data, new_block); + + // Verify all modules now have the correct initialBlock + if let Value::Dict(_, dict) = &data { + if let Some(Value::Array(_, modules)) = dict.get("modules") { + for module in modules { + if let Value::Dict(_, module_dict) = module { + if let Some(Value::Num(_, block)) = module_dict.get("initialBlock") { + assert_eq!(block.to_u128().unwrap(), new_block as u128); + } else { + panic!("initialBlock not found or has wrong type"); + } + } + } + } else { + panic!("modules not found or has wrong type"); + } + } + } +}