axum_api_kit/
router.rs

1use std::future::Future;
2
3use axum::{routing::get, Router};
4
5use crate::HealthResponse;
6
7/// Liveness probe handler: always returns [`HealthResponse::ok`].
8pub async fn liveness() -> HealthResponse {
9    HealthResponse::ok()
10}
11
12/// Build a [`Router`] exposing `/healthz` (liveness) and `/readyz` (readiness) probes.
13///
14/// `/healthz` always reports `ok` (the process is running). `/readyz` runs `readiness` on
15/// each request and returns whatever [`HealthResponse`] it produces, so a failing dependency
16/// can report `unhealthy` (HTTP 503). For a service with no dependencies to check, pass
17/// `|| async { HealthResponse::ok() }`.
18///
19/// The router is generic over the state type `S`, so it can be merged into a stateful app
20/// with `app.merge(health_routes(...))`. Capture whatever the readiness check needs (a
21/// database pool, etc.) in the closure.
22///
23/// Requires the `router` feature.
24///
25/// # Example
26///
27/// ```rust,no_run
28/// use axum::Router;
29/// use axum_api_kit::{health_routes, HealthResponse};
30///
31/// let app: Router = Router::new().merge(health_routes(|| async {
32///     // probe your dependencies here...
33///     HealthResponse::ok()
34/// }));
35/// ```
36pub fn health_routes<S, F, Fut>(readiness: F) -> Router<S>
37where
38    S: Clone + Send + Sync + 'static,
39    F: Fn() -> Fut + Clone + Send + Sync + 'static,
40    Fut: Future<Output = HealthResponse> + Send + 'static,
41{
42    // Wrap in a concrete zero-arg closure so axum can resolve the `Handler` marker; the
43    // returned future is the generic `Fut`, whose `Send` bound is explicit (an opaque async
44    // block here would not be provably `Send` in this generic context).
45    #[allow(clippy::redundant_closure)]
46    let readyz = move || readiness();
47
48    Router::new()
49        .route("/healthz", get(liveness))
50        .route("/readyz", get(readyz))
51}
52
53#[cfg(test)]
54mod tests {
55    use super::*;
56    use axum::{body::Body, http::Request, http::StatusCode};
57    use tower::ServiceExt;
58
59    async fn call(app: Router, path: &str) -> (StatusCode, serde_json::Value) {
60        let res = app
61            .oneshot(Request::builder().uri(path).body(Body::empty()).unwrap())
62            .await
63            .unwrap();
64        let status = res.status();
65        let bytes = axum::body::to_bytes(res.into_body(), usize::MAX)
66            .await
67            .unwrap();
68        let body = serde_json::from_slice(&bytes).unwrap_or(serde_json::Value::Null);
69        (status, body)
70    }
71
72    #[tokio::test]
73    async fn healthz_is_always_ok() {
74        let app: Router = health_routes(|| async { HealthResponse::unhealthy() });
75        let (status, body) = call(app, "/healthz").await;
76        assert_eq!(status, StatusCode::OK);
77        assert_eq!(body["status"], "ok");
78    }
79
80    #[tokio::test]
81    async fn readyz_reports_ready() {
82        let app: Router = health_routes(|| async { HealthResponse::ok() });
83        let (status, body) = call(app, "/readyz").await;
84        assert_eq!(status, StatusCode::OK);
85        assert_eq!(body["status"], "ok");
86    }
87
88    #[tokio::test]
89    async fn readyz_reports_unhealthy() {
90        let app: Router = health_routes(|| async { HealthResponse::unhealthy() });
91        let (status, body) = call(app, "/readyz").await;
92        assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
93        assert_eq!(body["status"], "unhealthy");
94    }
95
96    #[tokio::test]
97    async fn merges_into_stateful_app() {
98        #[derive(Clone)]
99        struct AppState;
100        let app: Router<AppState> =
101            Router::new().merge(health_routes(|| async { liveness().await }));
102        let app: Router = app.with_state(AppState);
103        let (status, _) = call(app, "/healthz").await;
104        assert_eq!(status, StatusCode::OK);
105    }
106}
axum_api_kit/router.rs

axum_api_kit/
router.rs
