eae3c31e53
Patterns applied: - Pattern 2: Float comparison (2x: utils.rs, var_edge_cases_tests.rs) - Pattern 7: Date/time construction (2x: production_streaming.rs, streaming.rs) - Pattern 1: Duration/time ops (2x: rate limiter, semaphore) - Pattern 4: Optional field access (1x: position_tracker.rs) Changes: - data/src/utils.rs: Float sort with NaN handling - data/src/providers/benzinga/production_streaming.rs: Rate limiter + semaphore + date/time - data/src/providers/benzinga/streaming.rs: Date/time construction - risk/src/position_tracker.rs: Emergency fallback counter - risk/tests/var_edge_cases_tests.rs: Test helper float sort Test impact: 0 failures (182/182 passing) Compilation: Clean (0 errors, 0 warnings) Time: 25 min (44% under budget)
632 lines
17 KiB
Rust
632 lines
17 KiB
Rust
//! Comprehensive health check tests for ML Training Service
|
|
//!
|
|
//! Tests cover:
|
|
//! - HTTP /health and /ready endpoints
|
|
//! - Service startup state transitions
|
|
//! - GPU availability checks
|
|
//! - Model checkpoint accessibility
|
|
//! - Database connectivity
|
|
//! - Training job status
|
|
//! - Dependency cascade failures
|
|
//! - Health check latency
|
|
//! - Concurrent health checks
|
|
//! - Health during model training
|
|
//! - Resource exhaustion (GPU memory, disk space)
|
|
|
|
use axum::{
|
|
body::Body,
|
|
http::{Request, StatusCode},
|
|
};
|
|
use std::sync::Arc;
|
|
use std::time::{Duration, Instant};
|
|
use tokio::sync::RwLock;
|
|
use tower::ServiceExt;
|
|
|
|
/// Mock health state for ML training service
|
|
#[derive(Clone)]
|
|
struct MockMlHealthState {
|
|
healthy: Arc<RwLock<bool>>,
|
|
ready: Arc<RwLock<bool>>,
|
|
gpu_available: Arc<RwLock<bool>>,
|
|
checkpoints_accessible: Arc<RwLock<bool>>,
|
|
database_connected: Arc<RwLock<bool>>,
|
|
training_active: Arc<RwLock<bool>>,
|
|
gpu_memory_available: Arc<RwLock<bool>>,
|
|
}
|
|
|
|
impl MockMlHealthState {
|
|
fn new() -> Self {
|
|
Self {
|
|
healthy: Arc::new(RwLock::new(true)),
|
|
ready: Arc::new(RwLock::new(true)),
|
|
gpu_available: Arc::new(RwLock::new(true)),
|
|
checkpoints_accessible: Arc::new(RwLock::new(true)),
|
|
database_connected: Arc::new(RwLock::new(true)),
|
|
training_active: Arc::new(RwLock::new(false)),
|
|
gpu_memory_available: Arc::new(RwLock::new(true)),
|
|
}
|
|
}
|
|
|
|
async fn set_healthy(&self, healthy: bool) {
|
|
*self.healthy.write().await = healthy;
|
|
}
|
|
|
|
async fn set_ready(&self, ready: bool) {
|
|
*self.ready.write().await = ready;
|
|
}
|
|
|
|
async fn set_gpu_available(&self, available: bool) {
|
|
*self.gpu_available.write().await = available;
|
|
}
|
|
|
|
async fn set_checkpoints_accessible(&self, accessible: bool) {
|
|
*self.checkpoints_accessible.write().await = accessible;
|
|
}
|
|
|
|
async fn set_database_connected(&self, connected: bool) {
|
|
*self.database_connected.write().await = connected;
|
|
}
|
|
|
|
async fn set_training_active(&self, active: bool) {
|
|
*self.training_active.write().await = active;
|
|
}
|
|
|
|
async fn set_gpu_memory_available(&self, available: bool) {
|
|
*self.gpu_memory_available.write().await = available;
|
|
}
|
|
|
|
async fn is_healthy(&self) -> bool {
|
|
*self.healthy.read().await
|
|
}
|
|
|
|
async fn is_ready(&self) -> bool {
|
|
*self.ready.read().await
|
|
}
|
|
|
|
async fn is_gpu_available(&self) -> bool {
|
|
*self.gpu_available.read().await
|
|
}
|
|
|
|
async fn is_checkpoints_accessible(&self) -> bool {
|
|
*self.checkpoints_accessible.read().await
|
|
}
|
|
|
|
async fn is_database_connected(&self) -> bool {
|
|
*self.database_connected.read().await
|
|
}
|
|
|
|
async fn is_training_active(&self) -> bool {
|
|
*self.training_active.read().await
|
|
}
|
|
|
|
async fn is_gpu_memory_available(&self) -> bool {
|
|
*self.gpu_memory_available.read().await
|
|
}
|
|
}
|
|
|
|
/// Create health router for ML training service
|
|
fn create_ml_health_router(state: MockMlHealthState) -> axum::Router {
|
|
use axum::{extract::State, routing::get, Json, Router};
|
|
use serde_json::json;
|
|
|
|
async fn health_handler(
|
|
State(state): State<MockMlHealthState>,
|
|
) -> Result<Json<serde_json::Value>, StatusCode> {
|
|
if state.is_healthy().await {
|
|
Ok(Json(json!({
|
|
"status": "healthy",
|
|
"service": "ml_training",
|
|
"version": env!("CARGO_PKG_VERSION")
|
|
})))
|
|
} else {
|
|
Err(StatusCode::SERVICE_UNAVAILABLE)
|
|
}
|
|
}
|
|
|
|
async fn ready_handler(
|
|
State(state): State<MockMlHealthState>,
|
|
) -> Result<Json<serde_json::Value>, StatusCode> {
|
|
if state.is_ready().await {
|
|
Ok(Json(json!({
|
|
"status": "ready",
|
|
"service": "ml_training",
|
|
"version": env!("CARGO_PKG_VERSION")
|
|
})))
|
|
} else {
|
|
Err(StatusCode::SERVICE_UNAVAILABLE)
|
|
}
|
|
}
|
|
|
|
async fn deep_health_handler(
|
|
State(state): State<MockMlHealthState>,
|
|
) -> Result<Json<serde_json::Value>, StatusCode> {
|
|
let gpu_ok = state.is_gpu_available().await;
|
|
let checkpoints_ok = state.is_checkpoints_accessible().await;
|
|
let db_ok = state.is_database_connected().await;
|
|
let gpu_mem_ok = state.is_gpu_memory_available().await;
|
|
let healthy = state.is_healthy().await;
|
|
let training_active = state.is_training_active().await;
|
|
|
|
if healthy && gpu_ok && checkpoints_ok && db_ok && gpu_mem_ok {
|
|
Ok(Json(json!({
|
|
"status": "healthy",
|
|
"service": "ml_training",
|
|
"version": env!("CARGO_PKG_VERSION"),
|
|
"dependencies": {
|
|
"gpu": "available",
|
|
"checkpoints": "accessible",
|
|
"database": "connected",
|
|
"gpu_memory": "available"
|
|
},
|
|
"training_active": training_active
|
|
})))
|
|
} else {
|
|
Err(StatusCode::SERVICE_UNAVAILABLE)
|
|
}
|
|
}
|
|
|
|
Router::new()
|
|
.route("/health", get(health_handler))
|
|
.route("/ready", get(ready_handler))
|
|
.route("/health/deep", get(deep_health_handler))
|
|
.with_state(state)
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ml_health_basic_healthy() {
|
|
let state = MockMlHealthState::new();
|
|
let app = create_ml_health_router(state.clone());
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/health")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ml_health_unhealthy() {
|
|
let state = MockMlHealthState::new();
|
|
state.set_healthy(false).await;
|
|
let app = create_ml_health_router(state.clone());
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/health")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ml_readiness_check() {
|
|
let state = MockMlHealthState::new();
|
|
let app = create_ml_health_router(state.clone());
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/ready")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ml_service_startup() {
|
|
let state = MockMlHealthState::new();
|
|
state.set_ready(false).await;
|
|
|
|
// Initially not ready
|
|
assert!(!state.is_ready().await);
|
|
assert!(state.is_healthy().await);
|
|
|
|
// Simulate GPU initialization
|
|
tokio::time::sleep(Duration::from_millis(100)).await;
|
|
state.set_ready(true).await;
|
|
|
|
assert!(state.is_ready().await);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ml_gpu_unavailable() {
|
|
let state = MockMlHealthState::new();
|
|
state.set_gpu_available(false).await;
|
|
let app = create_ml_health_router(state.clone());
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/health/deep")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ml_checkpoints_inaccessible() {
|
|
let state = MockMlHealthState::new();
|
|
state.set_checkpoints_accessible(false).await;
|
|
let app = create_ml_health_router(state.clone());
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/health/deep")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ml_database_disconnection() {
|
|
let state = MockMlHealthState::new();
|
|
state.set_database_connected(false).await;
|
|
let app = create_ml_health_router(state.clone());
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/health/deep")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ml_gpu_memory_exhausted() {
|
|
let state = MockMlHealthState::new();
|
|
state.set_gpu_memory_available(false).await;
|
|
let app = create_ml_health_router(state.clone());
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/health/deep")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ml_health_during_training() {
|
|
let state = MockMlHealthState::new();
|
|
state.set_training_active(true).await;
|
|
|
|
let app = create_ml_health_router(state.clone());
|
|
|
|
// Service should still be healthy during training
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/health")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ml_dependency_cascade_failure() {
|
|
let state = MockMlHealthState::new();
|
|
|
|
// All dependencies fail
|
|
state.set_gpu_available(false).await;
|
|
state.set_checkpoints_accessible(false).await;
|
|
state.set_database_connected(false).await;
|
|
|
|
let app = create_ml_health_router(state.clone());
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/health/deep")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ml_health_check_latency() {
|
|
let state = MockMlHealthState::new();
|
|
let app = create_ml_health_router(state.clone());
|
|
|
|
let start = Instant::now();
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/health")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
let latency = start.elapsed();
|
|
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
assert!(
|
|
latency < Duration::from_millis(100),
|
|
"Health check latency: {:?}",
|
|
latency
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ml_concurrent_health_checks() {
|
|
let state = MockMlHealthState::new();
|
|
|
|
let mut handles = vec![];
|
|
for _ in 0..100 {
|
|
let state_clone = state.clone();
|
|
let handle = tokio::spawn(async move {
|
|
let app = create_ml_health_router(state_clone);
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/health")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
});
|
|
handles.push(handle);
|
|
}
|
|
|
|
for handle in handles {
|
|
handle.await.unwrap();
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ml_health_during_shutdown() {
|
|
let state = MockMlHealthState::new();
|
|
|
|
// Graceful shutdown
|
|
state.set_ready(false).await;
|
|
state.set_healthy(true).await;
|
|
|
|
let app = create_ml_health_router(state.clone());
|
|
|
|
// Ready check fails
|
|
let response = app
|
|
.clone()
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/ready")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
|
|
|
// Health check passes
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/health")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ml_rapid_health_checks() {
|
|
let state = MockMlHealthState::new();
|
|
|
|
let start = Instant::now();
|
|
for _ in 0..500 {
|
|
let app = create_ml_health_router(state.clone());
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/health")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
}
|
|
let duration = start.elapsed();
|
|
|
|
assert!(
|
|
duration < Duration::from_secs(1),
|
|
"500 health checks took: {:?}",
|
|
duration
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ml_deep_vs_shallow_health() {
|
|
let state = MockMlHealthState::new();
|
|
let app = create_ml_health_router(state.clone());
|
|
|
|
// Shallow health
|
|
let start = Instant::now();
|
|
let response = app
|
|
.clone()
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/health")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
let shallow_latency = start.elapsed();
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
// Deep health
|
|
let start = Instant::now();
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/health/deep")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
let deep_latency = start.elapsed();
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
assert!(shallow_latency < Duration::from_millis(50));
|
|
assert!(deep_latency < Duration::from_millis(100));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ml_partial_availability() {
|
|
let state = MockMlHealthState::new();
|
|
|
|
// GPU down, checkpoints up
|
|
state.set_gpu_available(false).await;
|
|
state.set_checkpoints_accessible(true).await;
|
|
let app = create_ml_health_router(state.clone());
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/health/deep")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ml_recovery_after_failure() {
|
|
let state = MockMlHealthState::new();
|
|
|
|
// Service fails
|
|
state.set_healthy(false).await;
|
|
let app = create_ml_health_router(state.clone());
|
|
|
|
let response = app
|
|
.clone()
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/health")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
|
|
|
// Service recovers
|
|
state.set_healthy(true).await;
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/health")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ml_health_json_format() {
|
|
let state = MockMlHealthState::new();
|
|
let app = create_ml_health_router(state.clone());
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/health")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
let content_type = response.headers().get("content-type");
|
|
assert!(content_type.is_some());
|
|
let content_type_str = content_type.unwrap().to_str().expect("INVARIANT: Path should be valid UTF-8");
|
|
assert!(content_type_str.contains("application/json"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ml_gpu_recovery_scenario() {
|
|
let state = MockMlHealthState::new();
|
|
|
|
// GPU fails
|
|
state.set_gpu_available(false).await;
|
|
let app = create_ml_health_router(state.clone());
|
|
|
|
let response = app
|
|
.clone()
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/health/deep")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
|
|
|
// GPU recovers
|
|
state.set_gpu_available(true).await;
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/health/deep")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
}
|