From 7d2cfdaa9d996a090b0f1a8f83e446662daf4005 Mon Sep 17 00:00:00 2001
From: Itay Tsabary <itayt@starkware.co>
Date: Sun, 22 Dec 2024 20:34:21 +0200
Subject: [PATCH] chore(starknet_monitoring_endpoint): add metric test utils

commit-id:523a1d3e
---
 Cargo.lock                                    |  2 +
 .../starknet_monitoring_endpoint/Cargo.toml   |  4 +-
 .../src/test_utils.rs                         | 49 ++++++++++++++++++-
 3 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 6096e4aaac..0643fd933b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -10658,10 +10658,12 @@ dependencies = [
  "infra_utils",
  "metrics 0.21.1",
  "metrics-exporter-prometheus",
+ "num-traits 0.2.19",
  "papyrus_config",
  "pretty_assertions",
  "serde",
  "starknet_sequencer_infra",
+ "thiserror",
  "tokio",
  "tower 0.4.13",
  "tracing",
diff --git a/crates/starknet_monitoring_endpoint/Cargo.toml b/crates/starknet_monitoring_endpoint/Cargo.toml
index 4bc35f0fb5..266cb77068 100644
--- 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
index 8d284fc5aa..22759aa86e 100644
--- a/crates/starknet_monitoring_endpoint/src/test_utils.rs
+++ b/crates/starknet_monitoring_endpoint/src/test_utils.rs
@@ -1,14 +1,31 @@
 use std::net::{IpAddr, SocketAddr};
+use std::str::FromStr;
 
 use axum::body::Body;
 use axum::http::Request;
+use hyper::body::to_bytes;
 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,8 +63,38 @@ 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 = 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() })
+    }
 }
 
+// TODO(Tsabary): use socket instead of ip and port.
 pub(crate) fn build_request(ip: &IpAddr, port: u16, method: &str) -> Request<Body> {
     Request::builder()
         .uri(format!("http://{ip}:{port}/{MONITORING_PREFIX}/{method}").as_str())