Added some more test

Author: milindsrivastava1997

asap-planner-rs/Cargo.toml

@@ -22,7 +22,7 @@ anyhow.workspace = true
 tracing.workspace = true
 tracing-subscriber.workspace = true
 clap.workspace = true
-indexmap = { version = "2.0", features = ["serde"] }
+indexmap.workspace = true
 promql-parser = "0.5.0"
 
 [dev-dependencies]

asap-planner-rs/src/lib.rs

@@ -92,6 +92,69 @@ impl PlannerOutput {
         false
     }
 
+    /// Returns the sorted labels for the first aggregation matching `agg_type`,
+    /// for the given `label_kind` ("rollup", "grouping", or "aggregated").
+    pub fn aggregation_labels(&self, agg_type: &str, label_kind: &str) -> Vec<String> {
+        if let YamlValue::Mapping(root) = &self.streaming_yaml {
+            if let Some(YamlValue::Sequence(aggs)) = root.get("aggregations") {
+                for agg in aggs {
+                    if let YamlValue::Mapping(m) = agg {
+                        if let Some(YamlValue::String(t)) = m.get("aggregationType") {
+                            if t == agg_type {
+                                if let Some(YamlValue::Mapping(labels)) = m.get("labels") {
+                                    if let Some(YamlValue::Sequence(seq)) = labels.get(label_kind) {
+                                        let mut result: Vec<String> = seq
+                                            .iter()
+                                            .filter_map(|v| {
+                                                if let YamlValue::String(s) = v {
+                                                    Some(s.clone())
+                                                } else {
+                                                    None
+                                                }
+                                            })
+                                            .collect();
+                                        result.sort();
+                                        return result;
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        }
+        vec![]
+    }
+
+    /// Returns the cleanup param (read_count_threshold or num_aggregates_to_retain)
+    /// for the first aggregation entry of the given query string.
+    pub fn inference_cleanup_param(&self, query: &str) -> Option<u64> {
+        if let YamlValue::Mapping(root) = &self.inference_yaml {
+            if let Some(YamlValue::Sequence(queries)) = root.get("queries") {
+                for q in queries {
+                    if let YamlValue::Mapping(qm) = q {
+                        if let Some(YamlValue::String(qs)) = qm.get("query") {
+                            if qs == query {
+                                if let Some(YamlValue::Sequence(aggs)) = qm.get("aggregations") {
+                                    if let Some(YamlValue::Mapping(agg)) = aggs.first() {
+                                        for key in ["read_count_threshold", "num_aggregates_to_retain"] {
+                                            if let Some(v) = agg.get(key) {
+                                                if let YamlValue::Number(n) = v {
+                                                    return n.as_u64();
+                                                }
+                                            }
+                                        }
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        }
+        None
+    }
+
     pub fn to_streaming_yaml_string(&self) -> Result<String, anyhow::Error> {
         Ok(serde_yaml::to_string(&self.streaming_yaml)?)
     }

asap-planner-rs/src/output/generator.rs

@@ -74,15 +74,13 @@ pub fn generate_plan(
 
             if should_process {
                 let (configs, cleanup_param) = processor.get_streaming_aggregation_configs()?;
-                query_keys_map.insert(query_string.clone(), Vec::new());
+                let mut keys_for_query = Vec::new();
                 for config in configs {
                     let key = config.identifying_key();
-                    query_keys_map
-                        .get_mut(query_string)
-                        .unwrap()
-                        .push((key.clone(), cleanup_param));
+                    keys_for_query.push((key.clone(), cleanup_param));
                     dedup_map.entry(key).or_insert(config);
                 }
+                query_keys_map.insert(query_string.clone(), keys_for_query);
             }
         }
     }

asap-planner-rs/src/planner/logics.rs

@@ -250,6 +250,117 @@ pub fn get_cleanup_param(
     Ok(result)
 }
 
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::planner::patterns::build_patterns;
+    use promql_utilities::ast_matching::PromQLMatchResult;
+    use promql_utilities::query_logics::enums::QueryPatternType;
+
+    fn match_query(query: &str) -> (QueryPatternType, PromQLMatchResult) {
+        let ast = promql_parser::parser::parse(query).unwrap();
+        let patterns = build_patterns();
+        for (pt, pattern) in &patterns {
+            let result = pattern.matches(&ast);
+            if result.matches {
+                return (*pt, result);
+            }
+        }
+        panic!("no pattern matched query: {}", query);
+    }
+
+    // --- get_effective_repeat ---
+
+    #[test]
+    fn effective_repeat_no_step() {
+        assert_eq!(get_effective_repeat(300, 0), 300);
+    }
+
+    #[test]
+    fn effective_repeat_step_smaller_than_t_repeat() {
+        assert_eq!(get_effective_repeat(300, 30), 30);
+    }
+
+    #[test]
+    fn effective_repeat_step_larger_than_t_repeat() {
+        assert_eq!(get_effective_repeat(30, 300), 30);
+    }
+
+    // --- get_cleanup_param ---
+
+    #[test]
+    fn cleanup_param_circular_buffer_spatial_instant_query() {
+        let (pt, mr) = match_query("sum(some_metric)");
+        assert_eq!(pt, QueryPatternType::OnlySpatial);
+        // t_lookback = t_repeat = 300 (OnlySpatial path)
+        // effective_repeat = 300 (step=0)
+        // ceil((300 + 0) / 300) = 1
+        let result =
+            get_cleanup_param(CleanupPolicy::CircularBuffer, pt, &mr, 300, "tumbling", 0, 0)
+                .unwrap();
+        assert_eq!(result, 1);
+    }
+
+    #[test]
+    fn cleanup_param_circular_buffer_spatial_range_query() {
+        let (pt, mr) = match_query("sum(some_metric)");
+        // t_lookback = t_repeat = 300, effective_repeat = min(300, 30) = 30
+        // ceil((300 + 3600) / 30) = ceil(130) = 130
+        let result =
+            get_cleanup_param(CleanupPolicy::CircularBuffer, pt, &mr, 300, "tumbling", 3600, 30)
+                .unwrap();
+        assert_eq!(result, 130);
+    }
+
+    #[test]
+    fn cleanup_param_read_based_spatial_instant_query() {
+        let (pt, mr) = match_query("sum(some_metric)");
+        // lookback_buckets = ceil(300/300) = 1, num_steps = 1 → result = 1
+        let result =
+            get_cleanup_param(CleanupPolicy::ReadBased, pt, &mr, 300, "tumbling", 0, 0).unwrap();
+        assert_eq!(result, 1);
+    }
+
+    #[test]
+    fn cleanup_param_read_based_spatial_range_query() {
+        let (pt, mr) = match_query("sum(some_metric)");
+        // lookback_buckets = ceil(300/30) = 10, num_steps = 3600/30 + 1 = 121
+        // result = 10 * 121 = 1210
+        let result =
+            get_cleanup_param(CleanupPolicy::ReadBased, pt, &mr, 300, "tumbling", 3600, 30)
+                .unwrap();
+        assert_eq!(result, 1210);
+    }
+
+    #[test]
+    fn cleanup_param_circular_buffer_temporal_instant_query() {
+        let (pt, mr) = match_query("rate(some_metric[5m])");
+        assert_eq!(pt, QueryPatternType::OnlyTemporal);
+        // t_lookback = 5m = 300s (from [5m] range vector), range_duration=0, step=0
+        // effective_repeat = 60, ceil((300 + 0) / 60) = 5
+        let result =
+            get_cleanup_param(CleanupPolicy::CircularBuffer, pt, &mr, 60, "tumbling", 0, 0)
+                .unwrap();
+        assert_eq!(result, 5);
+    }
+
+    #[test]
+    fn cleanup_param_no_cleanup_returns_error() {
+        let (pt, mr) = match_query("sum(some_metric)");
+        let result = get_cleanup_param(CleanupPolicy::NoCleanup, pt, &mr, 300, "tumbling", 0, 0);
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn cleanup_param_mismatched_range_and_step_returns_error() {
+        let (pt, mr) = match_query("sum(some_metric)");
+        // range_duration > 0 but step == 0 is invalid
+        let result =
+            get_cleanup_param(CleanupPolicy::CircularBuffer, pt, &mr, 300, "tumbling", 3600, 0);
+        assert!(result.is_err());
+    }
+}
+
 pub fn set_subpopulation_labels(
     statistic: Statistic,
     aggregation_type: &str,

asap-planner-rs/src/planner/single_query.rs

@@ -413,3 +413,75 @@ fn compute_labels(
 
     (rollup, grouping, aggregated)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::Value;
+    use std::collections::HashMap;
+
+    fn base_config() -> IntermediateAggConfig {
+        IntermediateAggConfig {
+            aggregation_type: "MultipleIncrease".to_string(),
+            aggregation_sub_type: "rate".to_string(),
+            window_type: "tumbling".to_string(),
+            window_size: 300,
+            slide_interval: 300,
+            tumbling_window_size: 300,
+            spatial_filter: String::new(),
+            metric: "http_requests_total".to_string(),
+            table_name: None,
+            value_column: None,
+            parameters: HashMap::new(),
+            rollup_labels: KeyByLabelNames::new(vec!["instance".to_string()]),
+            grouping_labels: KeyByLabelNames::empty(),
+            aggregated_labels: KeyByLabelNames::empty(),
+        }
+    }
+
+    #[test]
+    fn identifying_key_is_stable() {
+        let cfg = base_config();
+        assert_eq!(cfg.identifying_key(), cfg.identifying_key());
+    }
+
+    #[test]
+    fn identical_configs_have_same_key() {
+        assert_eq!(base_config().identifying_key(), base_config().identifying_key());
+    }
+
+    #[test]
+    fn different_aggregation_type_produces_different_key() {
+        let cfg1 = base_config();
+        let mut cfg2 = base_config();
+        cfg2.aggregation_type = "DatasketchesKLL".to_string();
+        assert_ne!(cfg1.identifying_key(), cfg2.identifying_key());
+    }
+
+    #[test]
+    fn different_window_size_produces_different_key() {
+        let cfg1 = base_config();
+        let mut cfg2 = base_config();
+        cfg2.window_size = 60;
+        assert_ne!(cfg1.identifying_key(), cfg2.identifying_key());
+    }
+
+    #[test]
+    fn different_rollup_labels_produce_different_key() {
+        let cfg1 = base_config();
+        let mut cfg2 = base_config();
+        cfg2.rollup_labels = KeyByLabelNames::new(vec!["job".to_string()]);
+        assert_ne!(cfg1.identifying_key(), cfg2.identifying_key());
+    }
+
+    #[test]
+    fn parameter_insertion_order_does_not_affect_key() {
+        let mut cfg1 = base_config();
+        let mut cfg2 = base_config();
+        cfg1.parameters.insert("depth".to_string(), Value::Number(3.into()));
+        cfg1.parameters.insert("width".to_string(), Value::Number(1024.into()));
+        cfg2.parameters.insert("width".to_string(), Value::Number(1024.into()));
+        cfg2.parameters.insert("depth".to_string(), Value::Number(3.into()));
+        assert_eq!(cfg1.identifying_key(), cfg2.identifying_key());
+    }
+}

asap-planner-rs/tests/comparison/test_cases.json

@@ -80,6 +80,51 @@
       "enable_punting": false,
       "range_duration": 0,
       "step": 0
+    },
+    {
+      "id": "increase",
+      "input_config": "test_data/configs/increase.yaml",
+      "prometheus_scrape_interval": 15,
+      "streaming_engine": "arroyo",
+      "enable_punting": false,
+      "range_duration": 0,
+      "step": 0
+    },
+    {
+      "id": "sum_over_time",
+      "input_config": "test_data/configs/sum_over_time.yaml",
+      "prometheus_scrape_interval": 15,
+      "streaming_engine": "arroyo",
+      "enable_punting": false,
+      "range_duration": 0,
+      "step": 0
+    },
+    {
+      "id": "sum_by",
+      "input_config": "test_data/configs/sum_by.yaml",
+      "prometheus_scrape_interval": 15,
+      "streaming_engine": "arroyo",
+      "enable_punting": false,
+      "range_duration": 0,
+      "step": 0
+    },
+    {
+      "id": "temporal_overlapping",
+      "input_config": "test_data/configs/temporal_overlapping.yaml",
+      "prometheus_scrape_interval": 15,
+      "streaming_engine": "arroyo",
+      "enable_punting": false,
+      "range_duration": 0,
+      "step": 0
+    },
+    {
+      "id": "sum_by_overlapping",
+      "input_config": "test_data/configs/sum_by_overlapping.yaml",
+      "prometheus_scrape_interval": 15,
+      "streaming_engine": "arroyo",
+      "enable_punting": false,
+      "range_duration": 0,
+      "step": 0
     }
   ]
 }

asap-planner-rs/tests/comparison/test_data/configs/increase.yaml

@@ -0,0 +1,13 @@
+query_groups:
+  - id: 1
+    queries:
+      - "increase(http_requests_total[5m])"
+    repetition_delay: 300
+    controller_options:
+      accuracy_sla: 0.99
+      latency_sla: 1.0
+metrics:
+  - metric: "http_requests_total"
+    labels: ["instance", "job", "method", "status"]
+aggregate_cleanup:
+  policy: "read_based"

asap-planner-rs/tests/comparison/test_data/configs/sum_by.yaml

@@ -0,0 +1,13 @@
+query_groups:
+  - id: 1
+    queries:
+      - "sum by (job, method) (http_requests_total)"
+    repetition_delay: 300
+    controller_options:
+      accuracy_sla: 0.99
+      latency_sla: 1.0
+metrics:
+  - metric: "http_requests_total"
+    labels: ["instance", "job", "method", "status"]
+aggregate_cleanup:
+  policy: "read_based"

asap-planner-rs/tests/comparison/test_data/configs/sum_by_overlapping.yaml

@@ -0,0 +1,13 @@
+query_groups:
+  - id: 1
+    queries:
+      - "sum by (job, method) (http_requests_total)"
+    repetition_delay: 60
+    controller_options:
+      accuracy_sla: 0.99
+      latency_sla: 1.0
+metrics:
+  - metric: "http_requests_total"
+    labels: ["instance", "job", "method", "status"]
+aggregate_cleanup:
+  policy: "read_based"

asap-planner-rs/tests/comparison/test_data/configs/sum_over_time.yaml

@@ -0,0 +1,13 @@
+query_groups:
+  - id: 1
+    queries:
+      - "sum_over_time(http_requests_total[5m])"
+    repetition_delay: 300
+    controller_options:
+      accuracy_sla: 0.99
+      latency_sla: 1.0
+metrics:
+  - metric: "http_requests_total"
+    labels: ["instance", "job", "method", "status"]
+aggregate_cleanup:
+  policy: "read_based"