diff --git a/lite_bootstrap/bootstrappers/faststream_bootstrapper.py b/lite_bootstrap/bootstrappers/faststream_bootstrapper.py index cb8dc40..e7e7dc0 100644 --- a/lite_bootstrap/bootstrappers/faststream_bootstrapper.py +++ b/lite_bootstrap/bootstrappers/faststream_bootstrapper.py @@ -147,7 +147,6 @@ class FastStreamHealthChecksInstrument(HealthChecksInstrument[FastStreamConfig]) def bootstrap(self) -> None: config = self.bootstrap_config - @handle_get async def check_health(_: object) -> "AsgiResponse": health_status = await self.define_health_status() return AsgiResponse( @@ -159,6 +158,7 @@ async def check_health(_: object) -> "AsgiResponse": if config.opentelemetry_generate_health_check_spans and import_checker.is_opentelemetry_installed: check_health = tracer.start_as_current_span(f"GET {config.health_checks_path}")(check_health) + handler = handle_get(check_health) for path in _route_path_variants(config.health_checks_path): if ( path != config.health_checks_path @@ -166,7 +166,7 @@ async def check_health(_: object) -> "AsgiResponse": and FastStreamPrometheusInstrument.dependencies_installed() ): continue - config.application.mount(path, check_health) + config.application.mount(path, handler) async def _define_framework_health_status(self) -> bool: config = self.bootstrap_config diff --git a/tests/test_faststream_bootstrap.py b/tests/test_faststream_bootstrap.py index e5b89cd..1848259 100644 --- a/tests/test_faststream_bootstrap.py +++ b/tests/test_faststream_bootstrap.py @@ -23,6 +23,7 @@ from opentelemetry.metrics import MeterProvider, get_meter_provider from opentelemetry.sdk import resources from opentelemetry.sdk.metrics import MeterProvider as SDKMeterProvider +from opentelemetry.sdk.trace import ReadableSpan from opentelemetry.sdk.trace import TracerProvider as SDKTracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter @@ -309,7 +310,7 @@ def test_faststream_bootstrap_without_prometheus_client(broker: RedisBroker) -> bootstrapper.bootstrap() -def _span_names_for_one_health_check(config: FastStreamConfig) -> list[str]: +def _spans_for_one_health_check(config: FastStreamConfig) -> tuple[ReadableSpan, ...]: bootstrapper = FastStreamBootstrapper(bootstrap_config=config) # The module's tracer stays bound to the first provider a test installed; this one resolves to this test's. with patch.object(faststream_bootstrapper, "tracer", get_tracer(__name__)): @@ -322,17 +323,37 @@ def _span_names_for_one_health_check(config: FastStreamConfig) -> list[str]: TestClient(app=application).get(config.health_checks_path) finally: bootstrapper.teardown() - return [span.name for span in exporter.get_finished_spans()] + return exporter.get_finished_spans() def test_faststream_health_check_records_no_span_by_default() -> None: - assert _span_names_for_one_health_check(build_faststream_config()) == [] + assert not _spans_for_one_health_check(build_faststream_config()) -def test_faststream_health_check_span_is_recorded_when_turned_on() -> None: +def test_faststream_health_check_span_parents_async_probe_span() -> None: + async def probe() -> bool: + with get_tracer(__name__).start_as_current_span("dependency probe"): + return True + + config = dataclasses.replace( + build_faststream_config(), + opentelemetry_generate_health_check_spans=True, + health_checks_probes={"dependency": probe}, + ) + + spans = {span.name: span for span in _spans_for_one_health_check(config)} + health_span = spans[f"GET {config.health_checks_path}"] + assert health_span.context is not None + assert spans["dependency probe"].parent == health_span.context + + +def test_faststream_traced_health_check_mounts_http_handler() -> None: config = dataclasses.replace(build_faststream_config(), opentelemetry_generate_health_check_spans=True) - assert _span_names_for_one_health_check(config) == [f"GET {config.health_checks_path}"] + _spans_for_one_health_check(config) + + routes = dict(config.application.routes) + assert isinstance(routes[config.health_checks_path], faststream.asgi.HttpHandler) def test_faststream_bootstrap_without_opentelemetry(broker: RedisBroker) -> None: