mockforge_observability/prometheus/
exporter.rs

1//! Prometheus metrics exporter
2//!
3//! Provides HTTP endpoints for Prometheus to scrape metrics
4
5use axum::{
6    extract::State,
7    http::StatusCode,
8    response::{IntoResponse, Response},
9    routing::get,
10    Router,
11};
12use prometheus::{Encoder, TextEncoder};
13use std::sync::Arc;
14use tracing::{debug, error};
15
16use super::metrics::MetricsRegistry;
17
18/// Handler for the /metrics endpoint
19pub async fn metrics_handler(
20    State(registry): State<Arc<MetricsRegistry>>,
21) -> Result<impl IntoResponse, MetricsError> {
22    debug!("Serving Prometheus metrics");
23
24    let encoder = TextEncoder::new();
25    let metric_families = registry.registry().gather();
26
27    let mut buffer = Vec::new();
28    encoder.encode(&metric_families, &mut buffer).map_err(|e| {
29        error!("Failed to encode metrics: {}", e);
30        MetricsError::EncodingError(e.to_string())
31    })?;
32
33    let body = String::from_utf8(buffer).map_err(|e| {
34        error!("Failed to convert metrics to UTF-8: {}", e);
35        MetricsError::EncodingError(e.to_string())
36    })?;
37
38    Ok((
39        StatusCode::OK,
40        [("content-type", "text/plain; version=0.0.4; charset=utf-8")],
41        body,
42    ))
43}
44
45/// Health check endpoint for the metrics server
46pub async fn health_handler() -> impl IntoResponse {
47    (StatusCode::OK, "OK")
48}
49
50/// Create a router for the Prometheus metrics endpoint
51pub fn prometheus_router(registry: Arc<MetricsRegistry>) -> Router {
52    Router::new()
53        .route("/metrics", get(metrics_handler))
54        .route("/health", get(health_handler))
55        .with_state(registry)
56}
57
58/// Error type for metrics operations
59#[derive(Debug)]
60pub enum MetricsError {
61    EncodingError(String),
62}
63
64impl IntoResponse for MetricsError {
65    fn into_response(self) -> Response {
66        let (status, message) = match self {
67            MetricsError::EncodingError(msg) => {
68                (StatusCode::INTERNAL_SERVER_ERROR, format!("Encoding error: {}", msg))
69            }
70        };
71
72        (status, message).into_response()
73    }
74}
75
76impl std::fmt::Display for MetricsError {
77    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
78        match self {
79            MetricsError::EncodingError(msg) => write!(f, "Encoding error: {}", msg),
80        }
81    }
82}
83
84impl std::error::Error for MetricsError {}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89    use crate::prometheus::MetricsRegistry;
90
91    #[tokio::test]
92    async fn test_metrics_handler() {
93        let registry = Arc::new(MetricsRegistry::new());
94
95        // Record some test metrics
96        registry.record_http_request("GET", 200, 0.045);
97        registry.record_http_request("POST", 201, 0.123);
98
99        // Call the handler
100        let result = metrics_handler(State(registry)).await;
101        assert!(result.is_ok());
102    }
103
104    #[tokio::test]
105    async fn test_health_handler() {
106        let response = health_handler().await.into_response();
107        assert_eq!(response.status(), StatusCode::OK);
108    }
109
110    #[test]
111    fn test_prometheus_router_creation() {
112        let registry = Arc::new(MetricsRegistry::new());
113        let _router = prometheus_router(registry);
114        // Router should be created successfully
115    }
116
117    #[test]
118    fn test_metrics_error_display() {
119        let error = MetricsError::EncodingError("test error message".to_string());
120        let display = format!("{}", error);
121        assert!(display.contains("Encoding error"));
122        assert!(display.contains("test error message"));
123    }
124
125    #[test]
126    fn test_metrics_error_debug() {
127        let error = MetricsError::EncodingError("test".to_string());
128        let debug = format!("{:?}", error);
129        assert!(debug.contains("EncodingError"));
130        assert!(debug.contains("test"));
131    }
132
133    #[test]
134    fn test_metrics_error_is_error_trait() {
135        let error = MetricsError::EncodingError("test".to_string());
136        // Verify it implements std::error::Error
137        let _: &dyn std::error::Error = &error;
138    }
139
140    #[tokio::test]
141    async fn test_metrics_error_into_response() {
142        let error = MetricsError::EncodingError("encoding failed".to_string());
143        let response = error.into_response();
144        assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
145    }
146
147    #[tokio::test]
148    async fn test_metrics_handler_with_various_metrics() {
149        let registry = Arc::new(MetricsRegistry::new());
150
151        // Record various types of metrics
152        registry.record_http_request("GET", 200, 0.045);
153        registry.record_http_request("POST", 201, 0.123);
154        registry.record_http_request("DELETE", 204, 0.015);
155        registry.record_http_request("PUT", 500, 0.500);
156        registry.record_grpc_request("ListUsers", "OK", 0.025);
157        registry.record_ws_message_sent();
158        registry.record_ws_message_received();
159        registry.record_plugin_execution("test-plugin", true, 0.010);
160        registry.record_error("http", "timeout");
161        registry.update_memory_usage(1024.0 * 1024.0 * 50.0);
162        registry.update_cpu_usage(25.5);
163
164        // Call the handler
165        let result = metrics_handler(State(registry)).await;
166        assert!(result.is_ok());
167    }
168
169    #[tokio::test]
170    async fn test_metrics_handler_empty_registry() {
171        let registry = Arc::new(MetricsRegistry::new());
172        // No metrics recorded
173        let result = metrics_handler(State(registry)).await;
174        assert!(result.is_ok());
175    }
176
177    #[tokio::test]
178    async fn test_metrics_handler_with_grpc_metrics() {
179        let registry = Arc::new(MetricsRegistry::new());
180        registry.record_grpc_request("GetUser", "OK", 0.025);
181        registry.record_grpc_request("CreateUser", "INTERNAL", 0.150);
182        registry.record_grpc_request_with_pillar("ListUsers", "OK", 0.050, "reality");
183
184        let result = metrics_handler(State(registry)).await;
185        assert!(result.is_ok());
186    }
187
188    #[tokio::test]
189    async fn test_metrics_handler_with_websocket_metrics() {
190        let registry = Arc::new(MetricsRegistry::new());
191        registry.record_ws_connection_established();
192        registry.record_ws_message_sent();
193        registry.record_ws_message_received();
194        registry.record_ws_error();
195        registry.record_ws_connection_closed(60.0, "normal");
196
197        let result = metrics_handler(State(registry)).await;
198        assert!(result.is_ok());
199    }
200
201    #[tokio::test]
202    async fn test_metrics_handler_with_smtp_metrics() {
203        let registry = Arc::new(MetricsRegistry::new());
204        registry.record_smtp_connection_established();
205        registry.record_smtp_message_received();
206        registry.record_smtp_message_stored();
207        registry.record_smtp_error("auth_failed");
208        registry.record_smtp_connection_closed();
209
210        let result = metrics_handler(State(registry)).await;
211        assert!(result.is_ok());
212    }
213
214    #[tokio::test]
215    async fn test_metrics_handler_with_marketplace_metrics() {
216        let registry = Arc::new(MetricsRegistry::new());
217        registry.record_marketplace_publish("plugin", true, 2.5);
218        registry.record_marketplace_download("template", true, 0.5);
219        registry.record_marketplace_search("scenario", true, 0.1);
220        registry.record_marketplace_error("plugin", "validation_failed");
221        registry.update_marketplace_items_total("plugin", 150);
222
223        let result = metrics_handler(State(registry)).await;
224        assert!(result.is_ok());
225    }
226
227    #[tokio::test]
228    async fn test_metrics_handler_with_workspace_metrics() {
229        let registry = Arc::new(MetricsRegistry::new());
230        registry.record_workspace_request("workspace-1", "GET", 200, 0.05);
231        registry.update_workspace_active_routes("workspace-1", 10);
232        registry.record_workspace_error("workspace-1", "timeout");
233        registry.increment_workspace_routes("workspace-1");
234        registry.decrement_workspace_routes("workspace-1");
235
236        let result = metrics_handler(State(registry)).await;
237        assert!(result.is_ok());
238    }
239
240    #[tokio::test]
241    async fn test_metrics_handler_with_scenario_metrics() {
242        let registry = Arc::new(MetricsRegistry::new());
243        registry.set_scenario_mode(0); // healthy
244        registry.record_chaos_trigger();
245        registry.set_scenario_mode(3); // chaos
246
247        let result = metrics_handler(State(registry)).await;
248        assert!(result.is_ok());
249    }
250
251    #[tokio::test]
252    async fn test_metrics_handler_with_path_metrics() {
253        let registry = Arc::new(MetricsRegistry::new());
254        registry.record_http_request_with_path("/api/users/123", "GET", 200, 0.05);
255        registry.record_http_request_with_path("/api/users/456", "GET", 200, 0.06);
256        registry.record_http_request_with_path_and_pillar(
257            "/api/items",
258            "POST",
259            201,
260            0.1,
261            "reality",
262        );
263
264        let result = metrics_handler(State(registry)).await;
265        assert!(result.is_ok());
266    }
267}