From d901795e6b6b3fe02d2869699d5a9ca5df2660a6 Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Thu, 22 Oct 2026 08:33:42 +0300 Subject: [PATCH] feat: run user-supplied health probes on every health check route (#294) --- CONTEXT.md | 7 +++ docs/introduction/configuration.md | 29 +++++++++ .../bootstrappers/fastapi_bootstrapper.py | 7 ++- .../bootstrappers/fastmcp_bootstrapper.py | 7 ++- .../bootstrappers/faststream_bootstrapper.py | 15 ++--- .../bootstrappers/litestar_bootstrapper.py | 9 ++- .../instruments/healthchecks_instrument.py | 36 ++++++++++- tests/conftest.py | 36 +++++++++++ .../test_healthchecks_instrument.py | 29 +++++++++ tests/test_fastapi_bootstrap.py | 34 ++++++++++- tests/test_fastmcp_bootstrap.py | 29 +++++++++ tests/test_faststream_bootstrap.py | 60 ++++++++++++++++++- tests/test_litestar_bootstrap.py | 27 +++++++++ 13 files changed, 306 insertions(+), 19 deletions(-) diff --git a/CONTEXT.md b/CONTEXT.md index 789d992..ca60147 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -46,6 +46,13 @@ signals: not configured is silent and lands in `skipped_instruments`, so it show raises `InstrumentDependencyMissingWarning` plus a `logger.warning` and lands in neither list. The quiet skip is the one you have to go looking for; the loud one finds you. +**Health probe**: +One check of something the service depends on, run each time the health check route is requested. A +single failing probe makes the route report the service unhealthy. The health check is the route; +health probes are what it runs. +_Avoid_: health check, for this sense; ping, which is FastStream's broker method and only one +possible probe. + **Free**: Without a web framework: `FreeBootstrapper`, `FreeConfig`, the `free-all` extra. It has no application, so nothing to attach teardown to and no double-bootstrap guard. diff --git a/docs/introduction/configuration.md b/docs/introduction/configuration.md index fba77e5..7927e92 100644 --- a/docs/introduction/configuration.md +++ b/docs/introduction/configuration.md @@ -391,6 +391,35 @@ Additional params: - `health_checks_path` - `health_checks_include_in_schema` +- `health_checks_probes`: named async callables, each returning `bool`, that the route runs on every request + +The health check route on Litestar, FastAPI, FastMCP and FastStream runs every probe concurrently. On FastStream the probes run alongside the broker pings. The route returns 200 with `health_status: true` when everything passes. If any probe returns `False` or raises, the route returns 500 with the same JSON body and `health_status: false`. Each failed probe logs one warning with its name; the response body does not say which probe failed. + +```python +import functools + +import sqlalchemy + +from lite_bootstrap import LitestarConfig + + +async def database_probe() -> bool: + async with engine.connect() as connection: + await connection.execute(sqlalchemy.text("SELECT 1")) + return True + + +config = LitestarConfig( + health_checks_probes={ + "database": database_probe, + "outbox-broker": functools.partial(outbox_broker.ping, timeout=5), + }, +) +``` + +lite-bootstrap puts no timeout on probes, so a probe that hangs makes the route hang too. Give each probe its own timeout, as `functools.partial(broker.ping, timeout=5)` does. + +Be careful with probes on dependencies outside the service, such as a database, when the route also serves a Kubernetes liveness probe. During the database outage every pod fails its liveness check and gets restarted, and the restarts add load once the database comes back. Readiness probes are the usual place for dependency checks. ## Skipped instruments diff --git a/lite_bootstrap/bootstrappers/fastapi_bootstrapper.py b/lite_bootstrap/bootstrappers/fastapi_bootstrapper.py index 0701f7f..fb38f11 100644 --- a/lite_bootstrap/bootstrappers/fastapi_bootstrapper.py +++ b/lite_bootstrap/bootstrappers/fastapi_bootstrapper.py @@ -178,8 +178,11 @@ def build_fastapi_health_check_router(self) -> "fastapi.APIRouter": ) @fastapi_router.get(self.bootstrap_config.health_checks_path) - async def health_check_handler() -> HealthCheckTypedDict: - return self.render_health_check_data() + async def health_check_handler(response: fastapi.Response) -> HealthCheckTypedDict: + health_status = await self.define_health_status() + if not health_status: + response.status_code = fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR + return self.render_health_check_data(health_status=health_status) return fastapi_router diff --git a/lite_bootstrap/bootstrappers/fastmcp_bootstrapper.py b/lite_bootstrap/bootstrappers/fastmcp_bootstrapper.py index 9c2dae4..a71085d 100644 --- a/lite_bootstrap/bootstrappers/fastmcp_bootstrapper.py +++ b/lite_bootstrap/bootstrappers/fastmcp_bootstrapper.py @@ -22,6 +22,7 @@ from fastmcp.server.http import StarletteWithLifespan from fastmcp.server.middleware import Middleware, MiddlewareContext from fastmcp.server.providers import Provider + from starlette import status from starlette.requests import Request from starlette.responses import JSONResponse, Response @@ -210,7 +211,11 @@ def bootstrap(self) -> None: include_in_schema=config.health_checks_include_in_schema, ) async def health_check_handler(_: "Request") -> "JSONResponse": - return JSONResponse(dict(self.render_health_check_data())) + health_status = await self.define_health_status() + return JSONResponse( + dict(self.render_health_check_data(health_status=health_status)), + status_code=status.HTTP_200_OK if health_status else status.HTTP_500_INTERNAL_SERVER_ERROR, + ) @dataclasses.dataclass(kw_only=True) diff --git a/lite_bootstrap/bootstrappers/faststream_bootstrapper.py b/lite_bootstrap/bootstrappers/faststream_bootstrapper.py index 5049e1e..0fc393b 100644 --- a/lite_bootstrap/bootstrappers/faststream_bootstrapper.py +++ b/lite_bootstrap/bootstrappers/faststream_bootstrapper.py @@ -88,14 +88,11 @@ def bootstrap(self) -> None: @handle_get async def check_health(_: object) -> "AsgiResponse": - return ( - AsgiResponse( - json.dumps(self.render_health_check_data()).encode(), - 200, - headers={"content-type": "application/json"}, - ) - if await self._define_health_status() - else AsgiResponse(b"Service is unhealthy", 500, headers={"content-type": "text/plain"}) + health_status = await self.define_health_status() + return AsgiResponse( + json.dumps(self.render_health_check_data(health_status=health_status)).encode(), + 200 if health_status else 500, + headers={"content-type": "application/json"}, ) if config.opentelemetry_generate_health_check_spans and import_checker.is_opentelemetry_installed: @@ -103,7 +100,7 @@ async def check_health(_: object) -> "AsgiResponse": config.application.mount(config.health_checks_path, check_health) - async def _define_health_status(self) -> bool: + async def _define_framework_health_status(self) -> bool: config = self.bootstrap_config if not config.application or not config.application.brokers: return False diff --git a/lite_bootstrap/bootstrappers/litestar_bootstrapper.py b/lite_bootstrap/bootstrappers/litestar_bootstrapper.py index 839c605..623f031 100644 --- a/lite_bootstrap/bootstrappers/litestar_bootstrapper.py +++ b/lite_bootstrap/bootstrappers/litestar_bootstrapper.py @@ -28,6 +28,7 @@ if import_checker.is_litestar_installed: import litestar + from litestar import status_codes from litestar.config.app import AppConfig from litestar.config.cors import CORSConfig from litestar.logging.config import StructLoggingConfig @@ -149,8 +150,12 @@ def bootstrap(self) -> None: class LitestarHealthChecksInstrument(HealthChecksInstrument[LitestarConfig]): def build_litestar_health_check_router(self) -> "litestar.Router": @litestar.get(media_type=litestar.MediaType.JSON) - async def health_check_handler() -> HealthCheckTypedDict: - return self.render_health_check_data() + async def health_check_handler() -> litestar.Response[HealthCheckTypedDict]: + health_status = await self.define_health_status() + return litestar.Response( + self.render_health_check_data(health_status=health_status), + status_code=status_codes.HTTP_200_OK if health_status else status_codes.HTTP_500_INTERNAL_SERVER_ERROR, + ) return litestar.Router( path=self.bootstrap_config.health_checks_path, diff --git a/lite_bootstrap/instruments/healthchecks_instrument.py b/lite_bootstrap/instruments/healthchecks_instrument.py index 035b91c..b8837ac 100644 --- a/lite_bootstrap/instruments/healthchecks_instrument.py +++ b/lite_bootstrap/instruments/healthchecks_instrument.py @@ -1,11 +1,19 @@ +import asyncio import dataclasses +import logging import typing +from collections.abc import Awaitable, Callable, Mapping import typing_extensions from lite_bootstrap.instruments.base import BaseConfig, BaseInstrument +logger = logging.getLogger(__name__) + +HealthProbe = Callable[[], Awaitable[bool]] + + class HealthCheckTypedDict(typing_extensions.TypedDict, total=False): service_version: str | None service_name: str | None @@ -17,6 +25,7 @@ class HealthChecksConfig(BaseConfig): health_checks_enabled: bool = True health_checks_path: str = "/health/" health_checks_include_in_schema: bool = False + health_checks_probes: Mapping[str, HealthProbe] = dataclasses.field(default_factory=dict) @property def health_check_data(self) -> HealthCheckTypedDict: @@ -38,5 +47,28 @@ class HealthChecksInstrument(BaseInstrument[HealthChecksConfigT]): def is_configured(cls, bootstrap_config: "HealthChecksConfig") -> bool: return bootstrap_config.health_checks_enabled - def render_health_check_data(self) -> HealthCheckTypedDict: - return self.bootstrap_config.health_check_data + def render_health_check_data(self, *, health_status: bool = True) -> HealthCheckTypedDict: + health_check_data = self.bootstrap_config.health_check_data + health_check_data["health_status"] = health_status + return health_check_data + + async def define_health_status(self) -> bool: + results = await asyncio.gather( + self._define_framework_health_status(), + *(self._run_probe(name, probe) for name, probe in self.bootstrap_config.health_checks_probes.items()), + ) + return all(results) + + async def _define_framework_health_status(self) -> bool: + return True + + @staticmethod + async def _run_probe(name: str, probe: HealthProbe) -> bool: + try: + passed = await probe() + except Exception: + logger.warning("Health probe %s raised", name, exc_info=True) + return False + if not passed: + logger.warning("Health probe %s failed", name) + return bool(passed) diff --git a/tests/conftest.py b/tests/conftest.py index 762947e..4682a26 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -4,6 +4,7 @@ import warnings from collections.abc import Generator from importlib import reload +from unittest.mock import MagicMock, patch import opentelemetry.metrics._internal import opentelemetry.trace @@ -16,6 +17,7 @@ from lite_bootstrap import import_checker from lite_bootstrap.exceptions import InstrumentSkippedWarning +from lite_bootstrap.instruments import healthchecks_instrument def pytest_configure() -> None: @@ -36,6 +38,40 @@ def _uninstrument(self, **kwargs: typing.Mapping[str, typing.Any]) -> None: pass +class ProbeError(RuntimeError): + pass + + +async def passing_probe() -> bool: + return True + + +async def failing_probe() -> bool: + return False + + +async def raising_probe() -> bool: + message = "dependency unreachable" + raise ProbeError(message) + + +HEALTH_PROBE_CASES: typing.Final = [ + pytest.param({"db": passing_probe, "cache": passing_probe}, 200, True, [], id="all-pass"), + pytest.param({"db": passing_probe, "cache": failing_probe}, 500, False, ["cache"], id="one-returns-false"), + pytest.param({"db": passing_probe, "cache": raising_probe}, 500, False, ["cache"], id="one-raises"), +] + + +@contextlib.contextmanager +def capture_probe_warnings() -> Generator[MagicMock]: + with patch.object(healthchecks_instrument, "logger") as mock_logger: + yield mock_logger.warning + + +def failed_probe_names(warning_mock: MagicMock) -> list[str]: + return [call.args[1] for call in warning_mock.call_args_list] + + class SentryTestTransport(sentry_sdk.Transport): def __init__(self, *args: typing.Any, **kwargs: typing.Any) -> None: # noqa: ANN401 super().__init__(*args, **kwargs) diff --git a/tests/instruments/test_healthchecks_instrument.py b/tests/instruments/test_healthchecks_instrument.py index ab438b9..178834b 100644 --- a/tests/instruments/test_healthchecks_instrument.py +++ b/tests/instruments/test_healthchecks_instrument.py @@ -1,3 +1,6 @@ +import asyncio +import time + from lite_bootstrap.instruments.healthchecks_instrument import HealthChecksConfig, HealthChecksInstrument @@ -39,7 +42,33 @@ def test_healthchecks_config_defaults() -> None: assert config.health_checks_enabled is True assert config.health_checks_path == "/health/" assert config.health_checks_include_in_schema is False + assert config.health_checks_probes == {} def test_healthchecks_dependencies_installed() -> None: assert HealthChecksInstrument.dependencies_installed() is True + + +async def test_healthchecks_probes_run_concurrently() -> None: + interval = 0.2 + + async def slow_probe() -> bool: + await asyncio.sleep(interval) + return True + + instrument = HealthChecksInstrument( + bootstrap_config=HealthChecksConfig(health_checks_probes={"first": slow_probe, "second": slow_probe}), + ) + started = time.perf_counter() + assert await instrument.define_health_status() + assert time.perf_counter() - started < interval * 1.5 + + +async def test_healthchecks_status_is_healthy_without_probes() -> None: + instrument = HealthChecksInstrument(bootstrap_config=HealthChecksConfig()) + assert await instrument.define_health_status() + + +def test_healthchecks_render_data_unhealthy() -> None: + instrument = HealthChecksInstrument(bootstrap_config=HealthChecksConfig()) + assert instrument.render_health_check_data(health_status=False)["health_status"] is False diff --git a/tests/test_fastapi_bootstrap.py b/tests/test_fastapi_bootstrap.py index dfda088..7760249 100644 --- a/tests/test_fastapi_bootstrap.py +++ b/tests/test_fastapi_bootstrap.py @@ -23,8 +23,17 @@ from lite_bootstrap import FastAPIBootstrapper, FastAPIConfig, import_checker from lite_bootstrap.bootstrappers import fastapi_bootstrapper from lite_bootstrap.exceptions import ConfigurationError, InstrumentDependencyMissingWarning +from lite_bootstrap.instruments.healthchecks_instrument import HealthProbe from lite_bootstrap.types import UNSET -from tests.conftest import CustomInstrumentor, SentryTestTransport, emulate_package_missing, warning_source_files +from tests.conftest import ( + HEALTH_PROBE_CASES, + CustomInstrumentor, + SentryTestTransport, + capture_probe_warnings, + emulate_package_missing, + failed_probe_names, + warning_source_files, +) logger = structlog.getLogger(__name__) @@ -501,3 +510,26 @@ def test_fastapi_otel_honors_excluded_urls_from_environment( span_names = _server_span_names_for(fastapi_config, ["/from-env", "/custom-healthy"]) assert span_names == ["GET /custom-healthy"] + + +@pytest.mark.parametrize(("probes", "expected_status", "expected_health_status", "expected_failed"), HEALTH_PROBE_CASES) +def test_fastapi_health_check_runs_probes( + fastapi_config: FastAPIConfig, + probes: dict[str, HealthProbe], + expected_status: int, + expected_health_status: bool, + expected_failed: list[str], +) -> None: + config = dataclasses.replace(fastapi_config, health_checks_probes=probes) + bootstrapper = FastAPIBootstrapper(bootstrap_config=config) + application = bootstrapper.bootstrap() + with capture_probe_warnings() as probe_warnings, TestClient(application) as test_client: + response = test_client.get(config.health_checks_path) + + assert response.status_code == expected_status + assert response.json() == { + "health_status": expected_health_status, + "service_name": "microservice", + "service_version": "2.0.0", + } + assert failed_probe_names(probe_warnings) == expected_failed diff --git a/tests/test_fastmcp_bootstrap.py b/tests/test_fastmcp_bootstrap.py index 45d4d78..783eace 100644 --- a/tests/test_fastmcp_bootstrap.py +++ b/tests/test_fastmcp_bootstrap.py @@ -30,9 +30,13 @@ _postprocess_http_apps, ) from lite_bootstrap.exceptions import ConfigurationError +from lite_bootstrap.instruments.healthchecks_instrument import HealthProbe from tests.conftest import ( + HEALTH_PROBE_CASES, + capture_probe_warnings, emulate_package_missing, emulate_package_missing_with_module_reload, + failed_probe_names, ) @@ -641,3 +645,28 @@ def test_fastmcp_otel_honors_excluded_urls_from_environment(monkeypatch: pytest. assert [span.name for span in _server_spans(exporter)] == [ f"GET {bootstrapper.bootstrap_config.health_checks_path}" ] + + +@pytest.mark.parametrize(("probes", "expected_status", "expected_health_status", "expected_failed"), HEALTH_PROBE_CASES) +def test_fastmcp_health_check_runs_probes( + probes: dict[str, HealthProbe], + expected_status: int, + expected_health_status: bool, + expected_failed: list[str], +) -> None: + config = _make_test_config(health_checks_probes=probes) + bootstrapper = FastMcpBootstrapper(bootstrap_config=config) + application = bootstrapper.bootstrap() + try: + with capture_probe_warnings() as probe_warnings, TestClient(application.http_app()) as test_client: + response = test_client.get(config.health_checks_path) + finally: + bootstrapper.teardown() + + assert response.status_code == expected_status + assert response.json() == { + "health_status": expected_health_status, + "service_name": "test-mcp", + "service_version": "1.2.3", + } + assert failed_probe_names(probe_warnings) == expected_failed diff --git a/tests/test_faststream_bootstrap.py b/tests/test_faststream_bootstrap.py index 25273cf..b87c52a 100644 --- a/tests/test_faststream_bootstrap.py +++ b/tests/test_faststream_bootstrap.py @@ -29,11 +29,16 @@ FastStreamPrometheusInstrument, ) from lite_bootstrap.exceptions import ConfigurationError +from lite_bootstrap.instruments.healthchecks_instrument import HealthProbe from tests.conftest import ( + HEALTH_PROBE_CASES, CustomInstrumentor, SentryTestTransport, + capture_probe_warnings, emulate_package_missing, emulate_package_missing_with_module_reload, + failed_probe_names, + passing_probe, ) @@ -122,7 +127,11 @@ async def test_faststream_bootstrap_health_check_wo_broker() -> None: response = test_client.get(bootstrap_config.health_checks_path) assert response.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR - assert response.text == "Service is unhealthy" + assert response.json() == { + "health_status": False, + "service_name": "microservice", + "service_version": "2.0.0", + } bootstrapper.teardown() @@ -243,7 +252,7 @@ async def test_faststream_health_check_fails_when_second_broker_is_down(broker: ): response = test_client.get(config.health_checks_path) assert response.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR - assert response.text == "Service is unhealthy" + assert response.json()["health_status"] is False finally: bootstrapper.teardown() @@ -455,3 +464,50 @@ def __init__( assert isinstance(get_meter_provider(), SDKMeterProvider) finally: bootstrapper.teardown() + + +@pytest.mark.parametrize(("probes", "expected_status", "expected_health_status", "expected_failed"), HEALTH_PROBE_CASES) +async def test_faststream_health_check_runs_probes( + broker: RedisBroker, + probes: dict[str, HealthProbe], + expected_status: int, + expected_health_status: bool, + expected_failed: list[str], +) -> None: + config = dataclasses.replace(build_faststream_config(broker=broker), health_checks_probes=probes) + bootstrapper = FastStreamBootstrapper(bootstrap_config=config) + application = bootstrapper.bootstrap() + try: + with ( + patch.object(broker, "ping", new=AsyncMock(return_value=True)), + capture_probe_warnings() as probe_warnings, + TestClient(app=application) as test_client, + ): + response = test_client.get(config.health_checks_path) + finally: + bootstrapper.teardown() + + assert response.status_code == expected_status + assert response.json() == { + "health_status": expected_health_status, + "service_name": "microservice", + "service_version": "2.0.0", + } + assert failed_probe_names(probe_warnings) == expected_failed + + +async def test_faststream_health_check_fails_on_broker_ping_despite_passing_probe(broker: RedisBroker) -> None: + config = dataclasses.replace(build_faststream_config(broker=broker), health_checks_probes={"db": passing_probe}) + bootstrapper = FastStreamBootstrapper(bootstrap_config=config) + application = bootstrapper.bootstrap() + try: + with ( + patch.object(broker, "ping", new=AsyncMock(return_value=False)), + TestClient(app=application) as test_client, + ): + response = test_client.get(config.health_checks_path) + finally: + bootstrapper.teardown() + + assert response.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR + assert response.json()["health_status"] is False diff --git a/tests/test_litestar_bootstrap.py b/tests/test_litestar_bootstrap.py index 81073b3..7cb5409 100644 --- a/tests/test_litestar_bootstrap.py +++ b/tests/test_litestar_bootstrap.py @@ -35,11 +35,15 @@ build_span_name, ) from lite_bootstrap.exceptions import ConfigurationError +from lite_bootstrap.instruments.healthchecks_instrument import HealthProbe from tests.conftest import ( + HEALTH_PROBE_CASES, CustomInstrumentor, SentryTestTransport, + capture_probe_warnings, emulate_package_missing, emulate_package_missing_with_module_reload, + failed_probe_names, ) @@ -697,3 +701,26 @@ def test_litestar_otel_is_skipped_without_litestars_opentelemetry_plugin(litesta assert "OpenTelemetryPlugin" not in { type(plugin).__name__ for plugin in litestar_config.application_config.plugins } + + +@pytest.mark.parametrize(("probes", "expected_status", "expected_health_status", "expected_failed"), HEALTH_PROBE_CASES) +def test_litestar_health_check_runs_probes( + litestar_config: LitestarConfig, + probes: dict[str, HealthProbe], + expected_status: int, + expected_health_status: bool, + expected_failed: list[str], +) -> None: + config = dataclasses.replace(litestar_config, health_checks_probes=probes) + bootstrapper = LitestarBootstrapper(bootstrap_config=config) + application = bootstrapper.bootstrap() + with capture_probe_warnings() as probe_warnings, TestClient(app=application) as test_client: + response = test_client.get(config.health_checks_path) + + assert response.status_code == expected_status + assert response.json() == { + "health_status": expected_health_status, + "service_name": "microservice", + "service_version": "2.0.0", + } + assert failed_probe_names(probe_warnings) == expected_failed