Merge pull request #36 from BTreeMap/copilot/optimize-algorithms-and-systems

Fix flaky analytics integration tests for incomplete GeoIP data

Author: MinecraftFuns

tests/analytics_integration.rs

@@ -202,10 +202,18 @@ async fn test_storage_integration() {
     println!("Country aggregates: {:?}", agg_country);
     assert!(!agg_country.is_empty());
 
-    // Verify totals match
+    // Verify totals match (only counting entries with non-NULL country_code,
+    // since the aggregate query filters out NULL country codes)
     let total_agg: i64 = agg_country.iter().map(|a| a.visit_count).sum();
-    let total_detail: i64 = analytics.iter().map(|a| a.visit_count).sum();
-    assert_eq!(total_agg, total_detail);
+    let total_detail_with_country: i64 = analytics
+        .iter()
+        .filter(|a| a.country_code.is_some())
+        .map(|a| a.visit_count)
+        .sum();
+    assert_eq!(
+        total_agg, total_detail_with_country,
+        "Aggregate total should match detail total for entries with country codes"
+    );
 }
 
 #[tokio::test]
@@ -306,20 +314,41 @@ async fn test_aggregation_dimensions() {
         ("9.9.9.9", 3),        // Quad9 US, AS19281 (different ASN)
         ("208.67.222.222", 2), // OpenDNS US, AS36692 (different ASN)
     ];
-    for (ip_str, count) in tests {
-        for _ in 0..count {
-            let ip: IpAddr = ip_str.parse().unwrap();
-            let geo = geoip.lookup(ip);
+
+    // Track expected counts per dimension based on actual GeoIP data availability
+    let mut expected_country_count: i64 = 0;
+    let mut expected_asn_count: i64 = 0;
+    let total_visits: i64 = 10; // 5 + 3 + 2
+
+    for (ip_str, count) in &tests {
+        let ip: IpAddr = ip_str.parse().unwrap();
+        let geo = geoip.lookup(ip);
+
+        // Track which dimension data is available
+        if geo.country_code.is_some() {
+            expected_country_count += *count as i64;
+        }
+        if geo.asn.is_some() {
+            expected_asn_count += *count as i64;
+        }
+
+        // Reuse the same geo lookup for all iterations of this IP
+        for _ in 0..*count {
             let rec = lynx::analytics::AnalyticsRecord {
                 short_code: "multi".to_string(),
                 timestamp: chrono::Utc::now().timestamp(),
-                geo_location: geo,
+                geo_location: geo.clone(),
                 client_ip: Some(ip),
             };
             agg.record(rec);
         }
     }
 
+    println!(
+        "Expected counts - country: {}, asn: {}, hour: {}",
+        expected_country_count, expected_asn_count, total_visits
+    );
+
     // Flush to storage
     let entries = agg.drain();
     let records: Vec<_> = entries
@@ -339,17 +368,31 @@ async fn test_aggregation_dimensions() {
         .collect();
     storage.upsert_analytics_batch(records).await.unwrap();
 
-    // Test different dimensions
-    for dim in &["country", "asn", "hour"] {
-        let agg = storage
+    // Test different dimensions with appropriate expected values
+    // The aggregate query filters out entries with NULL values for the dimension
+    let dimensions = vec![
+        ("country", expected_country_count),
+        ("asn", expected_asn_count),
+        ("hour", total_visits), // time_bucket is never NULL
+    ];
+
+    for (dim, expected) in dimensions {
+        let agg_result = storage
             .get_analytics_aggregate("multi", None, None, dim, 10)
             .await
             .unwrap();
-        println!("Aggregated by {}: {:?}", dim, agg);
-        assert!(!agg.is_empty());
-
-        let total: i64 = agg.iter().map(|a| a.visit_count).sum();
-        assert_eq!(total, 10, "Total should be 10 for {}", dim);
+        println!("Aggregated by {}: {:?}", dim, agg_result);
+
+        // For dimensions with expected data, verify the aggregate is non-empty and totals match
+        if expected > 0 {
+            assert!(!agg_result.is_empty(), "Should have aggregates for {}", dim);
+            let total: i64 = agg_result.iter().map(|a| a.visit_count).sum();
+            assert_eq!(
+                total, expected,
+                "Total for {} should match expected (got {}, expected {})",
+                dim, total, expected
+            );
+        }
     }
 }