fix: make three-hop put test deterministic

Author: sanity

crates/core/tests/operations.rs

@@ -1909,9 +1909,31 @@ async fn test_delegate_request(ctx: &mut TestContext) -> TestResult {
 /// Expected flow:
 /// 1. peer-a sends PUT → routes through gateway → stored on peer-c
 /// 2. peer-c sends PUT response → routes back through gateway → received by peer-a
+const THREE_HOP_TEST_CONTRACT: &str = "test-contract-integration";
+
+static THREE_HOP_CONTRACT: LazyLock<(ContractContainer, freenet::dev_tool::Location)> =
+    LazyLock::new(|| {
+        let contract =
+            test_utils::load_contract(THREE_HOP_TEST_CONTRACT, vec![].into()).expect("contract");
+        let location = freenet::dev_tool::Location::from(&contract.key());
+        (contract, location)
+    });
+
+fn three_hop_node_locations() -> Vec<f64> {
+    let (_, contract_location) = &*THREE_HOP_CONTRACT;
+    let gateway = freenet::dev_tool::Location::new_rounded(contract_location.as_f64() + 0.2);
+    let peer_a = freenet::dev_tool::Location::new_rounded(contract_location.as_f64() + 0.5);
+    vec![
+        gateway.as_f64(),
+        peer_a.as_f64(),
+        contract_location.as_f64(),
+    ]
+}
+
 #[freenet_test(
     nodes = ["gateway", "peer-a", "peer-c"],
     gateways = ["gateway"],
+    node_locations_fn = three_hop_node_locations,
     auto_connect_peers = true,
     timeout_secs = 240,
     startup_wait_secs = 15,
@@ -1922,10 +1944,12 @@ async fn test_delegate_request(ctx: &mut TestContext) -> TestResult {
 async fn test_put_contract_three_hop_returns_response(ctx: &mut TestContext) -> TestResult {
     use freenet::dev_tool::Location;
 
-    const TEST_CONTRACT: &str = "test-contract-integration";
-    let contract = test_utils::load_contract(TEST_CONTRACT, vec![].into())?;
+    let (contract, contract_location) = {
+        let (contract, location) = &*THREE_HOP_CONTRACT;
+        (contract.clone(), *location)
+    };
     let contract_key = contract.key();
-    let contract_location = Location::from(&contract_key);
+    let node_locations = three_hop_node_locations();
 
     let initial_state = test_utils::create_empty_todo_list();
     let wrapped_state = WrappedState::from(initial_state);
@@ -1935,16 +1959,36 @@ async fn test_put_contract_three_hop_returns_response(ctx: &mut TestContext) ->
     let peer_a = ctx.node("peer-a")?;
     let peer_c = ctx.node("peer-c")?;
 
-    // Note: We cannot modify node locations after they're created with the macro,
-    // so this test will use random locations. The original test had specific location
-    // requirements to ensure proper three-hop routing. For now, we'll proceed with
-    // the test and it should still validate PUT response routing.
+    assert_eq!(gateway.location, node_locations[0]);
+    assert_eq!(peer_a.location, node_locations[1]);
+    assert_eq!(peer_c.location, node_locations[2]);
 
     tracing::info!("Node A data dir: {:?}", peer_a.temp_dir_path);
     tracing::info!("Gateway node data dir: {:?}", gateway.temp_dir_path);
     tracing::info!("Node C data dir: {:?}", peer_c.temp_dir_path);
     tracing::info!("Contract location: {}", contract_location.as_f64());
 
+    let gateway_distance = Location::new(gateway.location).distance(contract_location);
+    let peer_a_distance = Location::new(peer_a.location).distance(contract_location);
+    let peer_c_distance = Location::new(peer_c.location).distance(contract_location);
+
+    // Ensure the contract should naturally route to peer-c to create the 3-hop path:
+    // peer-a (client) -> gateway -> peer-c (closest to contract).
+    assert!(
+        peer_c_distance.as_f64() < gateway_distance.as_f64(),
+        "peer-c must be closer to contract than the gateway for three-hop routing"
+    );
+    assert!(
+        peer_c_distance.as_f64() < peer_a_distance.as_f64(),
+        "peer-c must be closest node to the contract location"
+    );
+    tracing::info!(
+        "Distances to contract - gateway: {}, peer-a: {}, peer-c: {}",
+        gateway_distance.as_f64(),
+        peer_a_distance.as_f64(),
+        peer_c_distance.as_f64()
+    );
+
     // Connect to peer A's WebSocket API
     let uri_a = format!(
         "ws://127.0.0.1:{}/v1/contract/command?encodingProtocol=native",

crates/freenet-macros/README.md

@@ -82,6 +82,45 @@ async fn test_multi_gateway(ctx: &mut TestContext) -> TestResult {
 }
 ```
 
+#### `node_locations`
+Set explicit ring locations for nodes (order must match the `nodes` list). This is useful for deterministic routing topologies in multi-hop tests.
+
+```rust
+#[freenet_test(
+    nodes = ["gateway", "peer-a", "peer-b"],
+    node_locations = [0.2, 0.6, 0.85], // gateway, peer-a, peer-b
+    auto_connect_peers = true
+)]
+async fn test_with_locations(ctx: &mut TestContext) -> TestResult {
+    // Test logic here...
+    Ok(())
+}
+```
+
+**Rules:**
+- Provide one numeric value per node.
+- Values should be in the range `[0.0, 1.0]`.
+- Omit `node_locations` to use random locations (default behavior).
+
+#### `node_locations_fn`
+Provide a function that returns `Vec<f64>` (one per node) when locations need to be computed dynamically (for example, based on a contract’s ring location).
+
+```rust
+fn my_locations() -> Vec<f64> {
+    // Must return one entry per node in the same order.
+    vec![0.1, 0.6, 0.9]
+}
+
+#[freenet_test(
+    nodes = ["gateway", "peer-a", "peer-b"],
+    node_locations_fn = my_locations,
+    auto_connect_peers = true
+)]
+```
+
+- The function must return exactly as many values as there are nodes; otherwise the test fails early.
+- `node_locations` and `node_locations_fn` are mutually exclusive.
+
 #### `auto_connect_peers`
 Automatically configure all peer nodes to connect to all gateway nodes.
 

crates/freenet-macros/src/codegen.rs

@@ -5,6 +5,69 @@ use proc_macro2::TokenStream;
 use quote::{format_ident, quote};
 use syn::{ItemFn, LitInt, Result};
 
+/// Get the configured location for a node or fall back to randomness.
+fn node_location(idx: usize) -> TokenStream {
+    let idx_lit = syn::Index::from(idx);
+    quote! {{
+        if let Some(ref locs) = __node_locations {
+            locs[#idx_lit]
+        } else {
+            rand::Rng::random(&mut rand::rng())
+        }
+    }}
+}
+
+/// Generate node location initialization (literal list or function).
+fn generate_node_locations_init(args: &FreenetTestArgs) -> TokenStream {
+    let node_count = args.nodes.len();
+
+    if let Some(ref fn_path) = args.node_locations_fn {
+        quote! {
+            let __node_locations: Option<Vec<f64>> = {
+                let locs = #fn_path();
+                if locs.len() != #node_count {
+                    return Err(anyhow::anyhow!(
+                        "node_locations_fn returned {} locations, expected {}",
+                        locs.len(),
+                        #node_count
+                    ));
+                }
+                for (idx, loc) in locs.iter().enumerate() {
+                    if !(0.0..=1.0).contains(loc) {
+                        return Err(anyhow::anyhow!(
+                            "node_locations_fn value at index {} is out of range: {} (must be in [0.0, 1.0])",
+                            idx,
+                            loc
+                        ));
+                    }
+                }
+                Some(locs)
+            };
+        }
+    } else if let Some(ref locations) = args.node_locations {
+        let values: Vec<_> = locations.iter().map(|loc| quote! { #loc }).collect();
+        quote! {
+            let __node_locations: Option<Vec<f64>> = {
+                let locs = vec![#(#values),*];
+                for (idx, loc) in locs.iter().enumerate() {
+                    if !(0.0..=1.0).contains(loc) {
+                        return Err(anyhow::anyhow!(
+                            "node_locations value at index {} is out of range: {} (must be in [0.0, 1.0])",
+                            idx,
+                            loc
+                        ));
+                    }
+                }
+                Some(locs)
+            };
+        }
+    } else {
+        quote! {
+            let __node_locations: Option<Vec<f64>> = None;
+        }
+    }
+}
+
 /// Helper to determine if a node is a gateway
 fn is_gateway(args: &FreenetTestArgs, node_label: &str, node_idx: usize) -> bool {
     if let Some(ref gateways) = args.gateways {
@@ -27,6 +90,7 @@ pub fn generate_test_code(args: FreenetTestArgs, input_fn: ItemFn) -> Result<Tok
 
     // Generate node setup code
     let node_setup = generate_node_setup(&args);
+    let node_locations_init = generate_node_locations_init(&args);
 
     // Extract values before configs are moved
     let value_extraction = generate_value_extraction(&args);
@@ -71,6 +135,7 @@ pub fn generate_test_code(args: FreenetTestArgs, input_fn: ItemFn) -> Result<Tok
             tracing::info!("Starting test: {}", stringify!(#test_fn_name));
 
             // 2. Create node configurations
+            #node_locations_init
             #node_setup
 
             // 3. Extract values before configs are moved
@@ -117,6 +182,7 @@ fn generate_node_setup(args: &FreenetTestArgs) -> TokenStream {
 
         if is_gw {
             // Gateway node configuration
+            let location_expr = node_location(idx);
             setup_code.push(quote! {
                 let (#config_var, #temp_var) = {
                     let temp_dir = tempfile::tempdir()?;
@@ -128,7 +194,7 @@ fn generate_node_setup(args: &FreenetTestArgs) -> TokenStream {
                     let network_port = freenet::test_utils::reserve_local_port()?;
                     let ws_port = freenet::test_utils::reserve_local_port()?;
 
-                    let location: f64 = rand::Rng::random(&mut rand::rng());
+                    let location: f64 = #location_expr;
 
                     let config = freenet::config::ConfigArgs {
                         ws_api: freenet::config::WebsocketApiArgs {
@@ -193,6 +259,7 @@ fn generate_node_setup(args: &FreenetTestArgs) -> TokenStream {
         let is_gw = is_gateway(args, node_label, idx);
 
         if !is_gw {
+            let location_expr = node_location(idx);
             // Collect gateway info variables to serialize
             let gateways_config = if args.auto_connect_peers {
                 // Collect all gateway_info_X variables
@@ -238,7 +305,7 @@ fn generate_node_setup(args: &FreenetTestArgs) -> TokenStream {
                     let network_port = freenet::test_utils::reserve_local_port()?;
                     let ws_port = freenet::test_utils::reserve_local_port()?;
 
-                    let location: f64 = rand::Rng::random(&mut rand::rng());
+                    let location: f64 = #location_expr;
 
                     let config = freenet::config::ConfigArgs {
                         ws_api: freenet::config::WebsocketApiArgs {

crates/freenet-macros/src/parser.rs

@@ -7,6 +7,10 @@ pub struct FreenetTestArgs {
     pub nodes: Vec<String>,
     /// Which nodes are gateways (if not specified, first node is gateway)
     pub gateways: Option<Vec<String>>,
+    /// Optional explicit node locations (same order as nodes)
+    pub node_locations: Option<Vec<f64>>,
+    /// Optional function path that returns node locations (same order as nodes)
+    pub node_locations_fn: Option<syn::ExprPath>,
     /// Whether peers should auto-connect to gateways
     pub auto_connect_peers: bool,
     /// Test timeout in seconds
@@ -40,6 +44,8 @@ impl syn::parse::Parse for FreenetTestArgs {
     fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
         let mut nodes = None;
         let mut gateways = None;
+        let mut node_locations = None;
+        let mut node_locations_fn = None;
         let mut auto_connect_peers = false;
         let mut timeout_secs = 180;
         let mut startup_wait_secs = 15;
@@ -79,6 +85,47 @@ impl syn::parse::Parse for FreenetTestArgs {
 
                     nodes = Some(node_list);
                 }
+                "node_locations" => {
+                    // Parse array literal of floats: [0.1, 0.5, 0.9]
+                    let content;
+                    syn::bracketed!(content in input);
+
+                    let mut locations = Vec::new();
+                    while !content.is_empty() {
+                        let lit: syn::Lit = content.parse()?;
+                        let value = match lit {
+                            syn::Lit::Float(f) => f.base10_parse::<f64>()?,
+                            syn::Lit::Int(i) => i.base10_parse::<f64>()?,
+                            other => {
+                                return Err(syn::Error::new(
+                                    other.span(),
+                                    "node_locations must be numeric literals",
+                                ))
+                            }
+                        };
+
+                        locations.push(value);
+
+                        // Handle optional trailing comma
+                        if content.peek(syn::Token![,]) {
+                            content.parse::<syn::Token![,]>()?;
+                        }
+                    }
+
+                    if locations.is_empty() {
+                        return Err(syn::Error::new(
+                            key.span(),
+                            "node_locations array cannot be empty if specified",
+                        ));
+                    }
+
+                    node_locations = Some(locations);
+                }
+                "node_locations_fn" => {
+                    // Parse a path to a function returning Vec<f64>
+                    let path: syn::ExprPath = input.parse()?;
+                    node_locations_fn = Some(path);
+                }
                 "gateways" => {
                     // Parse array literal: ["gateway-1", "gateway-2", ...]
                     let content;
@@ -191,9 +238,37 @@ impl syn::parse::Parse for FreenetTestArgs {
             }
         }
 
+        // Validate node_locations if provided
+        if let Some(ref locations) = node_locations {
+            if locations.len() != nodes.len() {
+                return Err(input.error(format!(
+                    "node_locations length ({}) must match nodes length ({})",
+                    locations.len(),
+                    nodes.len()
+                )));
+            }
+            for (i, &loc) in locations.iter().enumerate() {
+                if !(0.0..=1.0).contains(&loc) {
+                    return Err(input.error(format!(
+                        "node_locations[{}] = {} is out of range. Values must be in [0.0, 1.0]",
+                        i, loc
+                    )));
+                }
+            }
+        }
+
+        // Only one of node_locations or node_locations_fn may be provided
+        if node_locations.is_some() && node_locations_fn.is_some() {
+            return Err(
+                input.error("Specify only one of node_locations or node_locations_fn (not both)")
+            );
+        }
+
         Ok(FreenetTestArgs {
             nodes,
             gateways,
+            node_locations,
+            node_locations_fn,
             auto_connect_peers,
             timeout_secs,
             startup_wait_secs,