test: add generation arguments

etorreborre · etorreborre · commit 190584bae6d9 · 2025-10-17T10:26:11.000+02:00
Signed-off-by: etorreborre &lt;etorreborre@yahoo.com&gt;
diff --git a/crates/amaru-consensus/src/consensus/headers_tree/data_generation/actions.rs b/crates/amaru-consensus/src/consensus/headers_tree/data_generation/actions.rs
@@ -44,6 +44,7 @@ use serde::{Deserialize, Deserializer, Serialize, Serializer};
 use std::cmp::Reverse;
 use std::collections::BTreeMap;
 use std::fmt::{Debug, Display, Formatter};
+use std::str::FromStr;
 use std::sync::Arc;
 
 /// This data type models the events sent by the ChainSync mini-protocol with simplified data for the tests.
@@ -315,6 +316,30 @@ pub fn random_walk(
 #[derive(Copy, Clone, PartialEq, Eq, Debug)]
 pub struct Ratio(pub u32, pub u32);
 
+impl Display for Ratio {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}/{}", self.0, self.1)
+    }
+}
+
+impl FromStr for Ratio {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let parts: Vec<&str> = s.split('/').collect();
+        if parts.len() != 2 {
+            return Err("Ratio must be in the form 'numerator/denominator'".to_string());
+        }
+        let numerator = parts[0]
+            .parse::<u32>()
+            .map_err(|_| format!("Invalid numerator in ratio: {}", parts[0]))?;
+        let denominator = parts[1]
+            .parse::<u32>()
+            .map_err(|_| format!("Invalid denominator in ratio: {}", parts[1]))?;
+        Ok(Ratio(numerator, denominator))
+    }
+}
+
 /// Generate random walks for a fixed number of peers on a given tree of headers.
 ///
 /// The returned list of actions is transposed so that the actions from different peers are interleaved.
@@ -358,8 +383,9 @@ pub fn any_select_chains(
     peers_nb: usize,
     depth: usize,
     rollback_ratio: Ratio,
+    branching_ratio: Ratio,
 ) -> impl Strategy<Value = Vec<Action>> {
-    any_tree_of_headers(depth, Ratio(1, 2)).prop_flat_map(move |tree| {
+    any_tree_of_headers(depth, branching_ratio).prop_flat_map(move |tree| {
         (1..u64::MAX)
             .prop_map(move |seed| generate_random_walks(&tree, peers_nb, rollback_ratio, seed))
     })
diff --git a/crates/amaru-consensus/src/consensus/headers_tree/headers_tree.rs b/crates/amaru-consensus/src/consensus/headers_tree/headers_tree.rs
@@ -1284,11 +1284,12 @@ mod tests {
     const MAX_LENGTH: usize = 5;
     const TEST_CASES_NB: u32 = 1000;
     const ROLLBACK_RATIO: Ratio = Ratio(1, 2);
+    const BRANCHING_RATIO: Ratio = Ratio(1, 2);
 
     proptest! {
         #![proptest_config(config_begin().no_shrink().with_cases(TEST_CASES_NB).end())]
         #[test]
-        fn run_chain_selection(actions in any_select_chains(DEPTH, 5, ROLLBACK_RATIO)) {
+        fn run_chain_selection(actions in any_select_chains(DEPTH, 5, ROLLBACK_RATIO, BRANCHING_RATIO)) {
             let results = execute_actions(MAX_LENGTH, &actions, false).unwrap();
             let actual_chains = make_best_chains_from_results(&results);
             let expected_chains = make_best_chains_from_actions(&actions);
diff --git a/simulation/amaru-sim/src/simulator/args.rs b/simulation/amaru-sim/src/simulator/args.rs
@@ -12,7 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use amaru_consensus::consensus::headers_tree::data_generation::Ratio;
 use clap::Parser;
+use std::str::FromStr;
 
 #[derive(Debug, Parser, Clone)]
 #[clap(name = "Amaru Simulator")]
@@ -39,6 +41,14 @@ pub struct Args {
     #[arg(long, default_value = "10")]
     pub generated_chain_depth: u64,
 
+    /// Ratio of rollbacks in the generated chain for a given peer
+    #[arg(long, default_value = "1/2", value_parser = Ratio::from_str)]
+    pub generated_chain_rollback_ratio: Ratio,
+
+    /// Ratio of branches generated from a central chain that can be explored by peers during input generation
+    #[arg(long, default_value = "1/2", value_parser = Ratio::from_str)]
+    pub generated_chain_branching_ratio: Ratio,
+
     #[arg(long)]
     pub disable_shrinking: bool,
 
diff --git a/simulation/amaru-sim/src/simulator/data_generation/generate.rs b/simulation/amaru-sim/src/simulator/data_generation/generate.rs
@@ -18,7 +18,7 @@ use crate::simulator::data_generation::base_generators::generate_arrival_times;
 use crate::simulator::{Entry, NodeConfig};
 use crate::sync::ChainSyncMessage;
 use amaru_consensus::IsHeader;
-use amaru_consensus::consensus::headers_tree::data_generation::{Action, Ratio, any_select_chains};
+use amaru_consensus::consensus::headers_tree::data_generation::{Action, any_select_chains};
 use amaru_kernel::peer::Peer;
 use amaru_kernel::to_cbor;
 use amaru_ouroboros_traits::tests::run_with_rng;
@@ -37,7 +37,8 @@ pub fn generate_entries<R: Rng>(
             any_select_chains(
                 node_config.number_of_upstream_peers as usize,
                 node_config.generated_chain_depth as usize,
-                Ratio(1, 2),
+                node_config.generated_chain_rollback_ratio,
+                node_config.generated_chain_branching_ratio,
             ),
         );
 
diff --git a/simulation/amaru-sim/src/simulator/node_config.rs b/simulation/amaru-sim/src/simulator/node_config.rs
@@ -13,13 +13,16 @@
 // limitations under the License.
 
 use crate::simulator::Args;
+use amaru_consensus::consensus::headers_tree::data_generation::Ratio;
 
 /// Configuration for a single node
 #[derive(Debug, Clone)]
 pub struct NodeConfig {
     pub number_of_upstream_peers: u8,
     pub number_of_downstream_peers: u8,
     pub generated_chain_depth: u64,
+    pub generated_chain_rollback_ratio: Ratio,
+    pub generated_chain_branching_ratio: Ratio,
 }
 
 impl Default for NodeConfig {
@@ -28,6 +31,8 @@ impl Default for NodeConfig {
             number_of_upstream_peers: 2,
             number_of_downstream_peers: 1,
             generated_chain_depth: 10,
+            generated_chain_rollback_ratio: Ratio(1, 2),
+            generated_chain_branching_ratio: Ratio(1, 2),
         }
     }
 }
@@ -38,6 +43,8 @@ impl NodeConfig {
             number_of_upstream_peers: args.number_of_upstream_peers,
             number_of_downstream_peers: args.number_of_downstream_peers,
             generated_chain_depth: args.generated_chain_depth,
+            generated_chain_rollback_ratio: args.generated_chain_rollback_ratio,
+            generated_chain_branching_ratio: args.generated_chain_branching_ratio,
         }
     }
 
@@ -55,4 +62,14 @@ impl NodeConfig {
         self.generated_chain_depth = depth;
         self
     }
+
+    pub fn with_generated_chain_rollback_ratio(mut self, ratio: Ratio) -> Self {
+        self.generated_chain_rollback_ratio = ratio;
+        self
+    }
+
+    pub fn with_generated_chain_branching_ratio(mut self, ratio: Ratio) -> Self {
+        self.generated_chain_branching_ratio = ratio;
+        self
+    }
 }
diff --git a/simulation/amaru-sim/src/simulator/run.rs b/simulation/amaru-sim/src/simulator/run.rs
@@ -85,7 +85,7 @@ pub fn run(rt: Runtime, args: Args) {
         trace_buffer.clone(),
         args.persist_on_success,
     )
-    .unwrap_or_else(|e| panic!("{e}"));
+        .unwrap_or_else(|e| panic!("{e}"));
 }
 
 /// Create and start a node
@@ -127,7 +127,7 @@ pub fn spawn_node(
                             },
                         },
                     )
-                    .await
+                        .await
                 }
                 ChainSyncMessage::InitOk { .. } => (),
                 ChainSyncMessage::Fwd {
@@ -142,7 +142,7 @@ pub fn spawn_node(
                             span: Span::current(),
                         },
                     )
-                    .await
+                        .await
                 }
                 ChainSyncMessage::Bck { slot, hash, .. } => {
                     eff.send(
@@ -153,7 +153,7 @@ pub fn spawn_node(
                             span: Span::current(),
                         },
                     )
-                    .await
+                        .await
                 }
             }
             (downstream, output)
@@ -238,9 +238,9 @@ fn make_chain_selector(
 }
 
 /// Property: at the end of the simulation, the chain built from the history of messages received
-/// must match one of the best chains that could be built from the entries.
+/// downstream must match the best chain built directly from messages coming from upstream peers.
 fn chain_property()
--> impl Fn(&[Entry<ChainSyncMessage>], &History<ChainSyncMessage>) -> Result<(), String> {
+    -> impl Fn(&[Entry<ChainSyncMessage>], &History<ChainSyncMessage>) -> Result<(), String> {
     move |entries, history| {
         let expected = make_best_chain_from_entries(
             &entries
@@ -289,7 +289,7 @@ pub fn make_best_chain_from_entries(messages: &[Envelope<ChainSyncMessage>]) ->
                 // make sure to skip header hashes that cannot be decoded
                 if let Some(header_hash) = msg.header_hash()
                     && let Some(rollback_position) =
-                        chain.iter().position(|h| h.hash() == header_hash)
+                    chain.iter().position(|h| h.hash() == header_hash)
                 {
                     chain.truncate(rollback_position + 1)
                 }
diff --git a/simulation/amaru-sim/tests/simulation.rs b/simulation/amaru-sim/tests/simulation.rs
@@ -40,6 +40,14 @@ fn run_simulator() {
             "AMARU_GENERATED_CHAIN_DEPTH",
             node_config.generated_chain_depth,
         ),
+        generated_chain_rollback_ratio: get_env_var(
+            "AMARU_GENERATED_CHAIN_ROLLBACK_RATIO",
+            node_config.generated_chain_rollback_ratio,
+        ),
+        generated_chain_branching_ratio: get_env_var(
+            "AMARU_GENERATED_CHAIN_BRANCHING_RATIO",
+            node_config.generated_chain_branching_ratio,
+        ),
         disable_shrinking: is_true("AMARU_DISABLE_SHRINKING"),
         seed: get_optional_env_var("AMARU_TEST_SEED"),
         persist_on_success: is_true("AMARU_PERSIST_ON_SUCCESS"),
diff --git a/simulation/amaru-sim/tests/traces.rs b/simulation/amaru-sim/tests/traces.rs
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use amaru_consensus::consensus::headers_tree::data_generation::Ratio;
 use amaru_sim::simulator::run::spawn_node;
 use amaru_sim::simulator::simulate::simulate;
 use amaru_sim::simulator::{Args, NodeConfig, NodeHandle, SimulateConfig, generate_entries};
@@ -31,6 +32,8 @@ fn run_simulator_with_traces() {
         number_of_upstream_peers: 1,
         number_of_downstream_peers: 1,
         generated_chain_depth: 1,
+        generated_chain_rollback_ratio: Ratio(1, 2),
+        generated_chain_branching_ratio: Ratio(1, 2),
         disable_shrinking: true,
         seed: None,
         persist_on_success: false,

Original file line number	Diff line number	Diff line change
`@@ -13,13 +13,16 @@`
`13`	`13`	`// limitations under the License.`
`14`	`14`
`15`	`15`	`use crate::simulator::Args;`
	`16`	`+use amaru_consensus::consensus::headers_tree::data_generation::Ratio;`
`16`	`17`
`17`	`18`	`/// Configuration for a single node`
`18`	`19`	`#[derive(Debug, Clone)]`
`19`	`20`	`pub struct NodeConfig {`
`20`	`21`	`pub number_of_upstream_peers: u8,`
`21`	`22`	`pub number_of_downstream_peers: u8,`
`22`	`23`	`pub generated_chain_depth: u64,`
	`24`	`+ pub generated_chain_rollback_ratio: Ratio,`
	`25`	`+ pub generated_chain_branching_ratio: Ratio,`
`23`	`26`	`}`
`24`	`27`
`25`	`28`	`impl Default for NodeConfig {`
`@@ -28,6 +31,8 @@ impl Default for NodeConfig {`
`28`	`31`	`number_of_upstream_peers: 2,`
`29`	`32`	`number_of_downstream_peers: 1,`
`30`	`33`	`generated_chain_depth: 10,`
	`34`	`+ generated_chain_rollback_ratio: Ratio(1, 2),`
	`35`	`+ generated_chain_branching_ratio: Ratio(1, 2),`
`31`	`36`	`}`
`32`	`37`	`}`
`33`	`38`	`}`
`@@ -38,6 +43,8 @@ impl NodeConfig {`
`38`	`43`	`number_of_upstream_peers: args.number_of_upstream_peers,`
`39`	`44`	`number_of_downstream_peers: args.number_of_downstream_peers,`
`40`	`45`	`generated_chain_depth: args.generated_chain_depth,`
	`46`	`+ generated_chain_rollback_ratio: args.generated_chain_rollback_ratio,`
	`47`	`+ generated_chain_branching_ratio: args.generated_chain_branching_ratio,`
`41`	`48`	`}`
`42`	`49`	`}`
`43`	`50`
`@@ -55,4 +62,14 @@ impl NodeConfig {`
`55`	`62`	`self.generated_chain_depth = depth;`
`56`	`63`	`self`
`57`	`64`	`}`
	`65`	`+`
	`66`	`+ pub fn with_generated_chain_rollback_ratio(mut self, ratio: Ratio) -> Self {`
	`67`	`+ self.generated_chain_rollback_ratio = ratio;`
	`68`	`+ self`
	`69`	`+ }`
	`70`	`+`
	`71`	`+ pub fn with_generated_chain_branching_ratio(mut self, ratio: Ratio) -> Self {`
	`72`	`+ self.generated_chain_branching_ratio = ratio;`
	`73`	`+ self`
	`74`	`+ }`
`58`	`75`	`}`
Original file line number	Diff line number	Diff line change
`@@ -85,7 +85,7 @@ pub fn run(rt: Runtime, args: Args) {`
`85`	`85`	`trace_buffer.clone(),`
`86`	`86`	`args.persist_on_success,`
`87`	`87`	`)`
`88`		`- .unwrap_or_else(\|e\| panic!("{e}"));`
	`88`	`+ .unwrap_or_else(\|e\| panic!("{e}"));`
`89`	`89`	`}`
`90`	`90`
`91`	`91`	`/// Create and start a node`
`@@ -127,7 +127,7 @@ pub fn spawn_node(`
`127`	`127`	`},`
`128`	`128`	`},`
`129`	`129`	`)`
`130`		`- .await`
	`130`	`+ .await`
`131`	`131`	`}`
`132`	`132`	`ChainSyncMessage::InitOk { .. } => (),`
`133`	`133`	`ChainSyncMessage::Fwd {`
`@@ -142,7 +142,7 @@ pub fn spawn_node(`
`142`	`142`	`span: Span::current(),`
`143`	`143`	`},`
`144`	`144`	`)`
`145`		`- .await`
	`145`	`+ .await`
`146`	`146`	`}`
`147`	`147`	`ChainSyncMessage::Bck { slot, hash, .. } => {`
`148`	`148`	`eff.send(`
`@@ -153,7 +153,7 @@ pub fn spawn_node(`
`153`	`153`	`span: Span::current(),`
`154`	`154`	`},`
`155`	`155`	`)`
`156`		`- .await`
	`156`	`+ .await`
`157`	`157`	`}`
`158`	`158`	`}`
`159`	`159`	`(downstream, output)`
`@@ -238,9 +238,9 @@ fn make_chain_selector(`
`238`	`238`	`}`
`239`	`239`
`240`	`240`	`/// Property: at the end of the simulation, the chain built from the history of messages received`
`241`		`-/// must match one of the best chains that could be built from the entries.`
	`241`	`+/// downstream must match the best chain built directly from messages coming from upstream peers.`
`242`	`242`	`fn chain_property()`
`243`		`--> impl Fn(&[Entry<ChainSyncMessage>], &History<ChainSyncMessage>) -> Result<(), String> {`
	`243`	`+ -> impl Fn(&[Entry<ChainSyncMessage>], &History<ChainSyncMessage>) -> Result<(), String> {`
`244`	`244`	`move \|entries, history\| {`
`245`	`245`	`let expected = make_best_chain_from_entries(`
`246`	`246`	`&entries`
`@@ -289,7 +289,7 @@ pub fn make_best_chain_from_entries(messages: &[Envelope<ChainSyncMessage>]) ->`
`289`	`289`	`// make sure to skip header hashes that cannot be decoded`
`290`	`290`	`if let Some(header_hash) = msg.header_hash()`
`291`	`291`	`&& let Some(rollback_position) =`
`292`		`- chain.iter().position(\|h\| h.hash() == header_hash)`
	`292`	`+ chain.iter().position(\|h\| h.hash() == header_hash)`
`293`	`293`	`{`
`294`	`294`	`chain.truncate(rollback_position + 1)`
`295`	`295`	`}`