use adk_graph::deferred::{FanInTracker, MergeStrategy};
use proptest::prelude::*;
use serde_json::{Value, json};
fn arb_num_paths() -> impl Strategy<Value = usize> {
2usize..=5usize
}
fn arb_output_value() -> impl Strategy<Value = Value> {
prop_oneof![
any::<i64>().prop_map(|n| json!(n)),
"[a-z]{1,10}".prop_map(|s| json!(s)),
any::<bool>().prop_map(|b| json!(b)),
(any::<i64>(), "[a-z]{1,5}").prop_map(|(n, k)| json!({ k: n })),
]
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(100))]
#[test]
fn prop_fan_in_not_ready_until_all_complete(
num_paths in arb_num_paths(),
outputs in proptest::collection::vec(arb_output_value(), 5),
) {
let sources: Vec<String> = (0..num_paths).map(|i| format!("source_{i}")).collect();
let source_refs: Vec<&str> = sources.iter().map(|s| s.as_str()).collect();
let mut tracker = FanInTracker::new(source_refs);
for i in 0..num_paths {
prop_assert!(
!tracker.is_ready(),
"tracker should NOT be ready after {}/{} sources recorded",
i,
num_paths
);
prop_assert_eq!(
tracker.received_count(),
i,
"received_count should be {} before recording source {}",
i,
i
);
tracker.record(&sources[i], outputs[i % outputs.len()].clone());
}
prop_assert!(
tracker.is_ready(),
"tracker should be ready after all {} sources recorded",
num_paths
);
prop_assert_eq!(tracker.received_count(), num_paths);
}
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(100))]
#[test]
fn prop_merge_collect_produces_ordered_vec(
num_paths in arb_num_paths(),
outputs in proptest::collection::vec(arb_output_value(), 2..=5),
) {
let sources: Vec<String> = (0..num_paths).map(|i| format!("src_{i}")).collect();
let source_refs: Vec<&str> = sources.iter().map(|s| s.as_str()).collect();
let mut tracker = FanInTracker::new(source_refs);
let mut expected_outputs = Vec::new();
for i in 0..num_paths {
let output = outputs[i % outputs.len()].clone();
expected_outputs.push(output.clone());
tracker.record(&sources[i], output);
}
prop_assert!(tracker.is_ready());
let merged = tracker.merge(&MergeStrategy::Collect);
let arr = merged.as_array().expect("Collect should produce an array");
prop_assert_eq!(
arr.len(),
num_paths,
"Collect should produce {} elements, got {}",
num_paths,
arr.len()
);
for (i, expected) in expected_outputs.iter().enumerate() {
prop_assert_eq!(
&arr[i],
expected,
"Element at index {} should match recorded output",
i
);
}
}
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(100))]
#[test]
fn prop_merge_map_combines_all_keys(
num_paths in arb_num_paths(),
values in proptest::collection::vec(any::<i64>(), 5),
) {
let sources: Vec<String> = (0..num_paths).map(|i| format!("node_{i}")).collect();
let source_refs: Vec<&str> = sources.iter().map(|s| s.as_str()).collect();
let mut tracker = FanInTracker::new(source_refs);
for i in 0..num_paths {
let key = format!("key_{i}");
let output = json!({ key: values[i % values.len()] });
tracker.record(&sources[i], output);
}
prop_assert!(tracker.is_ready());
let merged = tracker.merge(&MergeStrategy::MergeMap);
let obj = merged.as_object().expect("MergeMap should produce an object");
for i in 0..num_paths {
let key = format!("key_{i}");
prop_assert!(
obj.contains_key(&key),
"merged object should contain key '{}'",
key
);
let expected_val = json!(values[i % values.len()]);
prop_assert_eq!(
obj.get(&key).unwrap(),
&expected_val,
"value for key '{}' should match",
key
);
}
}
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(100))]
#[test]
fn prop_merge_first_returns_first_recorded(
num_paths in arb_num_paths(),
outputs in proptest::collection::vec(arb_output_value(), 2..=5),
) {
let sources: Vec<String> = (0..num_paths).map(|i| format!("path_{i}")).collect();
let source_refs: Vec<&str> = sources.iter().map(|s| s.as_str()).collect();
let mut tracker = FanInTracker::new(source_refs);
let first_output = outputs[0].clone();
for i in 0..num_paths {
tracker.record(&sources[i], outputs[i % outputs.len()].clone());
}
prop_assert!(tracker.is_ready());
let merged = tracker.merge(&MergeStrategy::First);
prop_assert_eq!(
&merged,
&first_output,
"First strategy should return the first recorded output"
);
}
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(100))]
#[test]
fn prop_expected_count_matches_construction(
num_paths in arb_num_paths(),
) {
let sources: Vec<String> = (0..num_paths).map(|i| format!("s_{i}")).collect();
let source_refs: Vec<&str> = sources.iter().map(|s| s.as_str()).collect();
let tracker = FanInTracker::new(source_refs);
prop_assert_eq!(
tracker.expected_count(),
num_paths,
"expected_count should equal the number of sources"
);
}
}