chore(starknet_monitoring_endpoint): add metric test utils

commit-id:523a1d3e
starkware-libs · Dec 23, 2024 · fd45754 · fd45754
1 parent fa49151
commit fd45754
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 2 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/starknet_monitoring_endpoint/Cargo.toml b/crates/starknet_monitoring_endpoint/Cargo.toml
@@ -6,7 +6,7 @@ repository.workspace = true
 license-file.workspace = true
 
 [features]
-testing = ["tokio", "tower"]
+testing = ["num-traits", "thiserror", "tokio", "tower"]
 
 [lints]
 workspace = true
@@ -16,9 +16,11 @@ axum.workspace = true
 hyper = { workspace = true }
 infra_utils.workspace = true
 metrics-exporter-prometheus.workspace = true
+num-traits = { workspace = true, optional = true }
 papyrus_config.workspace = true
 serde.workspace = true
 starknet_sequencer_infra.workspace = true
+thiserror = { workspace = true, optional = true }
 tokio = { workspace = true, optional = true }
 tower = { workspace = true, optional = true }
 tracing.workspace = true

diff --git a/crates/starknet_monitoring_endpoint/src/test_utils.rs b/crates/starknet_monitoring_endpoint/src/test_utils.rs
@@ -1,14 +1,30 @@
 use std::net::{IpAddr, SocketAddr};
+use std::str::FromStr;
 
 use axum::body::Body;
 use axum::http::Request;
 use hyper::client::HttpConnector;
 use hyper::Client;
+use infra_utils::metrics::parse_numeric_metric;
 use infra_utils::run_until::run_until;
 use infra_utils::tracing::{CustomLogger, TraceLevel};
+use num_traits::Num;
+use thiserror::Error;
 use tracing::info;
 
-use crate::monitoring_endpoint::{ALIVE, MONITORING_PREFIX};
+use crate::monitoring_endpoint::{ALIVE, METRICS, MONITORING_PREFIX};
+
+// TODO(Tsabary): rename IsAliveClient to MonitoringClient.
+
+#[derive(Clone, Debug, Error, PartialEq, Eq)]
+pub enum MonitoringClientError {
+    #[error("Failed to connect, error details: {}", connection_error)]
+    ConnectionError { connection_error: String },
+    #[error("Erroneous status: {}", status)]
+    ResponseStatusError { status: String },
+    #[error("Missing metric name: {}", metric_name)]
+    MetricNotFound { metric_name: String },
+}
 
 /// Client for querying 'alive' status of an http server.
 pub struct IsAliveClient {
@@ -46,6 +62,35 @@ impl IsAliveClient {
             .ok_or(())
             .map(|_| ())
     }
+
+    pub async fn get_metric<T: Num + FromStr>(
+        &self,
+        metric_name: &str,
+    ) -> Result<T, MonitoringClientError> {
+        // Query the server for metrics.
+        let response = self
+            .client
+            .request(build_request(&self.socket.ip(), self.socket.port(), METRICS))
+            .await
+            .map_err(|err| MonitoringClientError::ConnectionError {
+                connection_error: err.to_string(),
+            })?;
+
+        // Check response status.
+        if !response.status().is_success() {
+            return Err(MonitoringClientError::ResponseStatusError {
+                status: format!("{:?}", response.status()),
+            });
+        }
+
+        // Parse the response body.
+        let body_bytes = hyper::body::to_bytes(response.into_body()).await.unwrap();
+        let body_string = String::from_utf8(body_bytes.to_vec()).unwrap();
+
+        // Extract and return the metric value, or a suitable error.
+        parse_numeric_metric::<T>(&body_string, metric_name)
+            .ok_or(MonitoringClientError::MetricNotFound { metric_name: metric_name.to_string() })
+    }
 }
 
 pub(crate) fn build_request(ip: &IpAddr, port: u16, method: &str) -> Request<Body> {