|
| 1 | +use core::time::Duration; |
| 2 | +use std::borrow::Cow; |
| 3 | +use std::sync::Arc; |
| 4 | + |
| 5 | +use axum::http::Request; |
| 6 | +use hyper::StatusCode; |
| 7 | +use nativelink_config::cas_server::HealthConfig; |
| 8 | +use nativelink_macro::nativelink_test; |
| 9 | +use nativelink_service::health_server::HealthServer; |
| 10 | +use nativelink_util::health_utils::{ |
| 11 | + HealthRegistry, HealthRegistryBuilder, HealthStatus, HealthStatusIndicator, |
| 12 | +}; |
| 13 | +use pretty_assertions::assert_eq; |
| 14 | +use serde_json::Value; |
| 15 | +use tonic::async_trait; |
| 16 | +use tonic::body::Body; |
| 17 | +use tonic::service::Routes; |
| 18 | +use tower::{Service, ServiceExt}; |
| 19 | + |
| 20 | +async fn health_tester( |
| 21 | + health_registry: HealthRegistry, |
| 22 | + expected_status_code: StatusCode, |
| 23 | + expected_result: &str, |
| 24 | + config: HealthConfig, |
| 25 | +) -> Result<(), Box<dyn core::error::Error>> { |
| 26 | + let health_server = HealthServer::new(health_registry, &config); |
| 27 | + |
| 28 | + let tonic_services = Routes::builder().routes(); |
| 29 | + |
| 30 | + let mut svc = tonic_services |
| 31 | + .into_axum_router() |
| 32 | + .route_service("/status", health_server); |
| 33 | + |
| 34 | + let request = Request::builder() |
| 35 | + .method("GET") |
| 36 | + .uri("/status") |
| 37 | + .body(Body::empty())?; |
| 38 | + let response: hyper::Response<axum::body::Body> = |
| 39 | + svc.as_service().ready().await?.call(request).await?; |
| 40 | + assert_eq!(response.status(), expected_status_code); |
| 41 | + |
| 42 | + let raw_json = String::from_utf8( |
| 43 | + axum::body::to_bytes(response.into_body(), usize::MAX) |
| 44 | + .await? |
| 45 | + .to_vec(), |
| 46 | + )?; |
| 47 | + let parsed_json: Value = serde_json::from_str(&raw_json)?; |
| 48 | + assert_eq!( |
| 49 | + serde_json::to_string_pretty(&parsed_json)?, |
| 50 | + String::from(expected_result) |
| 51 | + ); |
| 52 | + Ok(()) |
| 53 | +} |
| 54 | + |
| 55 | +#[nativelink_test] |
| 56 | +async fn basic_health_test() -> Result<(), Box<dyn core::error::Error>> { |
| 57 | + let health_registry = HealthRegistryBuilder::new("foo").build(); |
| 58 | + health_tester( |
| 59 | + health_registry, |
| 60 | + StatusCode::OK, |
| 61 | + "[]", |
| 62 | + HealthConfig::default(), |
| 63 | + ) |
| 64 | + .await |
| 65 | +} |
| 66 | + |
| 67 | +struct TestIndicator {} |
| 68 | + |
| 69 | +#[async_trait] |
| 70 | +impl HealthStatusIndicator for TestIndicator { |
| 71 | + fn get_name(&self) -> &'static str { |
| 72 | + "test_indicator" |
| 73 | + } |
| 74 | + |
| 75 | + fn struct_name(&self) -> &'static str { |
| 76 | + "TestIndicator" |
| 77 | + } |
| 78 | + |
| 79 | + async fn check_health(&self, _namespace: Cow<'static, str>) -> HealthStatus { |
| 80 | + HealthStatus::Ok { |
| 81 | + struct_name: self.struct_name(), |
| 82 | + message: Cow::Borrowed("all good"), |
| 83 | + } |
| 84 | + } |
| 85 | +} |
| 86 | + |
| 87 | +#[nativelink_test] |
| 88 | +async fn health_test_with_item() -> Result<(), Box<dyn core::error::Error>> { |
| 89 | + let mut health_registry_builder = HealthRegistryBuilder::new("foo"); |
| 90 | + health_registry_builder.register_indicator(Arc::new(TestIndicator {})); |
| 91 | + let health_registry = health_registry_builder.build(); |
| 92 | + health_tester( |
| 93 | + health_registry, |
| 94 | + StatusCode::OK, |
| 95 | + r#"[ |
| 96 | + { |
| 97 | + "namespace": "/foo/test_indicator", |
| 98 | + "status": { |
| 99 | + "Ok": { |
| 100 | + "struct_name": "TestIndicator", |
| 101 | + "message": "all good" |
| 102 | + } |
| 103 | + } |
| 104 | + } |
| 105 | +]"#, |
| 106 | + HealthConfig::default(), |
| 107 | + ) |
| 108 | + .await |
| 109 | +} |
| 110 | + |
| 111 | +struct TestSleepIndicator {} |
| 112 | + |
| 113 | +#[async_trait] |
| 114 | +impl HealthStatusIndicator for TestSleepIndicator { |
| 115 | + fn get_name(&self) -> &'static str { |
| 116 | + "test_sleep_indicator" |
| 117 | + } |
| 118 | + |
| 119 | + async fn check_health(&self, _namespace: Cow<'static, str>) -> HealthStatus { |
| 120 | + tokio::time::sleep(Duration::MAX).await; |
| 121 | + unreachable!("Because we sleep forever"); |
| 122 | + } |
| 123 | + |
| 124 | + fn struct_name(&self) -> &'static str { |
| 125 | + "TestSleepIndicator" |
| 126 | + } |
| 127 | +} |
| 128 | + |
| 129 | +#[nativelink_test] |
| 130 | +async fn health_test_with_sleep() -> Result<(), Box<dyn core::error::Error>> { |
| 131 | + let mut health_registry_builder = HealthRegistryBuilder::new("foo"); |
| 132 | + health_registry_builder.register_indicator(Arc::new(TestSleepIndicator {})); |
| 133 | + let health_registry = health_registry_builder.build(); |
| 134 | + health_tester( |
| 135 | + health_registry, |
| 136 | + StatusCode::SERVICE_UNAVAILABLE, |
| 137 | + r#"[ |
| 138 | + { |
| 139 | + "namespace": "/foo/test_sleep_indicator", |
| 140 | + "status": { |
| 141 | + "Timeout": { |
| 142 | + "struct_name": "TestSleepIndicator" |
| 143 | + } |
| 144 | + } |
| 145 | + } |
| 146 | +]"#, |
| 147 | + HealthConfig { |
| 148 | + timeout_seconds: 1, |
| 149 | + ..Default::default() |
| 150 | + }, |
| 151 | + ) |
| 152 | + .await?; |
| 153 | + assert!(logs_contain( |
| 154 | + "Timeout during health check struct_name=\"TestSleepIndicator\"" |
| 155 | + )); |
| 156 | + Ok(()) |
| 157 | +} |
0 commit comments