diff --git a/lite_bootstrap/bootstrappers/faststream_bootstrapper.py b/lite_bootstrap/bootstrappers/faststream_bootstrapper.py index a794dcb..fc9dffb 100644 --- a/lite_bootstrap/bootstrappers/faststream_bootstrapper.py +++ b/lite_bootstrap/bootstrappers/faststream_bootstrapper.py @@ -68,6 +68,7 @@ class FastStreamConfig( opentelemetry_middleware_cls: type[FastStreamTelemetryMiddlewareProtocol] | None = None prometheus_middleware_cls: type[FastStreamPrometheusMiddlewareProtocol] | None = None faststream_log_level: int = logging.WARNING + faststream_health_check_broker_timeout: float = 5.0 @dataclasses.dataclass(kw_only=True, slots=True) @@ -101,7 +102,9 @@ async def _define_health_status(self) -> bool: if not self.bootstrap_config.application or not self.bootstrap_config.application.broker: return False - return await self.bootstrap_config.application.broker.ping(timeout=5) + return await self.bootstrap_config.application.broker.ping( + timeout=self.bootstrap_config.faststream_health_check_broker_timeout, + ) @dataclasses.dataclass(kw_only=True) diff --git a/tests/test_faststream_bootstrap.py b/tests/test_faststream_bootstrap.py index 2cf6211..b890db9 100644 --- a/tests/test_faststream_bootstrap.py +++ b/tests/test_faststream_bootstrap.py @@ -1,5 +1,7 @@ +import dataclasses import logging import typing +from unittest.mock import AsyncMock, patch import faststream.asgi import pytest @@ -159,3 +161,23 @@ def test_faststream_bootstrap_without_opentelemetry(broker: RedisBroker) -> None ): bootstrapper = FastStreamBootstrapper(bootstrap_config=bootstrap_config) bootstrapper.bootstrap() + + +async def test_faststream_health_check_uses_configured_broker_timeout(broker: RedisBroker) -> None: + expected_timeout = 12.5 + config = dataclasses.replace( + build_faststream_config(broker=broker), + faststream_health_check_broker_timeout=expected_timeout, + ) + bootstrapper = FastStreamBootstrapper(bootstrap_config=config) + application = bootstrapper.bootstrap() + try: + with ( + patch.object(broker, "ping", new=AsyncMock(return_value=True)) as mock_ping, + TestClient(app=application) as test_client, + ): + response = test_client.get(config.health_checks_path) + assert response.status_code == status.HTTP_200_OK + mock_ping.assert_called_once_with(timeout=expected_timeout) + finally: + bootstrapper.teardown()