diff --git a/lite_bootstrap/bootstrappers/faststream_bootstrapper.py b/lite_bootstrap/bootstrappers/faststream_bootstrapper.py index a815f0e..db21a06 100644 --- a/lite_bootstrap/bootstrappers/faststream_bootstrapper.py +++ b/lite_bootstrap/bootstrappers/faststream_bootstrapper.py @@ -12,10 +12,9 @@ if import_checker.is_faststream_installed: - import faststream + from faststream._internal.broker import BrokerUsecase from faststream.asgi import AsgiFastStream, AsgiResponse from faststream.asgi import get as handle_get - from faststream.broker.core.usecase import BrokerUsecase if import_checker.is_prometheus_client_installed: import prometheus_client @@ -33,8 +32,8 @@ def __init__( tracer_provider: typing.Optional["TracerProvider"] = None, meter_provider: typing.Optional["MeterProvider"] = None, meter: typing.Optional["Meter"] = None, + include_messages_counters: bool = True, ) -> None: ... - def __call__(self, msg: typing.Any | None) -> "faststream.BaseMiddleware": ... # noqa: ANN401 @typing.runtime_checkable @@ -47,7 +46,6 @@ def __init__( metrics_prefix: str = "faststream", received_messages_size_buckets: typing.Sequence[float] | None = None, ) -> None: ... - def __call__(self, msg: typing.Any | None) -> "faststream.BaseMiddleware": ... # noqa: ANN401 @dataclasses.dataclass(kw_only=True, slots=True, frozen=True) @@ -103,8 +101,9 @@ def is_ready(self) -> bool: def bootstrap(self) -> None: if self.bootstrap_config.opentelemetry_middleware_cls and self.bootstrap_config.application.broker: + self.bootstrap_config.opentelemetry_middleware_cls(tracer_provider=get_tracer_provider()) self.bootstrap_config.application.broker.add_middleware( - self.bootstrap_config.opentelemetry_middleware_cls(tracer_provider=get_tracer_provider()) + self.bootstrap_config.opentelemetry_middleware_cls(tracer_provider=get_tracer_provider()) # type: ignore[arg-type] ) @@ -139,7 +138,7 @@ def bootstrap(self) -> None: ) if self.bootstrap_config.prometheus_middleware_cls and self.bootstrap_config.application.broker: self.bootstrap_config.application.broker.add_middleware( - self.bootstrap_config.prometheus_middleware_cls(registry=self.collector_registry) + self.bootstrap_config.prometheus_middleware_cls(registry=self.collector_registry) # type: ignore[arg-type] ) @@ -162,7 +161,7 @@ def is_ready(self) -> bool: def __init__(self, bootstrap_config: FastStreamConfig) -> None: super().__init__(bootstrap_config) if self.bootstrap_config.broker: - self.bootstrap_config.application.broker = self.bootstrap_config.broker + self.bootstrap_config.application.set_broker(self.bootstrap_config.broker) self.bootstrap_config.application.on_shutdown(self.teardown) def _prepare_application(self) -> "AsgiFastStream": diff --git a/tests/test_faststream_bootstrap.py b/tests/test_faststream_bootstrap.py index 0c49730..0883e7f 100644 --- a/tests/test_faststream_bootstrap.py +++ b/tests/test_faststream_bootstrap.py @@ -2,7 +2,7 @@ import pytest import structlog -from faststream.broker.core.usecase import BrokerUsecase +from faststream._internal.broker import BrokerUsecase from faststream.redis import RedisBroker, TestRedisBroker from faststream.redis.opentelemetry import RedisTelemetryMiddleware from faststream.redis.prometheus import RedisPrometheusMiddleware