feat(iota-genesis-builder): snapshot test data setup (#761)

* initial add_snapshot_test_data setup

* increment snapshot header output count

* add new outputs to snapshot

* change the header total supply

* add some comments

* rename and add comment

* rename example

crates/iota-genesis-builder/examples/snapshot_test_outputs.rs

@@ -0,0 +1,44 @@
+// Copyright (c) 2024 IOTA Stiftung
+// SPDX-License-Identifier: Apache-2.0
+
+//! Example to add test outputs to a full snapshot.
+
+use std::{fs::File, path::Path};
+
+use iota_genesis_builder::stardust::{
+    parse::FullSnapshotParser, test_outputs::add_snapshot_test_outputs,
+};
+
+fn parse_snapshot<P: AsRef<Path>>(path: P) -> anyhow::Result<()> {
+    let file = File::open(path)?;
+    let parser = FullSnapshotParser::new(file)?;
+
+    println!("Output count: {}", parser.header.output_count());
+
+    let total_supply_header = parser.total_supply()?;
+    let total_supply_outputs = parser.outputs().try_fold(0, |acc, output| {
+        Ok::<_, anyhow::Error>(acc + output?.1.amount())
+    })?;
+
+    assert_eq!(total_supply_header, total_supply_outputs);
+
+    println!("Total supply: {total_supply_header}");
+
+    Ok(())
+}
+
+fn main() -> anyhow::Result<()> {
+    let Some(current_path) = std::env::args().nth(1) else {
+        anyhow::bail!("please provide path to the full-snapshot file");
+    };
+    let mut new_path = String::from("test-");
+    new_path.push_str(&current_path);
+
+    parse_snapshot(&current_path)?;
+
+    add_snapshot_test_outputs(&current_path, &new_path)?;
+
+    parse_snapshot(&new_path)?;
+
+    Ok(())
+}

crates/iota-genesis-builder/src/stardust/mod.rs

@@ -9,4 +9,5 @@ pub mod error;
 pub mod migration;
 pub mod native_token;
 pub mod parse;
+pub mod test_outputs;
 pub mod types;

crates/iota-genesis-builder/src/stardust/parse.rs

@@ -48,17 +48,22 @@ impl<R: Read> FullSnapshotParser<R> {
         self.header.target_milestone_timestamp()
     }
 
-    /// Provide the network main token total supply through the snapshot
-    /// protocol parameters.
-    pub fn total_supply(&self) -> Result<u64> {
+    /// Provide the protocol parameters extracted from the snapshot header.
+    pub fn protocol_parameters(&self) -> Result<ProtocolParameters> {
         if let MilestoneOption::Parameters(params) = self.header.parameters_milestone_option() {
             let protocol_params = <ProtocolParameters as packable::PackableExt>::unpack_unverified(
                 params.binary_parameters(),
             )
             .expect("invalid protocol params");
-            Ok(protocol_params.token_supply())
+            Ok(protocol_params)
         } else {
             Err(StardustError::HornetSnapshotParametersNotFound.into())
         }
     }
+
+    /// Provide the network main token total supply through the snapshot
+    /// protocol parameters.
+    pub fn total_supply(&self) -> Result<u64> {
+        self.protocol_parameters().map(|p| p.token_supply())
+    }
 }

crates/iota-genesis-builder/src/stardust/test_outputs/dummy.rs

@@ -0,0 +1,38 @@
+// Copyright (c) 2024 IOTA Stiftung
+// SPDX-License-Identifier: Apache-2.0
+
+use std::str::FromStr;
+
+use iota_sdk::types::block::{
+    address::Ed25519Address,
+    output::{unlock_condition::AddressUnlockCondition, BasicOutputBuilder, Output},
+    payload::transaction::TransactionId,
+};
+
+use crate::stardust::types::snapshot::OutputHeader;
+
+pub(crate) fn outputs() -> Vec<(OutputHeader, Output)> {
+    let mut outputs = Vec::new();
+
+    let output_header = OutputHeader::new_testing(
+        *TransactionId::from_str(
+            "0xb191c4bc825ac6983789e50545d5ef07a1d293a98ad974fc9498cb1812345678",
+        )
+        .unwrap(),
+        rand::random(),
+        rand::random(),
+        rand::random(),
+    );
+    let output = Output::from(
+        BasicOutputBuilder::new_with_amount(1_000_000)
+            .add_unlock_condition(AddressUnlockCondition::new(Ed25519Address::from(
+                rand::random::<[u8; Ed25519Address::LENGTH]>(),
+            )))
+            .finish()
+            .unwrap(),
+    );
+
+    outputs.push((output_header, output));
+
+    outputs
+}

crates/iota-genesis-builder/src/stardust/test_outputs/mod.rs

@@ -0,0 +1,73 @@
+// Copyright (c) 2024 IOTA Stiftung
+// SPDX-License-Identifier: Apache-2.0
+
+mod dummy;
+
+use std::{
+    fs::{File, OpenOptions},
+    io::BufWriter,
+    path::Path,
+};
+
+use iota_sdk::types::block::{
+    payload::milestone::{MilestoneOption, ParametersMilestoneOption},
+    protocol::ProtocolParameters,
+};
+use packable::{packer::IoPacker, Packable, PackableExt};
+
+use crate::stardust::parse::FullSnapshotParser;
+
+/// Adds outputs to test specific and intricate scenario in the full snapshot.
+pub fn add_snapshot_test_outputs<P: AsRef<Path> + core::fmt::Debug>(
+    current_path: P,
+    new_path: P,
+) -> anyhow::Result<()> {
+    let current_file = File::open(current_path)?;
+    let new_file = OpenOptions::new()
+        .write(true)
+        .create(true)
+        .truncate(true)
+        .open(new_path)?;
+    let mut writer = IoPacker::new(BufWriter::new(new_file));
+    let mut parser = FullSnapshotParser::new(current_file)?;
+
+    let new_outputs = dummy::outputs();
+
+    // Increments the output count according to newly generated outputs.
+    parser.header.output_count += new_outputs.len() as u64;
+
+    // Creates new protocol parameters to increase the total supply according to newly generated outputs.
+    let params = parser.protocol_parameters()?;
+    let new_params = ProtocolParameters::new(
+        params.protocol_version(),
+        params.network_name().to_owned(),
+        params.bech32_hrp(),
+        params.min_pow_score(),
+        params.below_max_depth(),
+        *params.rent_structure(),
+        params.token_supply() + new_outputs.iter().map(|o| o.1.amount()).sum::<u64>(),
+    )?;
+    if let MilestoneOption::Parameters(params) = &parser.header.parameters_milestone_option {
+        parser.header.parameters_milestone_option =
+            MilestoneOption::Parameters(ParametersMilestoneOption::new(
+                params.target_milestone_index(),
+                params.protocol_version(),
+                new_params.pack_to_vec(),
+            )?);
+    }
+
+    // Writes the new header.
+    parser.header.pack(&mut writer)?;
+
+    // Writes previous and new outputs.
+    parser
+        .outputs()
+        .filter_map(|o| o.ok())
+        .chain(new_outputs)
+        .for_each(|(output_header, output)| {
+            output_header.pack(&mut writer).unwrap();
+            output.pack(&mut writer).unwrap();
+        });
+
+    Ok(())
+}

crates/iota-genesis-builder/src/stardust/types/snapshot.rs

@@ -109,8 +109,8 @@ pub struct FullSnapshotHeader {
     ledger_milestone_index: MilestoneIndex,
     treasury_output_milestone_id: MilestoneId,
     treasury_output_amount: u64,
-    parameters_milestone_option: MilestoneOption,
-    output_count: u64,
+    pub(crate) parameters_milestone_option: MilestoneOption,
+    pub(crate) output_count: u64,
     milestone_diff_count: u32,
     sep_count: u16,
 }