|
1 | 1 | use asap_planner::{Controller, ControllerError, PromQLSchema, RuntimeOptions, StreamingEngine}; |
2 | 2 | use promql_utilities::data_model::KeyByLabelNames; |
3 | 3 | use promql_utilities::query_logics::enums::AggregationType; |
| 4 | +use std::io::{Read, Write}; |
| 5 | +use std::net::TcpListener; |
4 | 6 | use std::path::Path; |
5 | 7 |
|
6 | 8 | // ─── helpers ───────────────────────────────────────────────────────────────── |
@@ -28,6 +30,29 @@ fn http_requests_schema() -> PromQLSchema { |
28 | 30 | ) |
29 | 31 | } |
30 | 32 |
|
| 33 | +fn single_request_prometheus_server( |
| 34 | + status: &str, |
| 35 | + body: &str, |
| 36 | +) -> (String, std::thread::JoinHandle<String>) { |
| 37 | + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); |
| 38 | + let url = format!("http://{}", listener.local_addr().unwrap()); |
| 39 | + let status = status.to_string(); |
| 40 | + let body = body.to_string(); |
| 41 | + let server = std::thread::spawn(move || { |
| 42 | + let (mut stream, _) = listener.accept().unwrap(); |
| 43 | + let mut request = [0u8; 4096]; |
| 44 | + let bytes_read = stream.read(&mut request).unwrap(); |
| 45 | + let response = format!( |
| 46 | + "HTTP/1.1 {status}\r\nContent-Length: {}\r\nContent-Type: application/json\r\nConnection: close\r\n\r\n{}", |
| 47 | + body.len(), |
| 48 | + body |
| 49 | + ); |
| 50 | + stream.write_all(response.as_bytes()).unwrap(); |
| 51 | + String::from_utf8_lossy(&request[..bytes_read]).into_owned() |
| 52 | + }); |
| 53 | + (url, server) |
| 54 | +} |
| 55 | + |
31 | 56 | #[test] |
32 | 57 | fn config_file_sliding_window_override_generates_per_query_candidates() { |
33 | 58 | let controller = Controller::from_file_with_schema( |
@@ -806,6 +831,129 @@ metrics: |
806 | 831 | assert!(result.is_ok()); |
807 | 832 | } |
808 | 833 |
|
| 834 | +#[test] |
| 835 | +fn hinted_metric_skips_prometheus_schema_discovery() { |
| 836 | + let config = tempfile::NamedTempFile::new().unwrap(); |
| 837 | + std::fs::write( |
| 838 | + config.path(), |
| 839 | + r#" |
| 840 | +query_groups: |
| 841 | + - id: 1 |
| 842 | + queries: |
| 843 | + - "rate(http_requests_total[5m])" |
| 844 | + repetition_delay_ms: 300000 |
| 845 | +metrics: |
| 846 | + - metric: "http_requests_total" |
| 847 | + labels: ["instance"] |
| 848 | +"#, |
| 849 | + ) |
| 850 | + .unwrap(); |
| 851 | + |
| 852 | + let unused_listener = TcpListener::bind("127.0.0.1:0").unwrap(); |
| 853 | + let unused_port = unused_listener.local_addr().unwrap().port(); |
| 854 | + drop(unused_listener); |
| 855 | + |
| 856 | + // The port has no listener. A successful construction proves the hinted |
| 857 | + // metric did not attempt URL-based discovery. |
| 858 | + let controller = Controller::from_file( |
| 859 | + config.path(), |
| 860 | + default_opts(), |
| 861 | + &format!("http://127.0.0.1:{unused_port}"), |
| 862 | + ) |
| 863 | + .unwrap(); |
| 864 | + |
| 865 | + assert_eq!(controller.generate().unwrap().inference_query_count(), 1); |
| 866 | +} |
| 867 | + |
| 868 | +#[test] |
| 869 | +fn unhinted_metric_discovers_schema_from_prometheus() { |
| 870 | + let config = tempfile::NamedTempFile::new().unwrap(); |
| 871 | + std::fs::write( |
| 872 | + config.path(), |
| 873 | + r#" |
| 874 | +query_groups: |
| 875 | + - id: 1 |
| 876 | + queries: |
| 877 | + - "rate(http_requests_total[5m])" |
| 878 | + repetition_delay_ms: 300000 |
| 879 | +"#, |
| 880 | + ) |
| 881 | + .unwrap(); |
| 882 | + |
| 883 | + let (url, server) = single_request_prometheus_server( |
| 884 | + "200 OK", |
| 885 | + r#"{"status":"success","data":[{"__name__":"http_requests_total","instance":"one"}]}"#, |
| 886 | + ); |
| 887 | + |
| 888 | + let controller = Controller::from_file(config.path(), default_opts(), &url).unwrap(); |
| 889 | + |
| 890 | + assert_eq!(controller.generate().unwrap().inference_query_count(), 1); |
| 891 | + assert!(server |
| 892 | + .join() |
| 893 | + .unwrap() |
| 894 | + .starts_with("GET /api/v1/series?match%5B%5D=http_requests_total")); |
| 895 | +} |
| 896 | + |
| 897 | +#[test] |
| 898 | +fn hints_remain_authoritative_while_unhinted_metrics_are_discovered() { |
| 899 | + let config = tempfile::NamedTempFile::new().unwrap(); |
| 900 | + std::fs::write( |
| 901 | + config.path(), |
| 902 | + r#" |
| 903 | +query_groups: |
| 904 | + - id: 1 |
| 905 | + queries: |
| 906 | + - "sum by (hint_label) (hinted_metric)" |
| 907 | + - "rate(unhinted_metric[5m])" |
| 908 | + repetition_delay_ms: 300000 |
| 909 | +metrics: |
| 910 | + - metric: "hinted_metric" |
| 911 | + labels: ["hint_label"] |
| 912 | +"#, |
| 913 | + ) |
| 914 | + .unwrap(); |
| 915 | + let (url, server) = single_request_prometheus_server( |
| 916 | + "200 OK", |
| 917 | + r#"{"status":"success","data":[{"__name__":"unhinted_metric","instance":"one"}]}"#, |
| 918 | + ); |
| 919 | + |
| 920 | + let controller = Controller::from_file(config.path(), default_opts(), &url).unwrap(); |
| 921 | + |
| 922 | + assert_eq!(controller.generate().unwrap().inference_query_count(), 2); |
| 923 | + assert!(server |
| 924 | + .join() |
| 925 | + .unwrap() |
| 926 | + .starts_with("GET /api/v1/series?match%5B%5D=unhinted_metric")); |
| 927 | +} |
| 928 | + |
| 929 | +#[test] |
| 930 | +fn unhinted_metric_propagates_prometheus_discovery_error() { |
| 931 | + let config = tempfile::NamedTempFile::new().unwrap(); |
| 932 | + std::fs::write( |
| 933 | + config.path(), |
| 934 | + r#" |
| 935 | +query_groups: |
| 936 | + - id: 1 |
| 937 | + queries: |
| 938 | + - "rate(http_requests_total[5m])" |
| 939 | + repetition_delay_ms: 300000 |
| 940 | +"#, |
| 941 | + ) |
| 942 | + .unwrap(); |
| 943 | + let (url, server) = single_request_prometheus_server( |
| 944 | + "422 Unprocessable Entity", |
| 945 | + r#"{"status":"error","error":"too many series"}"#, |
| 946 | + ); |
| 947 | + |
| 948 | + let result = Controller::from_file(config.path(), default_opts(), &url); |
| 949 | + |
| 950 | + assert!(matches!(result, Err(ControllerError::PrometheusClient(_)))); |
| 951 | + assert!(server |
| 952 | + .join() |
| 953 | + .unwrap() |
| 954 | + .starts_with("GET /api/v1/series?match%5B%5D=http_requests_total")); |
| 955 | +} |
| 956 | + |
809 | 957 | // --- Overlapping window tests --- |
810 | 958 | // Queries where range vector > t_repeat: e.g. [5m] repeated every 60s. |
811 | 959 | // Windows are always tumbling (sliding disabled); the planner emits windowSize=t_repeat |
|
0 commit comments