|
@@ -1,6 +1,7 @@
|
|
|
use lazy_static::lazy_static;
|
|
use lazy_static::lazy_static;
|
|
|
use prometheus::{IntCounter, Registry};
|
|
use prometheus::{IntCounter, Registry};
|
|
|
use std::result::Result;
|
|
use std::result::Result;
|
|
|
|
|
+use std::sync::atomic::{AtomicBool, Ordering};
|
|
|
use warp::{Filter, Rejection, Reply};
|
|
use warp::{Filter, Rejection, Reply};
|
|
|
|
|
|
|
|
lazy_static! {
|
|
lazy_static! {
|
|
@@ -9,6 +10,32 @@ lazy_static! {
|
|
|
pub static ref REGISTRY: Registry = Registry::new();
|
|
pub static ref REGISTRY: Registry = Registry::new();
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
|
|
+#[cfg(test)]
|
|
|
|
|
+lazy_static! {
|
|
|
|
|
+ static ref METRICS_REGISTERED: AtomicBool = AtomicBool::new(false);
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+#[cfg(test)]
|
|
|
|
|
+fn register_custom_metrics() {
|
|
|
|
|
+ if METRICS_REGISTERED
|
|
|
|
|
+ .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
|
|
|
|
|
+ .is_ok()
|
|
|
|
|
+ {
|
|
|
|
|
+ REGISTRY
|
|
|
|
|
+ .register(Box::new(INCOMING_REQUESTS.clone()))
|
|
|
|
|
+ .expect("collector can be registered");
|
|
|
|
|
+ // Add more metrics here as needed
|
|
|
|
|
+ }
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+#[cfg(not(test))]
|
|
|
|
|
+fn register_custom_metrics() {
|
|
|
|
|
+ REGISTRY
|
|
|
|
|
+ .register(Box::new(INCOMING_REQUESTS.clone()))
|
|
|
|
|
+ .expect("collector can be registered");
|
|
|
|
|
+ // Add more metrics here as needed
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
pub async fn web_main() {
|
|
pub async fn web_main() {
|
|
|
register_custom_metrics();
|
|
register_custom_metrics();
|
|
|
|
|
|
|
@@ -20,12 +47,6 @@ pub async fn web_main() {
|
|
|
.await;
|
|
.await;
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
-fn register_custom_metrics() {
|
|
|
|
|
- REGISTRY
|
|
|
|
|
- .register(Box::new(INCOMING_REQUESTS.clone()))
|
|
|
|
|
- .expect("collector can be registered");
|
|
|
|
|
-}
|
|
|
|
|
-
|
|
|
|
|
async fn status_handler() -> Result<impl Reply, Rejection> {
|
|
async fn status_handler() -> Result<impl Reply, Rejection> {
|
|
|
Ok("ok")
|
|
Ok("ok")
|
|
|
}
|
|
}
|
|
@@ -63,3 +84,76 @@ async fn metrics_handler() -> Result<impl Reply, Rejection> {
|
|
|
res.push_str(&res_custom);
|
|
res.push_str(&res_custom);
|
|
|
Ok(res)
|
|
Ok(res)
|
|
|
}
|
|
}
|
|
|
|
|
+
|
|
|
|
|
+#[cfg(test)]
|
|
|
|
|
+mod tests {
|
|
|
|
|
+ use super::*;
|
|
|
|
|
+ use reqwest::{Client, StatusCode};
|
|
|
|
|
+ use warp::Filter;
|
|
|
|
|
+
|
|
|
|
|
+ // Start a Warp server for testing
|
|
|
|
|
+ async fn setup() -> String {
|
|
|
|
|
+ register_custom_metrics();
|
|
|
|
|
+
|
|
|
|
|
+ let metrics_route = warp::path!("metrics").and_then(metrics_handler);
|
|
|
|
|
+ let status_route = warp::path!("status").and_then(status_handler);
|
|
|
|
|
+
|
|
|
|
|
+ let routes = metrics_route.or(status_route);
|
|
|
|
|
+ let (addr, server) = warp::serve(routes).bind_ephemeral(([127, 0, 0, 1], 0)); // Bind to a random port
|
|
|
|
|
+
|
|
|
|
|
+ tokio::spawn(async move {
|
|
|
|
|
+ server.await;
|
|
|
|
|
+ }); // Spawn the server in a background task
|
|
|
|
|
+
|
|
|
|
|
+ format!("http://{}", addr) // Return the address
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ #[tokio::test]
|
|
|
|
|
+ async fn test_status_handler() {
|
|
|
|
|
+ let base_url = setup().await;
|
|
|
|
|
+ println!("{}", base_url);
|
|
|
|
|
+ let client = Client::builder().no_proxy().build().unwrap();
|
|
|
|
|
+
|
|
|
|
|
+ let response = client
|
|
|
|
|
+ .get(format!("{}/status", base_url))
|
|
|
|
|
+ .send()
|
|
|
|
|
+ .await
|
|
|
|
|
+ .unwrap();
|
|
|
|
|
+
|
|
|
|
|
+ assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
|
+ assert_eq!(response.text().await.unwrap(), "ok");
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ #[tokio::test]
|
|
|
|
|
+ async fn test_metrics_handler() {
|
|
|
|
|
+ let base_url = setup().await;
|
|
|
|
|
+ let client = Client::builder().no_proxy().build().unwrap();
|
|
|
|
|
+
|
|
|
|
|
+ let response = client
|
|
|
|
|
+ .get(format!("{}/metrics", base_url))
|
|
|
|
|
+ .send()
|
|
|
|
|
+ .await
|
|
|
|
|
+ .unwrap();
|
|
|
|
|
+
|
|
|
|
|
+ assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
|
+
|
|
|
|
|
+ assert!(response
|
|
|
|
|
+ .text()
|
|
|
|
|
+ .await
|
|
|
|
|
+ .unwrap()
|
|
|
|
|
+ .contains("incoming_requests 0"));
|
|
|
|
|
+
|
|
|
|
|
+ INCOMING_REQUESTS.inc();
|
|
|
|
|
+ let response = client
|
|
|
|
|
+ .get(format!("{}/metrics", base_url))
|
|
|
|
|
+ .send()
|
|
|
|
|
+ .await
|
|
|
|
|
+ .unwrap();
|
|
|
|
|
+
|
|
|
|
|
+ assert!(response
|
|
|
|
|
+ .text()
|
|
|
|
|
+ .await
|
|
|
|
|
+ .unwrap()
|
|
|
|
|
+ .contains("incoming_requests 1"));
|
|
|
|
|
+ }
|
|
|
|
|
+}
|