diff --git a/microbootstrap/instruments/opentelemetry_instrument.py b/microbootstrap/instruments/opentelemetry_instrument.py index 1d40301..ca93f3c 100644 --- a/microbootstrap/instruments/opentelemetry_instrument.py +++ b/microbootstrap/instruments/opentelemetry_instrument.py @@ -6,7 +6,6 @@ import pydantic import structlog -from faststream._internal.types import BrokerMiddleware from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter from opentelemetry.instrumentation.dependencies import DependencyConflictError from opentelemetry.instrumentation.environment_variables import OTEL_PYTHON_DISABLED_INSTRUMENTATIONS @@ -70,7 +69,7 @@ class OpentelemetryConfig(BaseInstrumentConfig): @typing.runtime_checkable -class FastStreamTelemetryMiddlewareProtocol(BrokerMiddleware, typing.Protocol): +class FastStreamTelemetryMiddlewareProtocol(typing.Protocol): def __init__( self, *, @@ -79,6 +78,14 @@ def __init__( meter: Meter | None = None, ) -> None: ... + def __call__( + self, + msg: typing.Any, # noqa: ANN401 + /, + *, + context: typing.Any, # noqa: ANN401 + ) -> typing.Any: ... # noqa: ANN401 + class FastStreamOpentelemetryConfig(OpentelemetryConfig): opentelemetry_middleware_cls: type[FastStreamTelemetryMiddlewareProtocol] | None = None diff --git a/microbootstrap/instruments/prometheus_instrument.py b/microbootstrap/instruments/prometheus_instrument.py index 9751d72..9936657 100644 --- a/microbootstrap/instruments/prometheus_instrument.py +++ b/microbootstrap/instruments/prometheus_instrument.py @@ -2,7 +2,6 @@ import typing import pydantic -from faststream._internal.types import BrokerMiddleware from microbootstrap.helpers import is_valid_path from microbootstrap.instruments.base import BaseInstrumentConfig, Instrument @@ -34,7 +33,7 @@ class FastApiPrometheusConfig(BasePrometheusConfig): @typing.runtime_checkable -class FastStreamPrometheusMiddlewareProtocol(BrokerMiddleware, typing.Protocol): +class FastStreamPrometheusMiddlewareProtocol(typing.Protocol): def __init__( self, *, @@ -44,6 +43,13 @@ def __init__( received_messages_size_buckets: typing.Sequence[float] | None = None, custom_labels: dict[str, str | typing.Callable[[typing.Any], str]] | None = None, ) -> None: ... + def __call__( + self, + msg: typing.Any, # noqa: ANN401 + /, + *, + context: typing.Any, # noqa: ANN401 + ) -> typing.Any: ... # noqa: ANN401 class FastStreamPrometheusConfig(BasePrometheusConfig): diff --git a/tests/instruments/test_prometheus.py b/tests/instruments/test_prometheus.py index 1769027..a539052 100644 --- a/tests/instruments/test_prometheus.py +++ b/tests/instruments/test_prometheus.py @@ -145,7 +145,7 @@ async def test_faststream_prometheus_custom_labels( expected_label_keys: set[str], ) -> None: minimal_faststream_prometheus_config.prometheus_custom_labels = custom_labels - minimal_faststream_prometheus_config.prometheus_middleware_cls = RedisPrometheusMiddleware # type: ignore[assignment] + minimal_faststream_prometheus_config.prometheus_middleware_cls = RedisPrometheusMiddleware broker: typing.Final = RedisBroker() (