use axum::{
http::StatusCode,
response::{IntoResponse, Response},
Json,
};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::time::SystemTime;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum HealthStatus {
Healthy,
Degraded,
Unhealthy,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ComponentHealth {
pub status: HealthStatus,
#[serde(skip_serializing_if = "Option::is_none")]
pub message: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub response_time_ms: Option<u64>,
}
impl ComponentHealth {
#[must_use]
pub const fn healthy() -> Self {
Self {
status: HealthStatus::Healthy,
message: None,
response_time_ms: None,
}
}
#[must_use]
pub fn healthy_with_message(message: impl Into<String>) -> Self {
Self {
status: HealthStatus::Healthy,
message: Some(message.into()),
response_time_ms: None,
}
}
#[must_use]
pub fn degraded(message: impl Into<String>) -> Self {
Self {
status: HealthStatus::Degraded,
message: Some(message.into()),
response_time_ms: None,
}
}
#[must_use]
pub fn unhealthy(message: impl Into<String>) -> Self {
Self {
status: HealthStatus::Unhealthy,
message: Some(message.into()),
response_time_ms: None,
}
}
#[must_use]
pub const fn with_response_time(mut self, ms: u64) -> Self {
self.response_time_ms = Some(ms);
self
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HealthCheckResponse {
pub status: HealthStatus,
pub version: String,
pub timestamp: u64,
pub components: HashMap<String, ComponentHealth>,
}
impl HealthCheckResponse {
#[must_use]
pub fn new(version: impl Into<String>) -> Self {
Self {
status: HealthStatus::Healthy,
version: version.into(),
timestamp: SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.map_or(0, |d| d.as_secs()),
components: HashMap::new(),
}
}
pub fn add_component(&mut self, name: impl Into<String>, health: ComponentHealth) {
self.components.insert(name.into(), health);
self.recalculate_status();
}
fn recalculate_status(&mut self) {
if self.components.values().any(|c| c.status == HealthStatus::Unhealthy) {
self.status = HealthStatus::Unhealthy;
} else if self.components.values().any(|c| c.status == HealthStatus::Degraded) {
self.status = HealthStatus::Degraded;
} else {
self.status = HealthStatus::Healthy;
}
}
#[must_use]
pub const fn status_code(&self) -> StatusCode {
match self.status {
HealthStatus::Healthy | HealthStatus::Degraded => StatusCode::OK, HealthStatus::Unhealthy => StatusCode::SERVICE_UNAVAILABLE,
}
}
}
impl IntoResponse for HealthCheckResponse {
fn into_response(self) -> Response {
let status = self.status_code();
(status, Json(self)).into_response()
}
}
#[allow(clippy::unused_async)]
pub async fn liveness() -> impl IntoResponse {
(StatusCode::OK, "OK")
}
#[allow(clippy::unused_async)]
pub async fn readiness() -> impl IntoResponse {
let mut response = HealthCheckResponse::new(env!("CARGO_PKG_VERSION"));
response.add_component("application", ComponentHealth::healthy());
response
}
#[allow(clippy::unused_async)]
pub async fn health_check() -> impl IntoResponse {
let mut response = HealthCheckResponse::new(env!("CARGO_PKG_VERSION"));
response.add_component("application", ComponentHealth::healthy());
response
}
#[allow(clippy::unused_async)]
pub async fn health_check_with_state<S: Send + Sync>(_state: &S) -> HealthCheckResponse {
let mut response = HealthCheckResponse::new(env!("CARGO_PKG_VERSION"));
response.add_component("application", ComponentHealth::healthy());
response
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_component_health_healthy() {
let health = ComponentHealth::healthy();
assert_eq!(health.status, HealthStatus::Healthy);
assert!(health.message.is_none());
assert!(health.response_time_ms.is_none());
}
#[test]
fn test_component_health_with_message() {
let health = ComponentHealth::healthy_with_message("All good");
assert_eq!(health.status, HealthStatus::Healthy);
assert_eq!(health.message, Some("All good".to_string()));
}
#[test]
fn test_component_health_degraded() {
let health = ComponentHealth::degraded("High latency");
assert_eq!(health.status, HealthStatus::Degraded);
assert_eq!(health.message, Some("High latency".to_string()));
}
#[test]
fn test_component_health_unhealthy() {
let health = ComponentHealth::unhealthy("Connection failed");
assert_eq!(health.status, HealthStatus::Unhealthy);
assert_eq!(health.message, Some("Connection failed".to_string()));
}
#[test]
fn test_component_health_with_response_time() {
let health = ComponentHealth::healthy().with_response_time(150);
assert_eq!(health.response_time_ms, Some(150));
}
#[test]
fn test_health_check_response_new() {
let response = HealthCheckResponse::new("1.0.0");
assert_eq!(response.status, HealthStatus::Healthy);
assert_eq!(response.version, "1.0.0");
assert!(response.components.is_empty());
}
#[test]
fn test_health_check_response_add_component() {
let mut response = HealthCheckResponse::new("1.0.0");
response.add_component("database", ComponentHealth::healthy());
assert_eq!(response.components.len(), 1);
assert_eq!(response.status, HealthStatus::Healthy);
}
#[test]
fn test_health_check_response_degraded_status() {
let mut response = HealthCheckResponse::new("1.0.0");
response.add_component("app", ComponentHealth::healthy());
response.add_component("cache", ComponentHealth::degraded("High latency"));
assert_eq!(response.status, HealthStatus::Degraded);
assert_eq!(response.status_code(), StatusCode::OK);
}
#[test]
fn test_health_check_response_unhealthy_status() {
let mut response = HealthCheckResponse::new("1.0.0");
response.add_component("app", ComponentHealth::healthy());
response.add_component("database", ComponentHealth::unhealthy("Connection failed"));
assert_eq!(response.status, HealthStatus::Unhealthy);
assert_eq!(response.status_code(), StatusCode::SERVICE_UNAVAILABLE);
}
#[tokio::test]
async fn test_liveness_handler() {
let response = liveness().await.into_response();
assert_eq!(response.status(), StatusCode::OK);
}
#[tokio::test]
async fn test_readiness_handler() {
let response = readiness().await.into_response();
assert_eq!(response.status(), StatusCode::OK);
}
#[tokio::test]
async fn test_health_check_handler() {
let response = health_check().await.into_response();
assert_eq!(response.status(), StatusCode::OK);
}
}