ai_vitals/
lib.rs

1//! # ai-vitals
2//!
3//! A monitoring tool for LLM endpoints that reports status to Cronitor.
4//!
5//! The library is split into a few main components:
6//!
7//! * monitor: Entrypoint for running the monitoring process. It orchestrates the probing of endpoints and exporting results.
8//! * cli: Handles command-line argument parsing and configuration setup.
9//! * probes: Contains implementations for probing different types of endpoints, such as OpenAI chat completions and embeddings.
10//! * exporters: Contains implementations for exporting monitoring results to different services, currently only Cronitor.
11//!
12//! ## Running Tests
13//!
14//! ```bash
15//! cargo test
16//! ```
17use anyhow::Result;
18use tracing::{error, info};
19
20/// Result of an LLM endpoint probe
21#[derive(Debug, PartialEq)]
22pub enum ProbeResult {
23    Success,
24    Error(u16),
25    Timeout,
26    NetworkError(String),
27}
28
29#[async_trait::async_trait]
30pub trait Probe {
31    fn new(config: cli::Config) -> Result<Self>
32    where
33        Self: std::marker::Sized;
34    async fn probe(&self) -> ProbeResult;
35}
36
37/// State of a Export ping
38#[derive(Debug, Clone, Copy, PartialEq)]
39pub enum PingState {
40    Run,
41    Complete,
42    Fail,
43}
44
45impl PingState {
46    pub fn as_str(&self) -> &'static str {
47        match self {
48            PingState::Run => "run",
49            PingState::Complete => "complete",
50            PingState::Fail => "fail",
51        }
52    }
53}
54
55#[async_trait::async_trait]
56pub trait Export {
57    fn new(config: cli::Config) -> Result<Self>
58    where
59        Self: std::marker::Sized;
60    async fn ping(&self, state: PingState, status_code: u16, message: Option<&str>);
61}
62
63/// Main monitoring orchestrator.
64///
65/// It holds the exporter and probe implementations and runs the monitoring process.
66pub struct Monitor {
67    exporter: Box<dyn Export>,
68    probe: Box<dyn Probe>,
69}
70
71impl Monitor {
72    pub fn new(config: cli::Config) -> Result<Self> {
73        Ok(Monitor {
74            exporter: Box::new(exporters::Cronitor::new(config.clone())?),
75            probe: match config.endpoint_type {
76                probes::Type::OpenAIChatCompletion | probes::Type::OpenAIEmbedding => {
77                    Box::new(probes::OpenAI::new(config.clone())?)
78                }
79                probes::Type::Newman => Box::new(probes::Newman::new(config.clone())?),
80            },
81        })
82    }
83
84    pub async fn run(&self) -> i32 {
85        // Send start ping
86        info!("Sending start ping to Cronitor");
87        self.exporter.ping(PingState::Run, 0, None).await;
88
89        // Probe the endpoint
90        match self.probe.probe().await {
91            ProbeResult::Success => {
92                info!("Sending success ping to Cronitor");
93                self.exporter.ping(PingState::Complete, 0, None).await;
94                info!("SUCCESS: Endpoint responded successfully");
95                0
96            }
97            ProbeResult::Error(status_code) => {
98                info!("Sending failure ping to Cronitor");
99                self.exporter.ping(PingState::Fail, status_code, None).await;
100                error!("FAILURE: Endpoint failed with HTTP {status_code}");
101                1
102            }
103            ProbeResult::Timeout => {
104                info!("Sending timeout ping to Cronitor");
105                self.exporter
106                    .ping(PingState::Fail, 124, Some("Request timeout"))
107                    .await;
108                error!("TIMEOUT: Request timed out");
109                124
110            }
111            ProbeResult::NetworkError(error) => {
112                info!("Sending failure ping to Cronitor");
113                self.exporter
114                    .ping(PingState::Fail, 1, Some(&format!("Network error: {error}")))
115                    .await;
116                error!("FAILURE: Network error: {error}");
117                1
118            }
119        }
120    }
121}
122
123pub mod cli {
124    use clap::Parser;
125
126    use super::probes::Type as ProbeType;
127
128    /// Configuration for the monitoring tool
129    #[derive(Parser, Debug, Clone, PartialEq)]
130    #[command(
131        author,
132        version,
133        about,
134        long_about = "Probe an LLM endpoint and report status to Cronitor."
135    )]
136    pub struct Config {
137        /// Base URL for Cronitor, e.g. https://cronitor.link
138        #[arg(long, env = "CRONITOR_BASE_URL")]
139        pub cronitor_base_url: String,
140
141        /// Base URL for Cronitor, e.g. https://cronitor.link
142        #[arg(long, env = "CRONITOR_API_KEY")]
143        pub cronitor_api_key: Option<String>,
144
145        /// Monitor name / code in Cronitor
146        #[arg(long, env = "MONITOR_NAME")]
147        pub monitor_name: String,
148
149        /// Base URL of the server to probe, e.g. https://my-openai-proxy
150        #[arg(long, env = "SERVER_URL", default_value = "http://localhost:8000/v1")]
151        pub server_url: String,
152
153        /// Optional: Probe type to use for the probe. Currently only "llm" is supported.
154        #[arg(long, env = "ENDPOINT_TYPE", default_value = ProbeType::OpenAIChatCompletion)]
155        pub endpoint_type: ProbeType,
156
157        /// Name of the model to query
158        #[arg(long, env = "MODEL_NAME", default_value = "gpt-4")]
159        pub model_name: String,
160
161        /// Environment descriptor (defaults to "production")
162        #[arg(long, env = "APP_ENV", default_value = "production")]
163        pub env: String,
164
165        /// Request timeout in seconds (default 10)
166        #[arg(long, env = "TIMEOUT_SECONDS", default_value_t = 10)]
167        pub timeout_seconds: u64,
168
169        /// The below all require an API key to be set to take effect.
170
171        /// minFreqRequiredMins catches inactive alerts - if an alert starts but never completes,
172        /// it'll be marked as inactive by Cronitor. To force this into raising an alert,
173        /// we require a successful ping once per any minFreqRequiredMins period.
174        /// Beware - this can cause some unpredictable behaviour in Cronitor
175        #[arg(long, env = "MIN_SUCCESS_FREQ")]
176        pub min_success_freq: Option<u8>,
177
178        /// Which schedule to display in the frontend and to guide CONSECUTIVE_FAILURES_FOR_ALERT.
179        /// If none, one isn't sent to cronitor but will still be running for a cronjob.
180        #[arg(long, env = "SCHEDULE")]
181        pub schedule: Option<String>,
182
183        /// How often we want to resend alerts after the first fails, integer in HOURS
184        #[arg(long, env = "REALERT_INTERVAL")]
185        pub realert_interval: Option<u16>,
186
187        /// Optional: how many failed pings are needed to trigger an alert. Cronitor assumes 1 if unset.
188        #[arg(long, env = "CONSECUTIVE_FAILURES_FOR_ALERT")]
189        pub consecutive_failures: Option<u8>,
190
191        /// Optional: how many missing pings are needed to trigger an alert. Cronitor disables this
192        /// unless specified here as > 0. Requires schedule to be set.
193        #[arg(long, env = "CONSECUTIVE_MISSING_FOR_ALERT")]
194        pub consecutive_missing: Option<u8>,
195
196        /// Optional: Group to put monitor in, mostly for frontend viewing.
197        #[arg(long, env = "MONITOR_GROUP")]
198        pub monitor_group: Option<String>,
199
200        /// Newman-specific options
201        // Path to the Postman collection JSON file
202        #[arg(long, env = "COLLECTION_PATH", default_value = "collection.json")]
203        pub collection_path: String,
204
205        // Path to the Postman environment JSON file
206        #[arg(long, env = "ENVIRONMENT_PATH", default_value = None)]
207        pub environment_path: Option<String>,
208
209        // Delay request by N milliseconds to avoid hitting rate limits
210        #[arg(long, env = "REQUEST_DELAY_MILLISECONDS", default_value = None)]
211        pub request_delay_milliseconds: Option<u64>,
212    }
213
214    impl Default for Config {
215        fn default() -> Self {
216            Config {
217                cronitor_base_url: "https://cronitor.link".to_string(),
218                cronitor_api_key: None,
219                monitor_name: "test-monitor".to_string(),
220                server_url: "https://api.openai.com".to_string(),
221                endpoint_type: ProbeType::OpenAIChatCompletion,
222                model_name: "gpt-4".to_string(),
223                env: "test".to_string(),
224                timeout_seconds: 10,
225                schedule: None,
226                realert_interval: Some(9999),
227                consecutive_failures: Some(1),
228                min_success_freq: Some(60),
229                monitor_group: None,
230                consecutive_missing: Some(1),
231                collection_path: "collection.json".to_string(),
232                environment_path: None,
233                request_delay_milliseconds: None,
234            }
235        }
236    }
237}
238
239pub mod exporters {
240    use anyhow::{Context, Result};
241    use chrono::Utc;
242    use hostname::get;
243    use reqwest::Client;
244    use serde_json::json;
245    use std::time::Duration;
246    use tracing::{error, info};
247
248    use crate::Export;
249
250    use super::{PingState, cli::Config};
251
252    /// Cronitor client to send pings
253    pub struct Cronitor {
254        config: Config,
255        client: Client,
256        host: String,
257        series_id: String,
258    }
259
260    /// Cronitor exporter implementation
261    #[async_trait::async_trait]
262    impl Export for Cronitor {
263        fn new(config: Config) -> Result<Self> {
264            let client = Client::builder()
265                .timeout(Duration::from_secs(config.timeout_seconds))
266                .build()
267                .context("building reqwest client")?;
268
269            let host = get().unwrap_or_default().to_string_lossy().into_owned();
270            let series_id = format!("{}-{}", Utc::now().timestamp(), std::process::id());
271
272            info!("Starting job with series ID: {series_id}");
273
274            Ok(Cronitor {
275                config,
276                client,
277                host,
278                series_id,
279            })
280        }
281
282        async fn ping(&self, state: PingState, status_code: u16, message: Option<&str>) {
283            let url = self.build_ping_url(state, status_code, message);
284
285            match self.client.get(&url).send().await {
286                Ok(resp) if resp.status().is_success() => {
287                    // success: optionally peek at body for debugging
288                    info!("Cronitor ping OK");
289                }
290                Ok(resp) => {
291                    // non-2xx: log status + response body (often has the reason)
292                    let status = resp.status();
293                    let body = resp.text().await.unwrap_or_default(); // consumes resp
294                    error!("Cronitor ping non-2xx {status}: {body}");
295                }
296                Err(e) => {
297                    // request failed before a response was received
298                    error!("Failed to send ping to Cronitor: {e}");
299                }
300            }
301
302            if state == PingState::Run {
303                // The above handles the ping. We also want to update the created monitor if we can.
304
305                let Some(api_key) = self.config.cronitor_api_key.as_deref() else {
306                    info!("No api key, skipping monitor enrichment");
307                    return; // no key => skip update
308                };
309
310                match self
311                    .client
312                    .put("https://cronitor.io/api/monitors")
313                    .basic_auth(api_key, Some("")) // username = API key, blank password
314                    .json(&self.get_monitor_update_payload())
315                    .send()
316                    .await
317                {
318                    Ok(resp) if resp.status().is_success() => {
319                        info!("Monitor enriched successful");
320                    }
321                    Ok(resp) => {
322                        if !resp.status().is_success() {
323                            error!(
324                                "Monitor enrichment failed {}: {}",
325                                resp.status(),
326                                resp.text().await.unwrap_or_default()
327                            );
328                        }
329                    }
330                    Err(err) => {
331                        error!("Failed to enrich Cronitor monitor: {err}");
332                    }
333                }
334            }
335        }
336    }
337
338    /// Internal methods for Cronitor
339    impl Cronitor {
340        pub fn build_ping_url(
341            &self,
342            state: PingState,
343            status_code: u16,
344            message: Option<&str>,
345        ) -> String {
346            let mut url = format!(
347                "{}/{}?state={}&series={}&status_code={}&env={}&host={}",
348                self.config.cronitor_base_url,
349                self.config.monitor_name,
350                state.as_str(),
351                self.series_id,
352                status_code,
353                self.config.env,
354                self.host
355            );
356            if let Some(msg) = message {
357                url.push_str("&message=");
358                url.push_str(&urlencoding::encode(msg));
359            }
360            url
361        }
362
363        pub fn get_monitor_update_payload(&self) -> serde_json::Value {
364            let mut monitor = serde_json::Map::new();
365            monitor.insert("type".into(), json!("job"));
366            monitor.insert("key".into(), json!(self.config.monitor_name));
367
368            if let Some(consecutive_failures) = self.config.consecutive_failures {
369                monitor.insert("failure_tolerance".into(), json!(consecutive_failures));
370            }
371
372            if let Some(schedule) = self.config.schedule.clone() {
373                monitor.insert("schedule".into(), json!(schedule));
374            }
375
376            if let Some(realert_interval) = self.config.realert_interval {
377                monitor.insert("realert_interval".into(), json!(realert_interval));
378            }
379
380            if let (Some(consecutive_missing), Some(_)) = (
381                self.config.consecutive_missing,
382                self.config.schedule.clone(),
383            ) {
384                monitor.insert("schedule_tolerance".into(), json!(consecutive_missing));
385            }
386
387            if let Some(group) = self.config.monitor_group.clone() {
388                monitor.insert("group".into(), json!(group));
389            }
390
391            // always include the duration assertion
392            let mut assertions: Vec<String> = vec![format!(
393                "metric.duration < {}s",
394                self.config.timeout_seconds * 2
395            )];
396
397            if let Some(min_success_freq) = self.config.min_success_freq {
398                assertions.push(format!("job.completes < {min_success_freq} minute"));
399            }
400            monitor.insert("assertions".into(), json!(assertions));
401
402            json!({ "monitors": [serde_json::Value::Object(monitor)] })
403        }
404    }
405
406    #[cfg(test)]
407    mod tests {
408        use super::*;
409
410        #[test]
411        fn test_cronitor_client_creation() {
412            let config = Config::default();
413            let client = Cronitor::new(config);
414            assert!(client.is_ok());
415        }
416
417        #[test]
418        fn test_cronitor_ping_url_construction_without_message() {
419            let config = Config::default();
420            let client = Cronitor::new(config).unwrap();
421
422            let url = client.build_ping_url(PingState::Run, 0, None);
423
424            assert!(url.contains("https://cronitor.link/test-monitor"));
425            assert!(url.contains("state=run"));
426            assert!(url.contains("status_code=0"));
427            assert!(url.contains("env=test"));
428            assert!(url.contains("series="));
429            assert!(url.contains("host="));
430            assert!(!url.contains("message="));
431        }
432
433        #[test]
434        fn test_cronitor_ping_url_construction_with_message() {
435            let config = Config::default();
436            let client = Cronitor::new(config).unwrap();
437
438            let url = client.build_ping_url(PingState::Fail, 500, Some("Test error"));
439
440            assert!(url.contains("https://cronitor.link/test-monitor"));
441            assert!(url.contains("state=fail"));
442            assert!(url.contains("status_code=500"));
443            assert!(url.contains("env=test"));
444            assert!(url.contains("message=Test%20error")); // URL encoded
445        }
446
447        #[test]
448        fn test_cronitor_ping_url_special_characters() {
449            let config = Config::default();
450            let client = Cronitor::new(config).unwrap();
451
452            let url = client.build_ping_url(PingState::Fail, 500, Some("Error: 500 & timeout!"));
453
454            assert!(url.contains("message=Error%3A%20500%20%26%20timeout%21"));
455        }
456    }
457}
458
459pub mod probes {
460    use anyhow::{Context, Result};
461    use reqwest::Client;
462    use serde_json::json;
463    use std::{
464        process::{Command, Stdio},
465        time::Duration,
466    };
467    use tracing::{error, info};
468
469    use super::{ProbeResult, cli::Config};
470
471    // Type of LLM endpoint to probe
472    #[derive(Debug, Clone, Copy, PartialEq, clap::ValueEnum)]
473    pub enum Type {
474        #[value(name = "openai-chat-completion")]
475        OpenAIChatCompletion,
476        #[value(name = "openai-embedding")]
477        OpenAIEmbedding,
478        #[value(name = "newman")]
479        Newman,
480    }
481
482    impl From<Type> for clap::builder::OsStr {
483        fn from(value: Type) -> Self {
484            match value {
485                Type::OpenAIChatCompletion => "openai-chat-completion".into(),
486                Type::OpenAIEmbedding => "openai-embedding".into(),
487                Type::Newman => "newman".into(),
488            }
489        }
490    }
491
492    /// LLM endpoint probe functionality
493    pub struct OpenAI {
494        client: Client,
495        config: Config,
496    }
497
498    /// LLM probe implementation
499    #[async_trait::async_trait]
500    impl super::Probe for OpenAI {
501        fn new(config: Config) -> Result<Self> {
502            let client = Client::builder()
503                .timeout(Duration::from_secs(config.timeout_seconds))
504                .build()
505                .context("building reqwest client")?;
506
507            Ok(OpenAI { client, config })
508        }
509
510        async fn probe(&self) -> ProbeResult {
511            let endpoint = self.build_endpoint_url();
512            let payload = self.build_payload();
513
514            info!("Querying {endpoint}");
515
516            match self.client.post(&endpoint).json(&payload).send().await {
517                Ok(resp) => {
518                    let status = resp.status();
519                    let body = resp.text().await.unwrap_or_default();
520                    info!("Response body: {body}");
521
522                    if status.is_success() {
523                        ProbeResult::Success
524                    } else {
525                        ProbeResult::Error(status.as_u16())
526                    }
527                }
528                Err(e) if e.is_timeout() => ProbeResult::Timeout,
529                Err(e) => ProbeResult::NetworkError(e.to_string()),
530            }
531        }
532    }
533
534    /// Internal methods for OpenAI probe
535    impl OpenAI {
536        pub fn build_endpoint_url(&self) -> String {
537            match self.config.endpoint_type {
538                Type::OpenAIChatCompletion => {
539                    format!("{}/v1/chat/completions", self.config.server_url)
540                }
541                Type::OpenAIEmbedding => format!("{}/v1/embeddings", self.config.server_url),
542                _ => panic!("Unsupported endpoint type"),
543            }
544        }
545
546        pub fn build_payload(&self) -> serde_json::Value {
547            match self.config.endpoint_type {
548                Type::OpenAIChatCompletion => json!({
549                    "model": self.config.model_name,
550                    "messages": [{ "role": "user", "content": "test" }],
551                    "max_tokens": 1,
552                    "priority": -100
553                }),
554                Type::OpenAIEmbedding => json!({
555                    "model": self.config.model_name,
556                    "input": "test",
557                    "priority": -100
558                }),
559                _ => panic!("Unsupported endpoint type"),
560            }
561        }
562    }
563
564    /// Newman probe functionality
565    pub struct Newman {
566        config: Config,
567    }
568
569    /// LLM probe implementation
570    #[async_trait::async_trait]
571    impl super::Probe for Newman {
572        fn new(config: Config) -> Result<Self> {
573            Ok(Newman { config })
574        }
575
576        async fn probe(&self) -> ProbeResult {
577            let mut newman = Command::new("newman");
578
579            newman.arg("run").arg(&self.config.collection_path);
580
581            // Set timeout - timeout is in seconds, but newman expects milliseconds
582            newman
583                .arg("--timeout-request")
584                .arg((self.config.timeout_seconds * 1000).to_string());
585
586            // Optional args if set in config
587            if let Some(env_path) = &self.config.environment_path {
588                newman.arg("-e").arg(env_path);
589            }
590            if let Some(delay) = self.config.request_delay_milliseconds {
591                newman.arg("--delay-request").arg(delay.to_string());
592            }
593
594            if let Ok(child) = newman
595                .stdout(Stdio::piped())
596                .stderr(Stdio::piped())
597                .spawn()
598                .context("spawning newman process")
599            {
600                match child.wait_with_output() {
601                    Ok(output) => {
602                        let status = output.status;
603                        let body = String::from_utf8_lossy(&output.stdout);
604                        info!("--- Newman stdout ---\n {body}");
605
606                        if status.success() {
607                            ProbeResult::Success
608                        } else {
609                            ProbeResult::Error(1)
610                        }
611                    }
612                    Err(e) => {
613                        error!("Failed to wait for newman process: {e}");
614                        ProbeResult::Error(1)
615                    }
616                }
617            } else {
618                error!("Failed to start newman process");
619                ProbeResult::Error(1)
620            }
621        }
622    }
623
624    #[cfg(test)]
625    mod tests {
626        use super::{super::Probe, *};
627        use httpmock::prelude::*;
628        use serde_json::json;
629
630        #[test]
631        fn test_openai_creation() {
632            let config = Config::default();
633            let probe = OpenAI::new(config);
634            assert!(probe.is_ok());
635        }
636
637        #[test]
638        fn test_openai_chat_endpoint_url() {
639            let config = Config {
640                endpoint_type: Type::OpenAIChatCompletion,
641                server_url: "https://api.openai.com".to_string(),
642                ..Default::default()
643            };
644            let probe = OpenAI::new(config).unwrap();
645
646            let url = probe.build_endpoint_url();
647            assert_eq!(url, "https://api.openai.com/v1/chat/completions");
648        }
649
650        #[test]
651        fn test_openai_embedding_endpoint_url() {
652            let config = Config {
653                endpoint_type: Type::OpenAIEmbedding,
654                server_url: "https://api.example.com".to_string(),
655                ..Default::default()
656            };
657            let probe = OpenAI::new(config).unwrap();
658
659            let url = probe.build_endpoint_url();
660            assert_eq!(url, "https://api.example.com/v1/embeddings");
661        }
662
663        #[test]
664        fn test_openai_chat_payload() {
665            let config = Config {
666                endpoint_type: Type::OpenAIChatCompletion,
667                model_name: "a-piece-of-cheese".to_string(),
668                ..Default::default()
669            };
670            let probe = OpenAI::new(config).unwrap();
671
672            let payload = probe.build_payload();
673            let expected = json!({
674                "model": "a-piece-of-cheese",
675                "messages": [{ "role": "user", "content": "test" }],
676                "max_tokens": 1,
677                "priority": -100
678            });
679
680            assert_eq!(payload, expected);
681        }
682
683        #[test]
684        fn test_openai_embedding_payload() {
685            let config = Config {
686                endpoint_type: Type::OpenAIEmbedding,
687                model_name: "text-embedding-ada-002".to_string(),
688                ..Default::default()
689            };
690            let probe = OpenAI::new(config).unwrap();
691
692            let payload = probe.build_payload();
693            let expected = json!({
694                "model": "text-embedding-ada-002",
695                "input": "test",
696                "priority": -100
697            });
698
699            assert_eq!(payload, expected);
700        }
701
702        #[tokio::test]
703        async fn test_openai_successful_response() {
704            let server = MockServer::start();
705
706            // Mock successful LLM response
707            let mock = server.mock(|when, then| {
708                when.method(POST)
709                    .path("/v1/chat/completions")
710                    .json_body(json!({
711                        "model": "gpt-4",
712                        "messages": [{ "role": "user", "content": "test" }],
713                        "max_tokens": 1,
714                        "priority": -100
715                    }));
716                then.status(200).json_body(json!({
717                    "choices": [{"message": {"role": "assistant", "content": "Hello"}}]
718                }));
719            });
720
721            let config = Config {
722                server_url: server.base_url(),
723                endpoint_type: Type::OpenAIChatCompletion,
724                model_name: "gpt-4".to_string(),
725                ..Default::default()
726            };
727
728            let probe = OpenAI::new(config).unwrap();
729            let result = probe.probe().await;
730
731            assert_eq!(result, ProbeResult::Success);
732
733            mock.assert();
734        }
735
736        #[tokio::test]
737        async fn test_openai_http_error_response() {
738            let server = MockServer::start();
739
740            // Mock failed LLM response
741            let mock = server.mock(|when, then| {
742                when.method(POST).path("/v1/embeddings");
743                then.status(420).json_body(json!({
744                    "error": {"message": "Internal server error"}
745                }));
746            });
747
748            let config = Config {
749                server_url: server.base_url(),
750                endpoint_type: Type::OpenAIEmbedding,
751                model_name: "text-embedding-ada-002".to_string(),
752                ..Default::default()
753            };
754
755            let probe = OpenAI::new(config).unwrap();
756            let result = probe.probe().await;
757
758            match result {
759                ProbeResult::Error(status_code) => {
760                    assert_eq!(status_code, 420);
761                }
762                _ => panic!("Expected HTTP error probe result"),
763            }
764
765            mock.assert();
766        }
767
768        #[tokio::test]
769        async fn test_openai_timeout() {
770            let config = Config {
771                server_url: "http://10.255.255.1:12345".to_string(), // Non-routable IP for timeout
772                timeout_seconds: 1,                                  // Very short timeout
773                ..Default::default()
774            };
775
776            let probe = OpenAI::new(config).unwrap();
777            let result = probe.probe().await;
778
779            assert!(matches!(result, ProbeResult::Timeout));
780        }
781
782        #[tokio::test]
783        async fn test_openai_network_error() {
784            let config = Config {
785                server_url: "http://localhost:99999".to_string(), // Invalid port
786                ..Default::default()
787            };
788
789            let probe = OpenAI::new(config).unwrap();
790            let result = probe.probe().await;
791
792            match result {
793                ProbeResult::NetworkError(error) => {
794                    assert!(!error.is_empty());
795                }
796                _ => panic!("Expected network error probe result"),
797            }
798        }
799
800        // Newman probe tests
801        use std::fs;
802        use tempfile::TempDir;
803
804        #[test]
805        fn test_newman_probe_creation() {
806            let config = Config {
807                endpoint_type: Type::Newman,
808                collection_path: "test-collection.json".to_string(),
809                environment_path: Some("test-environment.json".to_string()),
810                ..Default::default()
811            };
812            let probe = Newman::new(config);
813            assert!(probe.is_ok());
814        }
815
816        #[tokio::test]
817        async fn test_newman_probe_with_mock_endpoints() {
818            let server = MockServer::start();
819            let temp_dir = TempDir::new().unwrap();
820
821            // Mock the endpoints from our test collection
822            let health_mock = server.mock(|when, then| {
823                when.method(GET).path("/health");
824                then.status(200).json_body(json!({"status": "ok"}));
825            });
826
827            let user_mock = server.mock(|when, then| {
828                when.method(GET)
829                    .path("/api/v1/users/123")
830                    .header("Authorization", "Bearer test-token-12345");
831                then.status(200)
832                    .json_body(json!({"id": 123, "name": "Test User"}));
833            });
834
835            let create_mock = server.mock(|when, then| {
836                when.method(POST)
837                    .path("/api/v1/resources")
838                    .header("Authorization", "Bearer test-token-12345")
839                    .header("Content-Type", "application/json");
840                then.status(201).json_body(json!({
841                    "id": 789,
842                    "name": "Test Resource",
843                    "description": "Created at 1234567890",
844                    "active": true
845                }));
846            });
847
848            let delete_mock = server.mock(|when, then| {
849                when.method(DELETE)
850                    .path("/api/v1/resources/456")
851                    .header("Authorization", "Bearer test-token-12345");
852                then.status(204);
853            });
854
855            // Create test collection file with server base URL
856            let collection_path = temp_dir.path().join("collection.json");
857            let collection_content = fs::read_to_string("test-collection.json")
858                .unwrap_or_else(|_| include_str!("../test-collection.json").to_string());
859            fs::write(&collection_path, collection_content).unwrap();
860
861            // Create test environment file with mock server URL
862            let environment_path = temp_dir.path().join("environment.json");
863            let environment_content = json!({
864                "id": "test-env",
865                "name": "Test Environment",
866                "values": [
867                    {
868                        "key": "base_url",
869                        "value": server.base_url(),
870                        "enabled": true,
871                        "type": "default"
872                    },
873                    {
874                        "key": "api_token",
875                        "value": "test-token-12345",
876                        "enabled": true,
877                        "type": "secret"
878                    }
879                ],
880                "_postman_variable_scope": "environment"
881            });
882            fs::write(&environment_path, environment_content.to_string()).unwrap();
883
884            let config = Config {
885                endpoint_type: Type::Newman,
886                collection_path: collection_path.to_str().unwrap().to_string(),
887                environment_path: Some(environment_path.to_str().unwrap().to_string()),
888                ..Default::default()
889            };
890
891            let probe = Newman::new(config).unwrap();
892            let result = probe.probe().await;
893
894            // Newman should succeed if all tests pass
895            assert_eq!(result, ProbeResult::Success);
896
897            // Verify all endpoints were called
898            health_mock.assert();
899            user_mock.assert();
900            create_mock.assert();
901            delete_mock.assert();
902        }
903
904        #[tokio::test]
905        async fn test_newman_probe_with_failed_test() {
906            let server = MockServer::start();
907            let temp_dir = TempDir::new().unwrap();
908
909            // Mock health endpoint to return 500 (will fail the test)
910            let health_mock = server.mock(|when, then| {
911                when.method(GET).path("/health");
912                then.status(500).json_body(json!({"error": "Server error"}));
913            });
914
915            // Create a minimal collection with just the health check
916            let collection_path = temp_dir.path().join("collection.json");
917            let collection_content = json!({
918                "info": {
919                    "name": "Test Collection",
920                    "schema": "https://schema.getpostman.com/json/collection/v2.1.0/collection.json"
921                },
922                "item": [{
923                    "name": "Health Check",
924                    "event": [{
925                        "listen": "test",
926                        "script": {
927                            "exec": [
928                                "pm.test(\"Status code is 200\", function () {",
929                                "    pm.response.to.have.status(200);",
930                                "});"
931                            ],
932                            "type": "text/javascript"
933                        }
934                    }],
935                    "request": {
936                        "method": "GET",
937                        "url": "{{base_url}}/health"
938                    }
939                }]
940            });
941            fs::write(&collection_path, collection_content.to_string()).unwrap();
942
943            // Create environment file
944            let environment_path = temp_dir.path().join("environment.json");
945            let environment_content = json!({
946                "id": "test-env",
947                "name": "Test Environment",
948                "values": [{
949                    "key": "base_url",
950                    "value": server.base_url(),
951                    "enabled": true,
952                    "type": "default"
953                }],
954                "_postman_variable_scope": "environment"
955            });
956            fs::write(&environment_path, environment_content.to_string()).unwrap();
957
958            let config = Config {
959                endpoint_type: Type::Newman,
960                collection_path: collection_path.to_str().unwrap().to_string(),
961                environment_path: Some(environment_path.to_str().unwrap().to_string()),
962                ..Default::default()
963            };
964
965            let probe = Newman::new(config).unwrap();
966            let result = probe.probe().await;
967
968            // Newman should fail if any test fails
969            assert_eq!(result, ProbeResult::Error(1));
970
971            health_mock.assert();
972        }
973
974        #[tokio::test]
975        async fn test_newman_probe_without_environment() {
976            let server = MockServer::start();
977            let temp_dir = TempDir::new().unwrap();
978
979            // Mock endpoint without auth
980            let health_mock = server.mock(|when, then| {
981                when.method(GET).path("/health");
982                then.status(200).json_body(json!({"status": "ok"}));
983            });
984
985            // Create collection with hardcoded URL (no environment needed)
986            let collection_path = temp_dir.path().join("collection.json");
987            let collection_content = json!({
988                "info": {
989                    "name": "Test Collection",
990                    "schema": "https://schema.getpostman.com/json/collection/v2.1.0/collection.json"
991                },
992                "item": [{
993                    "name": "Health Check",
994                    "event": [{
995                        "listen": "test",
996                        "script": {
997                            "exec": [
998                                "pm.test(\"Status code is 200\", function () {",
999                                "    pm.response.to.have.status(200);",
1000                                "});"
1001                            ],
1002                            "type": "text/javascript"
1003                        }
1004                    }],
1005                    "request": {
1006                        "method": "GET",
1007                        "url": format!("{}/health", server.base_url())
1008                    }
1009                }]
1010            });
1011            fs::write(&collection_path, collection_content.to_string()).unwrap();
1012
1013            let config = Config {
1014                endpoint_type: Type::Newman,
1015                collection_path: collection_path.to_str().unwrap().to_string(),
1016                environment_path: None, // No environment file
1017                ..Default::default()
1018            };
1019
1020            let probe = Newman::new(config).unwrap();
1021            let result = probe.probe().await;
1022
1023            assert_eq!(result, ProbeResult::Success);
1024            health_mock.assert();
1025        }
1026    }
1027}
1028
1029#[cfg(test)]
1030mod tests {
1031    use super::{Monitor, cli::Config};
1032    use httpmock::prelude::*;
1033    use serde_json::json;
1034
1035    #[tokio::test]
1036    async fn test_monitor_creation() {
1037        let config = Config::default();
1038        let monitor = Monitor::new(config);
1039        assert!(monitor.is_ok());
1040    }
1041
1042    #[tokio::test]
1043    async fn test_monitor_run_success() {
1044        let server = MockServer::start();
1045
1046        // Mock successful LLM response
1047        let llm_mock = server.mock(|when, then| {
1048            when.method(POST).path("/v1/chat/completions");
1049            then.status(200).json_body(
1050                json!({"choices": [{"message": {"role": "assistant", "content": "OK"}}]}),
1051            );
1052        });
1053
1054        // Mock Cronitor pings
1055        let cronitor_run_mock = server.mock(|when, then| {
1056            when.method(GET)
1057                .path("/test-monitor")
1058                .query_param("state", "run");
1059            then.status(200);
1060        });
1061
1062        let cronitor_complete_mock = server.mock(|when, then| {
1063            when.method(GET)
1064                .path("/test-monitor")
1065                .query_param("state", "complete");
1066            then.status(200);
1067        });
1068
1069        let config = Config {
1070            cronitor_base_url: server.base_url(),
1071            server_url: server.base_url(),
1072            ..Default::default()
1073        };
1074
1075        let monitor = Monitor::new(config).unwrap();
1076        let exit_code = monitor.run().await;
1077
1078        assert_eq!(exit_code, 0);
1079        llm_mock.assert();
1080        cronitor_run_mock.assert();
1081        cronitor_complete_mock.assert();
1082    }
1083
1084    #[tokio::test]
1085    async fn test_monitor_run_http_error() {
1086        let server = MockServer::start();
1087
1088        // Mock failed LLM response
1089        let llm_mock = server.mock(|when, then| {
1090            when.method(POST).path("/v1/chat/completions");
1091            then.status(500)
1092                .json_body(json!({"error": {"message": "Server error"}}));
1093        });
1094
1095        // Mock Cronitor pings
1096        let cronitor_run_mock = server.mock(|when, then| {
1097            when.method(GET)
1098                .path("/test-monitor")
1099                .query_param("state", "run");
1100            then.status(200);
1101        });
1102
1103        let cronitor_fail_mock = server.mock(|when, then| {
1104            when.method(GET)
1105                .path("/test-monitor")
1106                .query_param("state", "fail")
1107                .query_param("status_code", "500");
1108            then.status(200);
1109        });
1110
1111        let config = Config {
1112            cronitor_base_url: server.base_url(),
1113            server_url: server.base_url(),
1114            ..Default::default()
1115        };
1116
1117        let monitor = Monitor::new(config).unwrap();
1118        let exit_code = monitor.run().await;
1119
1120        assert_eq!(exit_code, 1);
1121        llm_mock.assert();
1122        cronitor_run_mock.assert();
1123        cronitor_fail_mock.assert();
1124    }
1125
1126    #[tokio::test]
1127    async fn test_monitor_run_timeout() {
1128        let server = MockServer::start();
1129
1130        // Mock Cronitor pings
1131        let cronitor_run_mock = server.mock(|when, then| {
1132            when.method(GET)
1133                .path("/test-monitor")
1134                .query_param("state", "run");
1135            then.status(200);
1136        });
1137
1138        let cronitor_fail_mock = server.mock(|when, then| {
1139            when.method(GET)
1140                .path("/test-monitor")
1141                .query_param("state", "fail")
1142                .query_param("status_code", "124")
1143                .query_param("message", "Request timeout");
1144            then.status(200);
1145        });
1146
1147        let config = Config {
1148            cronitor_base_url: server.base_url(),
1149            server_url: "http://10.255.255.1:12345".to_string(), // Non-routable for timeout
1150            timeout_seconds: 1,
1151            ..Default::default()
1152        };
1153
1154        let monitor = Monitor::new(config).unwrap();
1155        let exit_code = monitor.run().await;
1156
1157        assert_eq!(exit_code, 124); // Timeout exit code
1158        cronitor_run_mock.assert();
1159        cronitor_fail_mock.assert();
1160    }
1161
1162    #[tokio::test]
1163    async fn test_monitor_run_network_error() {
1164        let server = MockServer::start();
1165
1166        // Mock Cronitor pings
1167        let cronitor_run_mock = server.mock(|when, then| {
1168            when.method(GET)
1169                .path("/test-monitor")
1170                .query_param("state", "run");
1171            then.status(200);
1172        });
1173
1174        let cronitor_fail_mock = server.mock(|when, then| {
1175            when.method(GET)
1176                .path("/test-monitor")
1177                .query_param("state", "fail")
1178                .query_param("status_code", "1")
1179                .query_param_exists("message"); // Just check that message parameter exists
1180            then.status(200);
1181        });
1182
1183        let config = Config {
1184            cronitor_base_url: server.base_url(),
1185            server_url: "http://localhost:99999".to_string(), // Invalid port
1186            ..Default::default()
1187        };
1188
1189        let monitor = Monitor::new(config).unwrap();
1190        let exit_code = monitor.run().await;
1191
1192        assert_eq!(exit_code, 1); // Network error exit code
1193        cronitor_run_mock.assert();
1194        cronitor_fail_mock.assert();
1195    }
1196
1197    #[tokio::test]
1198    async fn test_monitor_cronitor_message_validation() {
1199        let server = MockServer::start();
1200
1201        // Mock specific Cronitor ping calls with exact message validation
1202        let cronitor_run_mock = server.mock(|when, then| {
1203            when.method(GET)
1204                .path("/test-monitor")
1205                .query_param("state", "run")
1206                .query_param("status_code", "0")
1207                .query_param("env", "test");
1208            then.status(200);
1209        });
1210
1211        let cronitor_timeout_mock = server.mock(|when, then| {
1212            when.method(GET)
1213                .path("/test-monitor")
1214                .query_param("state", "fail")
1215                .query_param("status_code", "124")
1216                .query_param("message", "Request timeout")
1217                .query_param("env", "test");
1218            then.status(200);
1219        });
1220
1221        let config = Config {
1222            cronitor_base_url: server.base_url(),
1223            server_url: "http://10.255.255.1:12345".to_string(), // Non-routable for timeout
1224            timeout_seconds: 1,
1225            ..Default::default()
1226        };
1227
1228        let monitor = Monitor::new(config).unwrap();
1229        let exit_code = monitor.run().await;
1230
1231        assert_eq!(exit_code, 124);
1232        cronitor_run_mock.assert();
1233        cronitor_timeout_mock.assert();
1234    }
1235}
ai_vitals/lib.rs

ai_vitals/
lib.rs
