[router] add ut for pd request, metrics and config (#8184)

sgl-router/src/metrics.rs

@@ -322,3 +322,414 @@ impl RouterMetrics {
         .set(count as f64);
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::net::TcpListener;
+
+    // ============= PrometheusConfig Tests =============
+
+    #[test]
+    fn test_prometheus_config_default() {
+        let config = PrometheusConfig::default();
+        assert_eq!(config.port, 29000);
+        assert_eq!(config.host, "0.0.0.0");
+    }
+
+    #[test]
+    fn test_prometheus_config_custom() {
+        let config = PrometheusConfig {
+            port: 8080,
+            host: "127.0.0.1".to_string(),
+        };
+        assert_eq!(config.port, 8080);
+        assert_eq!(config.host, "127.0.0.1");
+    }
+
+    #[test]
+    fn test_prometheus_config_clone() {
+        let config = PrometheusConfig {
+            port: 9090,
+            host: "192.168.1.1".to_string(),
+        };
+        let cloned = config.clone();
+        assert_eq!(cloned.port, config.port);
+        assert_eq!(cloned.host, config.host);
+    }
+
+    // ============= IP Address Parsing Tests =============
+
+    #[test]
+    fn test_valid_ipv4_parsing() {
+        let test_cases = vec!["127.0.0.1", "192.168.1.1", "0.0.0.0"];
+
+        for ip_str in test_cases {
+            let config = PrometheusConfig {
+                port: 29000,
+                host: ip_str.to_string(),
+            };
+
+            let ip_addr: IpAddr = config.host.parse().unwrap();
+            assert!(matches!(ip_addr, IpAddr::V4(_)));
+        }
+    }
+
+    #[test]
+    fn test_valid_ipv6_parsing() {
+        let test_cases = vec!["::1", "2001:db8::1", "::"];
+
+        for ip_str in test_cases {
+            let config = PrometheusConfig {
+                port: 29000,
+                host: ip_str.to_string(),
+            };
+
+            let ip_addr: IpAddr = config.host.parse().unwrap();
+            assert!(matches!(ip_addr, IpAddr::V6(_)));
+        }
+    }
+
+    #[test]
+    fn test_invalid_ip_parsing() {
+        let test_cases = vec!["invalid", "256.256.256.256", "hostname"];
+
+        for ip_str in test_cases {
+            let config = PrometheusConfig {
+                port: 29000,
+                host: ip_str.to_string(),
+            };
+
+            let ip_addr: IpAddr = config
+                .host
+                .parse()
+                .unwrap_or(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)));
+
+            // Should fall back to 0.0.0.0
+            assert_eq!(ip_addr, IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)));
+        }
+    }
+
+    // ============= Socket Address Creation Tests =============
+
+    #[test]
+    fn test_socket_addr_creation() {
+        let test_cases = vec![("127.0.0.1", 8080), ("0.0.0.0", 29000), ("::1", 9090)];
+
+        for (host, port) in test_cases {
+            let config = PrometheusConfig {
+                port,
+                host: host.to_string(),
+            };
+
+            let ip_addr: IpAddr = config.host.parse().unwrap();
+            let socket_addr = SocketAddr::new(ip_addr, config.port);
+
+            assert_eq!(socket_addr.port(), port);
+            assert_eq!(socket_addr.ip().to_string(), host);
+        }
+    }
+
+    #[test]
+    fn test_socket_addr_with_different_ports() {
+        let ports = vec![0, 80, 8080, 65535];
+
+        for port in ports {
+            let config = PrometheusConfig {
+                port,
+                host: "127.0.0.1".to_string(),
+            };
+
+            let ip_addr: IpAddr = config.host.parse().unwrap();
+            let socket_addr = SocketAddr::new(ip_addr, config.port);
+
+            assert_eq!(socket_addr.port(), port);
+        }
+    }
+
+    // ============= Duration Bucket Tests =============
+
+    #[test]
+    fn test_duration_bucket_values() {
+        let expected_buckets = vec![
+            0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 15.0, 30.0, 45.0,
+            60.0, 90.0, 120.0, 180.0, 240.0,
+        ];
+
+        // The buckets are defined in start_prometheus function
+        assert_eq!(expected_buckets.len(), 20);
+
+        // Verify proper ordering
+        for i in 1..expected_buckets.len() {
+            assert!(expected_buckets[i] > expected_buckets[i - 1]);
+        }
+    }
+
+    #[test]
+    fn test_duration_bucket_coverage() {
+        let test_cases = vec![
+            (0.0005, "sub-millisecond"),
+            (0.005, "5ms"),
+            (0.05, "50ms"),
+            (1.0, "1s"),
+            (10.0, "10s"),
+            (60.0, "1m"),
+            (240.0, "4m"),
+        ];
+
+        let buckets = vec![
+            0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 15.0, 30.0, 45.0,
+            60.0, 90.0, 120.0, 180.0, 240.0,
+        ];
+
+        for (duration, label) in test_cases {
+            let bucket_found = buckets
+                .iter()
+                .any(|&b| ((b - duration) as f64).abs() < 0.0001 || b > duration);
+            assert!(bucket_found, "No bucket found for {} ({})", duration, label);
+        }
+    }
+
+    // ============= Matcher Configuration Tests =============
+
+    #[test]
+    fn test_duration_suffix_matcher() {
+        let matcher = Matcher::Suffix(String::from("duration_seconds"));
+
+        // Test matching behavior
+        let _matching_metrics = vec![
+            "request_duration_seconds",
+            "response_duration_seconds",
+            "sgl_router_request_duration_seconds",
+        ];
+
+        let _non_matching_metrics =
+            vec!["duration_total", "duration_seconds_total", "other_metric"];
+
+        // Note: We can't directly test Matcher matching without the internals,
+        // but we can verify the matcher is created correctly
+        match matcher {
+            Matcher::Suffix(suffix) => assert_eq!(suffix, "duration_seconds"),
+            _ => panic!("Expected Suffix matcher"),
+        }
+    }
+
+    // ============= Builder Configuration Tests =============
+
+    #[test]
+    fn test_prometheus_builder_configuration() {
+        // This test verifies the builder configuration without actually starting Prometheus
+        let _config = PrometheusConfig::default();
+
+        let duration_matcher = Matcher::Suffix(String::from("duration_seconds"));
+        let duration_bucket = [
+            0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 15.0, 30.0, 45.0,
+            60.0, 90.0, 120.0, 180.0, 240.0,
+        ];
+
+        // Verify bucket configuration
+        assert_eq!(duration_bucket.len(), 20);
+
+        // Verify matcher is suffix type
+        match duration_matcher {
+            Matcher::Suffix(s) => assert_eq!(s, "duration_seconds"),
+            _ => panic!("Expected Suffix matcher"),
+        }
+    }
+
+    // ============= Upkeep Timeout Tests =============
+
+    #[test]
+    fn test_upkeep_timeout_duration() {
+        let timeout = Duration::from_secs(5 * 60);
+        assert_eq!(timeout.as_secs(), 300);
+    }
+
+    // ============= Custom Bucket Tests =============
+
+    #[test]
+    fn test_custom_buckets_for_different_metrics() {
+        // Test that we can create different bucket configurations
+        let request_buckets = vec![0.001, 0.01, 0.1, 1.0, 10.0];
+        let generate_buckets = vec![0.1, 0.5, 1.0, 5.0, 30.0, 60.0];
+
+        assert_eq!(request_buckets.len(), 5);
+        assert_eq!(generate_buckets.len(), 6);
+
+        // Verify each set is sorted
+        for i in 1..request_buckets.len() {
+            assert!(request_buckets[i] > request_buckets[i - 1]);
+        }
+
+        for i in 1..generate_buckets.len() {
+            assert!(generate_buckets[i] > generate_buckets[i - 1]);
+        }
+    }
+
+    // ============= RouterMetrics Tests =============
+
+    #[test]
+    fn test_metrics_static_methods() {
+        // Test that all static methods can be called without panic
+        RouterMetrics::record_request("/generate");
+        RouterMetrics::record_request_duration("/generate", Duration::from_millis(100));
+        RouterMetrics::record_request_error("/generate", "timeout");
+        RouterMetrics::record_retry("/generate");
+
+        RouterMetrics::set_active_workers(5);
+        RouterMetrics::set_worker_health("http://worker1", true);
+        RouterMetrics::set_worker_load("http://worker1", 10);
+        RouterMetrics::record_processed_request("http://worker1");
+
+        RouterMetrics::record_policy_decision("random", "http://worker1");
+        RouterMetrics::record_cache_hit();
+        RouterMetrics::record_cache_miss();
+        RouterMetrics::set_tree_size("http://worker1", 1000);
+        RouterMetrics::record_load_balancing_event();
+        RouterMetrics::set_load_range(20, 5);
+
+        RouterMetrics::record_pd_request("/v1/chat/completions");
+        RouterMetrics::record_pd_request_duration("/v1/chat/completions", Duration::from_secs(1));
+        RouterMetrics::record_pd_prefill_request("http://prefill1");
+        RouterMetrics::record_pd_decode_request("http://decode1");
+        RouterMetrics::record_pd_error("invalid_request");
+        RouterMetrics::record_pd_prefill_error("http://prefill1");
+        RouterMetrics::record_pd_decode_error("http://decode1");
+        RouterMetrics::record_pd_stream_error("http://decode1");
+
+        RouterMetrics::record_discovery_update(3, 1);
+        RouterMetrics::record_generate_duration(Duration::from_secs(2));
+        RouterMetrics::set_running_requests("http://worker1", 15);
+    }
+
+    // ============= Port Availability Tests =============
+
+    #[test]
+    fn test_port_already_in_use() {
+        // Skip this test if we can't bind to the port
+        let port = 29123; // Use a different port to avoid conflicts
+
+        if let Ok(_listener) = TcpListener::bind(("127.0.0.1", port)) {
+            // Port is available, we can test
+            let config = PrometheusConfig {
+                port,
+                host: "127.0.0.1".to_string(),
+            };
+
+            // Just verify config is created correctly
+            assert_eq!(config.port, port);
+        }
+    }
+
+    // ============= Integration Test Helpers =============
+
+    #[test]
+    fn test_metrics_endpoint_accessibility() {
+        // This would be an integration test in practice
+        // Here we just verify the configuration
+        let config = PrometheusConfig {
+            port: 29000,
+            host: "127.0.0.1".to_string(),
+        };
+
+        let ip_addr: IpAddr = config.host.parse().unwrap();
+        let socket_addr = SocketAddr::new(ip_addr, config.port);
+
+        assert_eq!(socket_addr.to_string(), "127.0.0.1:29000");
+    }
+
+    #[test]
+    fn test_concurrent_metric_updates() {
+        // Test that metric updates can be called concurrently
+        use std::sync::atomic::{AtomicBool, Ordering};
+        use std::sync::Arc;
+        use std::thread;
+
+        let done = Arc::new(AtomicBool::new(false));
+        let mut handles = vec![];
+
+        for i in 0..3 {
+            let done_clone = done.clone();
+            let handle = thread::spawn(move || {
+                let worker = format!("http://worker{}", i);
+                while !done_clone.load(Ordering::Relaxed) {
+                    RouterMetrics::set_worker_load(&worker, i * 10);
+                    RouterMetrics::record_processed_request(&worker);
+                    thread::sleep(Duration::from_millis(1));
+                }
+            });
+            handles.push(handle);
+        }
+
+        // Let threads run briefly
+        thread::sleep(Duration::from_millis(10));
+        done.store(true, Ordering::Relaxed);
+
+        // Wait for all threads
+        for handle in handles {
+            handle.join().unwrap();
+        }
+
+        // If we get here without panic, concurrent access works
+        assert!(true);
+    }
+
+    // ============= Edge Cases Tests =============
+
+    #[test]
+    fn test_empty_string_metrics() {
+        // Test that empty strings don't cause issues
+        RouterMetrics::record_request("");
+        RouterMetrics::set_worker_health("", true);
+        RouterMetrics::record_policy_decision("", "");
+
+        // If we get here without panic, empty strings are handled
+        assert!(true);
+    }
+
+    #[test]
+    fn test_very_long_metric_labels() {
+        let long_label = "a".repeat(1000);
+
+        RouterMetrics::record_request(&long_label);
+        RouterMetrics::set_worker_health(&long_label, false);
+
+        // If we get here without panic, long labels are handled
+        assert!(true);
+    }
+
+    #[test]
+    fn test_special_characters_in_labels() {
+        let special_labels = vec![
+            "test/with/slashes",
+            "test-with-dashes",
+            "test_with_underscores",
+            "test.with.dots",
+            "test:with:colons",
+        ];
+
+        for label in special_labels {
+            RouterMetrics::record_request(label);
+            RouterMetrics::set_worker_health(label, true);
+        }
+
+        // If we get here without panic, special characters are handled
+        assert!(true);
+    }
+
+    #[test]
+    fn test_extreme_metric_values() {
+        // Test extreme values
+        RouterMetrics::set_active_workers(0);
+        RouterMetrics::set_active_workers(usize::MAX);
+
+        RouterMetrics::set_worker_load("worker", 0);
+        RouterMetrics::set_worker_load("worker", usize::MAX);
+
+        RouterMetrics::record_request_duration("route", Duration::from_nanos(1));
+        RouterMetrics::record_request_duration("route", Duration::from_secs(86400)); // 24 hours
+
+        // If we get here without panic, extreme values are handled
+        assert!(true);
+    }
+}

sgl-router/src/routers/pd_types.rs

@@ -58,7 +58,7 @@ pub enum PDSelectionPolicy {
     },
 }
 // Bootstrap types from PDLB
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Debug, Deserialize, Serialize, PartialEq)]
 #[serde(untagged)]
 pub enum SingleOrBatch<T> {
     Single(T),