tidy + fix test

evervault · Aug 16, 2024 · 1a6e0b5 · 1a6e0b5
1 parent 813c85f
commit 1a6e0b5
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 72 deletions.
diff --git a/control-plane/src/health.rs b/control-plane/src/health.rs
@@ -2,8 +2,12 @@ use crate::enclave_connection::get_connection_to_enclave;
 use crate::error::ServerError;
 use hyper::{Body, Request, Response};
 use serde::{Deserialize, Serialize};
-use shared::server::health::{ControlPlaneState, DataPlaneHealthCheck, DataPlaneState, Health};
-use shared::server::{error::ServerResult, tcp::TcpServer, Listener};
+use shared::server::{
+    error::ServerResult,
+    health::{ControlPlaneState, DataPlaneState, HealthCheck},
+    tcp::TcpServer,
+    Listener,
+};
 use shared::ENCLAVE_HEALTH_CHECK_PORT;
 use std::net::SocketAddr;
 use std::sync::OnceLock;
@@ -20,7 +24,7 @@ pub struct HealthCheckServer {
 #[serde(rename_all = "camelCase")]
 struct CombinedHealthCheckLog {
     control_plane: ControlPlaneState,
-    data_plane: DataPlaneHealthCheck,
+    data_plane: DataPlaneState,
 }
 
 pub async fn run_ecs_health_check_service(
@@ -29,9 +33,9 @@ pub async fn run_ecs_health_check_service(
     if is_draining {
         let combined_log = CombinedHealthCheckLog {
             control_plane: ControlPlaneState::Draining,
-            data_plane: DataPlaneHealthCheck::Ok(DataPlaneState::Unknown(
+            data_plane: DataPlaneState::Unknown(
                 "Enclave is draining, data-plane health will not be checked".into(),
-            )),
+            ),
         };
 
         let combined_log_json = serde_json::to_string(&combined_log)?;
@@ -43,7 +47,12 @@ pub async fn run_ecs_health_check_service(
     };
 
     let control_plane = ControlPlaneState::Ok;
-    let data_plane = health_check_data_plane().await;
+    let data_plane = match health_check_data_plane().await {
+        Ok(state) => state,
+        Err(e) => {
+            DataPlaneState::Unknown(format!("Failed to contact data-plane for healthcheck: {e}"))
+        }
+    };
 
     let status_to_return = std::cmp::max(control_plane.status_code(), data_plane.status_code());
 
@@ -60,14 +69,12 @@ pub async fn run_ecs_health_check_service(
         .map_err(ServerError::from)
 }
 
-async fn health_check_data_plane() -> DataPlaneHealthCheck {
-    let stream = get_connection_to_enclave(ENCLAVE_HEALTH_CHECK_PORT)
-        .await
-        .map_err(|e| format!("Error connecting to enclave for healthcheck ({e})"))?;
+type EcsHealthCheckResult = Result<DataPlaneState, ServerError>;
 
-    let (mut sender, connection) = hyper::client::conn::handshake(stream)
-        .await
-        .map_err(|e| format!("Error performing handshake with enclave for healthcheck ({e})"))?;
+async fn health_check_data_plane() -> EcsHealthCheckResult {
+    let stream = get_connection_to_enclave(ENCLAVE_HEALTH_CHECK_PORT).await?;
+
+    let (mut sender, connection) = hyper::client::conn::handshake(stream).await?;
 
     tokio::spawn(connection);
     let request = Request::builder()
@@ -76,16 +83,11 @@ async fn health_check_data_plane() -> DataPlaneHealthCheck {
         .body(Body::empty())
         .expect("Cannot fail");
 
-    let response = sender
-        .send_request(request)
-        .await
-        .map_err(|e| format!("Error sending health check to enclave ({e})"))?;
+    let response = sender.send_request(request).await?;
 
-    let bytes = hyper::body::to_bytes(response.into_parts().1)
-        .await
-        .map_err(|e| format!("Error parsing health check body ({e})"))?;
+    let bytes = hyper::body::to_bytes(response.into_parts().1).await?;
 
-    serde_json::from_slice(&bytes[..]).map_err(|e| format!("Error deserializing health check {e}"))
+    Ok(serde_json::from_slice(&bytes[..])?)
 }
 
 impl HealthCheckServer {
@@ -150,63 +152,26 @@ mod health_check_tests {
         println!("deep response: {response:?}");
         let health_check_log = response_to_health_check_log(response).await;
         assert!(matches!(
-            health_check_log.data_plane.status,
-            HealthCheckStatus::Err
-        ));
-    }
-
-    #[tokio::test]
-    async fn test_enclave_health_check_service_with_skip_deep_set_to_true() {
-        // the data-plane status should error, as its not running
-        let response = run_ecs_health_check_service(true, false).await.unwrap();
-        assert_eq!(response.status(), 200);
-        println!("deep response: {response:?}");
-        let health_check_log = response_to_health_check_log(response).await;
-        assert!(matches!(
-            health_check_log.data_plane.status,
-            HealthCheckStatus::Ignored
+            health_check_log.data_plane,
+            DataPlaneState::Unknown(_)
         ));
     }
 
     #[tokio::test]
     async fn test_enclave_health_check_service_with_draining_set_to_true() {
         // the data-plane status should error, as its not running
         IS_DRAINING.set(true).unwrap();
-        let response = run_ecs_health_check_service(false, true).await.unwrap();
+        let response = run_ecs_health_check_service(true).await.unwrap();
         assert_eq!(response.status(), 500);
         println!("deep response: {response:?}");
         let health_check_log = response_to_health_check_log(response).await;
         assert!(matches!(
-            health_check_log.data_plane.status,
-            HealthCheckStatus::Err
+            health_check_log.data_plane,
+            DataPlaneState::Unknown(_)
         ));
         assert!(matches!(
-            health_check_log.control_plane.status,
-            HealthCheckStatus::Err
+            health_check_log.control_plane,
+            ControlPlaneState::Draining
         ));
     }
-
-    #[tokio::test]
-    async fn test_max_of_enum() {
-        // used in run_ecs_health_check_service to ensure non-success codes have priority
-        let max = [
-            HealthCheckStatus::Err,
-            HealthCheckStatus::Ok,
-            HealthCheckStatus::Unknown,
-            HealthCheckStatus::Ignored,
-        ]
-        .iter()
-        .max()
-        .unwrap();
-        assert!(matches!(max, HealthCheckStatus::Err));
-        let max = [
-            HealthCheckStatus::Ok,
-            HealthCheckStatus::Unknown,
-            HealthCheckStatus::Ignored,
-        ]
-        .iter()
-        .max()
-        .unwrap();
-        assert!(matches!(max, HealthCheckStatus::Unknown));
-    }
 }
diff --git a/shared/src/server/health.rs b/shared/src/server/health.rs
@@ -65,15 +65,10 @@ impl HealthCheck for ControlPlaneState {
     }
 }
 
-pub type DataPlaneHealthCheck = Result<DataPlaneState, String>;
-
-impl HealthCheck for DataPlaneHealthCheck {
+impl HealthCheck for DataPlaneState {
     fn status_code(&self) -> u16 {
         match self {
-            Ok(state) => match state {
-                DataPlaneState::Initialized(diagnostic) if diagnostic.is_healthy => 200,
-                _ => 500,
-            },
+            DataPlaneState::Initialized(diagnostic) if diagnostic.is_healthy => 200,
             _ => 500,
         }
     }