From da2507b64bdcab39dab0ab18b4ceb3771a508955 Mon Sep 17 00:00:00 2001 From: Lancetnik Date: Wed, 9 Sep 2026 11:33:19 +0300 Subject: [PATCH 1/8] feat(testing): Kafka handlers and publishers assert on `key` and `partition` `Endpoint` reads the wrapper class it builds from `_call_wrapper_class`, and `ensure_call_wrapper` takes that class while still returning an already wrapped handler untouched, so the first decorator on a shared handler decides. Both Kafka packages set `KafkaHandlerCallWrapper` on their `LogicSubscriber` and `LogicPublisher`, and repeat the `__call__` overloads with the return type narrowed to it. `KafkaCallAssertions` in `_internal/kafka` declares the three Call assertions with `key=` and `partition=` and forwards them as Broker fields. `ExpectedCall` carries the fields with a reader `(name, message) -> value` that each package supplies from its `call_wrapper` module: attributes on an aiokafka record, method calls on a confluent message, a `SetupError` naming Kafka for a raw message of another broker. A batch refuses the two fields as it refuses `headers`. The in-memory confluent message is now a `confluent_kafka.Message` subclass, so the reader takes it for the client's own. Co-Authored-By: Claude Fable 5.1 --- faststream/_internal/endpoint/call_wrapper.py | 8 +- faststream/_internal/endpoint/usecase.py | 8 +- faststream/_internal/kafka/calls.py | 168 ++++++++++++++++++ faststream/_internal/testing/calls.py | 53 ++++-- faststream/confluent/call_wrapper.py | 35 ++++ faststream/confluent/publisher/usecase.py | 21 ++- faststream/confluent/subscriber/usecase.py | 68 ++++++- faststream/confluent/testing.py | 14 +- faststream/kafka/call_wrapper.py | 34 ++++ faststream/kafka/publisher/usecase.py | 21 ++- faststream/kafka/subscriber/usecase.py | 68 ++++++- tests/brokers/confluent/test_test_client.py | 35 ++++ tests/brokers/kafka/test_test_client.py | 54 ++++++ .../subscription/test_advanced_testing.py | 4 + tests/mypy/confluent.py | 25 +++ tests/mypy/kafka.py | 25 ++- 16 files changed, 608 insertions(+), 33 deletions(-) create mode 100644 faststream/_internal/kafka/calls.py create mode 100644 faststream/confluent/call_wrapper.py create mode 100644 faststream/kafka/call_wrapper.py diff --git a/faststream/_internal/endpoint/call_wrapper.py b/faststream/_internal/endpoint/call_wrapper.py index eabefaf1668..32eb1a1073b 100644 --- a/faststream/_internal/endpoint/call_wrapper.py +++ b/faststream/_internal/endpoint/call_wrapper.py @@ -5,6 +5,7 @@ Any, Generic, Optional, + cast, ) import anyio @@ -28,11 +29,16 @@ def ensure_call_wrapper( call: Callable[P_HandlerParams, T_HandlerReturn], outer_config: BrokerConfig, + wrapper_class: "type[HandlerCallWrapper[..., Any]]", ) -> "HandlerCallWrapper[P_HandlerParams, T_HandlerReturn]": + # An already wrapped handler keeps its class: the first decorator decides it if isinstance(call, HandlerCallWrapper): return call - return HandlerCallWrapper(call, outer_config) + return cast( + "HandlerCallWrapper[P_HandlerParams, T_HandlerReturn]", + wrapper_class(call, outer_config), + ) class HandlerCallWrapper(CallAssertions, Generic[P_HandlerParams, T_HandlerReturn]): diff --git a/faststream/_internal/endpoint/usecase.py b/faststream/_internal/endpoint/usecase.py index f5cf4719f2b..b7eff58bf8f 100644 --- a/faststream/_internal/endpoint/usecase.py +++ b/faststream/_internal/endpoint/usecase.py @@ -1,5 +1,5 @@ from collections.abc import Callable -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from faststream._internal.types import P_HandlerParams, T_HandlerReturn @@ -13,6 +13,10 @@ class Endpoint: + # A broker's subscriber and publisher classes override it with a subclass + # whose Call assertions take the broker's own fields + _call_wrapper_class: type[HandlerCallWrapper[..., Any]] = HandlerCallWrapper + def __init__(self, config: "BrokerConfig") -> None: self._outer_config = config @@ -21,6 +25,6 @@ def __call__( func: Callable[P_HandlerParams, T_HandlerReturn], ) -> HandlerCallWrapper[P_HandlerParams, T_HandlerReturn]: handler: HandlerCallWrapper[P_HandlerParams, T_HandlerReturn] = ( - ensure_call_wrapper(func, self._outer_config) + ensure_call_wrapper(func, self._outer_config, self._call_wrapper_class) ) return handler diff --git a/faststream/_internal/kafka/calls.py b/faststream/_internal/kafka/calls.py new file mode 100644 index 00000000000..e8f42b40f50 --- /dev/null +++ b/faststream/_internal/kafka/calls.py @@ -0,0 +1,168 @@ +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any + +from faststream._internal.constants import EMPTY +from faststream._internal.testing.calls import CallAssertions, ExpectedCall + +if TYPE_CHECKING: + from faststream.message import StreamMessage + + +class KafkaCallAssertions(CallAssertions): + """The Call assertions of a Kafka endpoint: the message fields, `key` and `partition`.""" + + __slots__ = () + + async def assert_called_once_with( + self, + body: Any = EMPTY, + /, + *, + headers: Any = EMPTY, + correlation_id: Any = EMPTY, + reply_to: Any = EMPTY, + content_type: Any = EMPTY, + path: Any = EMPTY, + context: Mapping[str, Any] = EMPTY, + key: Any = EMPTY, + partition: Any = EMPTY, + ) -> None: + """Assert the endpoint was called once, with the message described here. + + Args: + body: The body as a dict, a model or a matcher; it goes through the codec. + headers: Headers the message must carry; the rest may carry more. + correlation_id: The exact correlation id. + reply_to: The exact reply-to destination. + content_type: The exact content type. + path: The exact path parameters the subject template matched. + context: Context paths, as given to `Context()`, mapped to their values. + key: The exact record key, as the bytes the client delivered. + partition: The exact partition the record was read from. + """ + await self._assert_called_once_with( + self._describe( + body, + headers=headers, + correlation_id=correlation_id, + reply_to=reply_to, + content_type=content_type, + path=path, + context=context, + key=key, + partition=partition, + ) + ) + + async def assert_called_with( + self, + body: Any = EMPTY, + /, + *, + headers: Any = EMPTY, + correlation_id: Any = EMPTY, + reply_to: Any = EMPTY, + content_type: Any = EMPTY, + path: Any = EMPTY, + context: Mapping[str, Any] = EMPTY, + key: Any = EMPTY, + partition: Any = EMPTY, + ) -> None: + """Assert the last message the endpoint saw is the one described here. + + Args: + body: The body as a dict, a model or a matcher; it goes through the codec. + headers: Headers the message must carry; the rest may carry more. + correlation_id: The exact correlation id. + reply_to: The exact reply-to destination. + content_type: The exact content type. + path: The exact path parameters the subject template matched. + context: Context paths, as given to `Context()`, mapped to their values. + key: The exact record key, as the bytes the client delivered. + partition: The exact partition the record was read from. + """ + await self._assert_called_with( + self._describe( + body, + headers=headers, + correlation_id=correlation_id, + reply_to=reply_to, + content_type=content_type, + path=path, + context=context, + key=key, + partition=partition, + ) + ) + + async def assert_any_call( + self, + body: Any = EMPTY, + /, + *, + headers: Any = EMPTY, + correlation_id: Any = EMPTY, + reply_to: Any = EMPTY, + content_type: Any = EMPTY, + path: Any = EMPTY, + context: Mapping[str, Any] = EMPTY, + key: Any = EMPTY, + partition: Any = EMPTY, + ) -> None: + """Assert one of the messages the endpoint saw is the one described here. + + Args: + body: The body as a dict, a model or a matcher; it goes through the codec. + headers: Headers the message must carry; the rest may carry more. + correlation_id: The exact correlation id. + reply_to: The exact reply-to destination. + content_type: The exact content type. + path: The exact path parameters the subject template matched. + context: Context paths, as given to `Context()`, mapped to their values. + key: The exact record key, as the bytes the client delivered. + partition: The exact partition the record was read from. + """ + await self._assert_any_call( + self._describe( + body, + headers=headers, + correlation_id=correlation_id, + reply_to=reply_to, + content_type=content_type, + path=path, + context=context, + key=key, + partition=partition, + ) + ) + + @staticmethod + def _read_field(name: str, message: "StreamMessage[Any]") -> Any: + """Take a Kafka field off the raw message; each Kafka package knows its client's.""" + raise NotImplementedError + + def _describe( + self, + body: Any = EMPTY, + /, + *, + headers: Any = EMPTY, + correlation_id: Any = EMPTY, + reply_to: Any = EMPTY, + content_type: Any = EMPTY, + path: Any = EMPTY, + context: Mapping[str, Any] = EMPTY, + key: Any = EMPTY, + partition: Any = EMPTY, + ) -> ExpectedCall: + return ExpectedCall( + body=body, + headers=headers, + correlation_id=correlation_id, + reply_to=reply_to, + content_type=content_type, + path=path, + context=context, + broker_fields={"key": key, "partition": partition}, + read_field=self._read_field, + ) diff --git a/faststream/_internal/testing/calls.py b/faststream/_internal/testing/calls.py index bd98f693483..c7ae2ecc431 100644 --- a/faststream/_internal/testing/calls.py +++ b/faststream/_internal/testing/calls.py @@ -1,7 +1,7 @@ -from collections.abc import Mapping +from collections.abc import Callable, Mapping from copy import copy -from dataclasses import dataclass, field, fields -from typing import TYPE_CHECKING, Any +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, TypeAlias from unittest.mock import MagicMock from faststream._internal.constants import EMPTY @@ -13,6 +13,9 @@ from faststream._internal.configs import BrokerConfig from faststream.message import StreamMessage +# Takes a Broker field off the raw message, under the name the broker's `publish()` uses +FieldReader: TypeAlias = Callable[[str, "StreamMessage[Any]"], Any] + class CallAssertions: """The mock and the Call assertions an endpoint answers with under a test broker.""" @@ -50,9 +53,7 @@ async def assert_called_once_with( path: The exact path parameters the subject template matched. context: Context paths, as given to `Context()`, mapped to their values. """ - recorder = self._recorder_with_calls() - recorder.mock.assert_called_once() - await recorder.assert_last_call( + await self._assert_called_once_with( ExpectedCall( body=body, headers=headers, @@ -87,8 +88,7 @@ async def assert_called_with( path: The exact path parameters the subject template matched. context: Context paths, as given to `Context()`, mapped to their values. """ - recorder = self._recorder_with_calls() - await recorder.assert_last_call( + await self._assert_called_with( ExpectedCall( body=body, headers=headers, @@ -123,8 +123,7 @@ async def assert_any_call( path: The exact path parameters the subject template matched. context: Context paths, as given to `Context()`, mapped to their values. """ - recorder = self._recorder_with_calls() - await recorder.assert_any_call( + await self._assert_any_call( ExpectedCall( body=body, headers=headers, @@ -136,6 +135,18 @@ async def assert_any_call( ) ) + async def _assert_called_once_with(self, expected: "ExpectedCall") -> None: + """The check behind `assert_called_once_with`; a broker's mixin enters here.""" + recorder = self._recorder_with_calls() + recorder.mock.assert_called_once() + await recorder.assert_last_call(expected) + + async def _assert_called_with(self, expected: "ExpectedCall") -> None: + await self._recorder_with_calls().assert_last_call(expected) + + async def _assert_any_call(self, expected: "ExpectedCall") -> None: + await self._recorder_with_calls().assert_any_call(expected) + def _recorder_with_calls(self) -> "CallRecorder": recorder = self._recorder_under_test() # Every Call assertion answers alike for an endpoint nobody called @@ -250,6 +261,11 @@ async def _mismatches( actual = EMPTY checks.compare(f"context[{key!r}]", value, actual) + if expected.read_field is not None: + for name, value in expected.broker_fields.items(): + if value is not EMPTY: + checks.compare(name, value, expected.read_field(name, call.message)) + return checks.lines async def _expected_body(self, body: Any, message: "StreamMessage[Any]") -> Any: @@ -290,13 +306,24 @@ class ExpectedCall: content_type: Any path: Any context: Mapping[str, Any] + # A broker's own fields under the names its `publish()` takes, and the reader + # its wrapper hands over to take each one off the raw message + broker_fields: Mapping[str, Any] = field(default_factory=dict) + read_field: FieldReader | None = None def message_fields(self) -> list[str]: """The names asked of the message beside its body.""" + named = ( + "headers", + "correlation_id", + "reply_to", + "content_type", + "path", + "context", + ) return [ - f.name - for f in fields(self) - if f.name != "body" and getattr(self, f.name) is not EMPTY + *(name for name in named if getattr(self, name) is not EMPTY), + *(name for name, value in self.broker_fields.items() if value is not EMPTY), ] diff --git a/faststream/confluent/call_wrapper.py b/faststream/confluent/call_wrapper.py new file mode 100644 index 00000000000..c8003255d4a --- /dev/null +++ b/faststream/confluent/call_wrapper.py @@ -0,0 +1,35 @@ +from typing import TYPE_CHECKING, Any + +from confluent_kafka import Message + +from faststream._internal.endpoint.call_wrapper import HandlerCallWrapper +from faststream._internal.kafka.calls import KafkaCallAssertions +from faststream._internal.types import P_HandlerParams, T_HandlerReturn +from faststream.exceptions import SetupError + +if TYPE_CHECKING: + from faststream.message import StreamMessage + + +def read_field(name: str, message: "StreamMessage[Any]") -> Any: + """Take a Kafka field off a confluent message, under the name `publish()` gives it.""" + raw = message.raw_message + if not isinstance(raw, Message): + msg = ( + f"`{name}` is a Kafka field, and this message did not come from Kafka: " + f"its raw message is a `{type(raw).__name__}`." + ) + raise SetupError(msg) + # The confluent client answers with methods where aiokafka has attributes + return getattr(raw, name)() + + +class KafkaHandlerCallWrapper( + KafkaCallAssertions, + HandlerCallWrapper[P_HandlerParams, T_HandlerReturn], +): + """The wrapper of a Kafka handler: its Call assertions take `key` and `partition`.""" + + __slots__ = () + + _read_field = staticmethod(read_field) diff --git a/faststream/confluent/publisher/usecase.py b/faststream/confluent/publisher/usecase.py index 3d680e47b69..59335bfb68c 100644 --- a/faststream/confluent/publisher/usecase.py +++ b/faststream/confluent/publisher/usecase.py @@ -1,5 +1,5 @@ import asyncio -from collections.abc import Iterable +from collections.abc import Callable, Iterable from typing import TYPE_CHECKING, Any, Literal, Union, cast, overload from confluent_kafka import Message @@ -9,6 +9,9 @@ PublisherSpecification, PublisherUsecase, ) +from faststream._internal.kafka.calls import KafkaCallAssertions +from faststream._internal.types import P_HandlerParams, T_HandlerReturn +from faststream.confluent.call_wrapper import KafkaHandlerCallWrapper, read_field from faststream.confluent.response import KafkaPublishCommand from faststream.response.publish_type import PublishType @@ -22,9 +25,12 @@ from .producer import AsyncConfluentFastProducer -class LogicPublisher(PublisherUsecase): +class LogicPublisher(KafkaCallAssertions, PublisherUsecase): """A class to publish messages to a Kafka topic.""" + _call_wrapper_class = KafkaHandlerCallWrapper + _read_field = staticmethod(read_field) + def __init__( self, config: "KafkaPublisherConfig", @@ -37,6 +43,17 @@ def __init__( self.reply_to = config.reply_to self.headers = config.headers or {} + @override + def __call__( + self, + func: Callable[P_HandlerParams, T_HandlerReturn], + ) -> "KafkaHandlerCallWrapper[P_HandlerParams, T_HandlerReturn]": + # The base builds the wrapper from `_call_wrapper_class`; this only narrows the name + return cast( + "KafkaHandlerCallWrapper[P_HandlerParams, T_HandlerReturn]", + super().__call__(func), + ) + @property def topic(self) -> str: return f"{self._outer_config.prefix}{self._topic}" diff --git a/faststream/confluent/subscriber/usecase.py b/faststream/confluent/subscriber/usecase.py index b260159cb40..019a2c033c5 100644 --- a/faststream/confluent/subscriber/usecase.py +++ b/faststream/confluent/subscriber/usecase.py @@ -1,29 +1,35 @@ import logging from abc import abstractmethod -from collections.abc import AsyncIterator, Sequence +from collections.abc import AsyncIterator, Callable, Iterable, Sequence from typing import ( TYPE_CHECKING, Any, Optional, + Union, cast, ) import anyio from confluent_kafka import KafkaException, Message -from typing_extensions import override +from typing_extensions import overload, override from faststream._internal.endpoint.subscriber import SubscriberUsecase from faststream._internal.endpoint.subscriber.mixins import ConcurrentMixin, TasksMixin +from faststream._internal.endpoint.subscriber.utils import default_filter from faststream._internal.endpoint.utils import process_msg -from faststream._internal.types import MsgType +from faststream._internal.types import MsgType, P_HandlerParams, T_HandlerReturn +from faststream.confluent.call_wrapper import KafkaHandlerCallWrapper from faststream.confluent.parser import AsyncConfluentParser from faststream.confluent.publisher.fake import KafkaFakePublisher from faststream.confluent.schemas import Topic, TopicPartition if TYPE_CHECKING: + from fast_depends.dependencies import Dependant + from faststream._internal.endpoint.publisher import PublisherProto from faststream._internal.endpoint.subscriber import SubscriberSpecification from faststream._internal.endpoint.subscriber.call_item import CallsCollection + from faststream._internal.types import CustomCallable, Filter from faststream.confluent.configs import KafkaBrokerConfig from faststream.confluent.helpers.client import AsyncConfluentConsumer from faststream.confluent.message import KafkaMessage @@ -36,6 +42,7 @@ class LogicSubscriber(TasksMixin, SubscriberUsecase[MsgType]): """A class to handle logic for consuming messages from Kafka.""" _outer_config: "KafkaBrokerConfig" + _call_wrapper_class = KafkaHandlerCallWrapper group_id: str | None @@ -60,6 +67,61 @@ def __init__( self.consumer = None self.polling_interval = config.polling_interval + @overload + def __call__( + self, + func: Callable[P_HandlerParams, T_HandlerReturn], + *, + filter: "Filter[Any]" = default_filter, + parser: Optional["CustomCallable"] = None, + decoder: Optional["CustomCallable"] = None, + dependencies: Iterable["Dependant"] = (), + ) -> "KafkaHandlerCallWrapper[P_HandlerParams, T_HandlerReturn]": ... + + @overload + def __call__( + self, + func: None = None, + *, + filter: "Filter[Any]" = default_filter, + parser: Optional["CustomCallable"] = None, + decoder: Optional["CustomCallable"] = None, + dependencies: Iterable["Dependant"] = (), + ) -> Callable[ + [Callable[P_HandlerParams, T_HandlerReturn]], + "KafkaHandlerCallWrapper[P_HandlerParams, T_HandlerReturn]", + ]: ... + + @override + def __call__( + self, + func: Callable[P_HandlerParams, T_HandlerReturn] | None = None, + *, + filter: "Filter[Any]" = default_filter, + parser: Optional["CustomCallable"] = None, + decoder: Optional["CustomCallable"] = None, + dependencies: Iterable["Dependant"] = (), + ) -> Union[ + "KafkaHandlerCallWrapper[P_HandlerParams, T_HandlerReturn]", + Callable[ + [Callable[P_HandlerParams, T_HandlerReturn]], + "KafkaHandlerCallWrapper[P_HandlerParams, T_HandlerReturn]", + ], + ]: + # The base builds the wrapper from `_call_wrapper_class`; this only narrows the name + return cast( + "KafkaHandlerCallWrapper[P_HandlerParams, T_HandlerReturn] | Callable[" + "[Callable[P_HandlerParams, T_HandlerReturn]], " + "KafkaHandlerCallWrapper[P_HandlerParams, T_HandlerReturn]]", + super().__call__( + func, + filter=filter, + parser=parser, + decoder=decoder, + dependencies=dependencies, + ), + ) + @property def client_id(self) -> str | None: return self._outer_config.client_id diff --git a/faststream/confluent/testing.py b/faststream/confluent/testing.py index 5f7d0f2ec01..964b5d4bf8d 100644 --- a/faststream/confluent/testing.py +++ b/faststream/confluent/testing.py @@ -5,6 +5,7 @@ from unittest.mock import AsyncMock, MagicMock import anyio +from confluent_kafka import KafkaError, Message from typing_extensions import override from faststream._internal.endpoint.utils import ParserComposition @@ -281,18 +282,21 @@ async def _execute_handler( ) -class MockConfluentMessage: +# A `Message` so that the reader of the Kafka fields takes it for the client's own +class MockConfluentMessage(Message): + """The message the in-memory broker delivers, answering as the client's does.""" + def __init__( self, raw_msg: bytes | None, topic: str, key: bytes | str, - headers: list[tuple[str, bytes]], + headers: list[tuple[str, str | bytes | None]], offset: int, partition: int, timestamp_type: int, timestamp_ms: int, - error: str | None = None, + error: KafkaError | None = None, ) -> None: self._raw_msg = raw_msg self._topic = topic @@ -311,10 +315,10 @@ def __init__( def len(self) -> int: return 0 if self._raw_msg is None else len(self._raw_msg) - def error(self) -> str | None: + def error(self) -> KafkaError | None: return self._error - def headers(self) -> list[tuple[str, bytes]]: + def headers(self) -> list[tuple[str, str | bytes | None]]: return self._headers def key(self) -> bytes: diff --git a/faststream/kafka/call_wrapper.py b/faststream/kafka/call_wrapper.py new file mode 100644 index 00000000000..945c97eaa4e --- /dev/null +++ b/faststream/kafka/call_wrapper.py @@ -0,0 +1,34 @@ +from typing import TYPE_CHECKING, Any + +from aiokafka import ConsumerRecord + +from faststream._internal.endpoint.call_wrapper import HandlerCallWrapper +from faststream._internal.kafka.calls import KafkaCallAssertions +from faststream._internal.types import P_HandlerParams, T_HandlerReturn +from faststream.exceptions import SetupError + +if TYPE_CHECKING: + from faststream.message import StreamMessage + + +def read_field(name: str, message: "StreamMessage[Any]") -> Any: + """Take a Kafka field off an aiokafka record, under the name `publish()` gives it.""" + raw = message.raw_message + if not isinstance(raw, ConsumerRecord): + msg = ( + f"`{name}` is a Kafka field, and this message did not come from Kafka: " + f"its raw message is a `{type(raw).__name__}`." + ) + raise SetupError(msg) + return getattr(raw, name) + + +class KafkaHandlerCallWrapper( + KafkaCallAssertions, + HandlerCallWrapper[P_HandlerParams, T_HandlerReturn], +): + """The wrapper of a Kafka handler: its Call assertions take `key` and `partition`.""" + + __slots__ = () + + _read_field = staticmethod(read_field) diff --git a/faststream/kafka/publisher/usecase.py b/faststream/kafka/publisher/usecase.py index da5d1b8eec5..7427f5f83ea 100644 --- a/faststream/kafka/publisher/usecase.py +++ b/faststream/kafka/publisher/usecase.py @@ -1,9 +1,12 @@ -from collections.abc import Iterable +from collections.abc import Callable, Iterable from typing import TYPE_CHECKING, Any, Literal, Union, cast, overload from typing_extensions import override from faststream._internal.endpoint.publisher import PublisherUsecase +from faststream._internal.kafka.calls import KafkaCallAssertions +from faststream._internal.types import P_HandlerParams, T_HandlerReturn +from faststream.kafka.call_wrapper import KafkaHandlerCallWrapper, read_field from faststream.kafka.response import KafkaPublishCommand from faststream.response.publish_type import PublishType @@ -22,9 +25,12 @@ from .producer import AioKafkaFastProducer -class LogicPublisher(PublisherUsecase): +class LogicPublisher(KafkaCallAssertions, PublisherUsecase): """A class to publish messages to a Kafka topic.""" + _call_wrapper_class = KafkaHandlerCallWrapper + _read_field = staticmethod(read_field) + def __init__( self, config: "KafkaPublisherConfig", @@ -37,6 +43,17 @@ def __init__( self.reply_to = config.reply_to self.headers = config.headers or {} + @override + def __call__( + self, + func: Callable[P_HandlerParams, T_HandlerReturn], + ) -> "KafkaHandlerCallWrapper[P_HandlerParams, T_HandlerReturn]": + # The base builds the wrapper from `_call_wrapper_class`; this only narrows the name + return cast( + "KafkaHandlerCallWrapper[P_HandlerParams, T_HandlerReturn]", + super().__call__(func), + ) + @property def topic(self) -> str: return f"{self._outer_config.prefix}{self._topic}" diff --git a/faststream/kafka/subscriber/usecase.py b/faststream/kafka/subscriber/usecase.py index 83770f0874e..338c798f35b 100644 --- a/faststream/kafka/subscriber/usecase.py +++ b/faststream/kafka/subscriber/usecase.py @@ -1,8 +1,8 @@ import logging from abc import abstractmethod -from collections.abc import AsyncIterator, Callable, Sequence +from collections.abc import AsyncIterator, Callable, Iterable, Sequence from itertools import chain -from typing import TYPE_CHECKING, Any, Optional, cast +from typing import TYPE_CHECKING, Any, Optional, Union, cast import anyio from aiokafka import ( @@ -10,13 +10,15 @@ TopicPartition as AIOKafkaTopicPartition, ) from aiokafka.errors import ConsumerStoppedError, KafkaError, UnsupportedCodecError -from typing_extensions import override +from typing_extensions import overload, override from faststream._internal.endpoint.subscriber.mixins import ConcurrentMixin, TasksMixin from faststream._internal.endpoint.subscriber.usecase import SubscriberUsecase +from faststream._internal.endpoint.subscriber.utils import default_filter from faststream._internal.endpoint.utils import process_msg -from faststream._internal.types import MsgType +from faststream._internal.types import MsgType, P_HandlerParams, T_HandlerReturn from faststream._internal.utils.path import Address, AddressSyntax +from faststream.kafka.call_wrapper import KafkaHandlerCallWrapper from faststream.kafka.helpers import make_logging_listener from faststream.kafka.message import KafkaAckableMessage, KafkaMessage, KafkaRawMessage from faststream.kafka.parser import AioKafkaBatchParser, AioKafkaParser @@ -24,10 +26,12 @@ if TYPE_CHECKING: from aiokafka import AIOKafkaConsumer + from fast_depends.dependencies import Dependant from faststream._internal.endpoint.publisher import PublisherProto from faststream._internal.endpoint.subscriber import SubscriberSpecification from faststream._internal.endpoint.subscriber.call_item import CallsCollection + from faststream._internal.types import CustomCallable, Filter from faststream.kafka.configs import KafkaBrokerConfig from faststream.message import StreamMessage @@ -49,6 +53,7 @@ class LogicSubscriber(TasksMixin, SubscriberUsecase[MsgType]): parser: AioKafkaParser _outer_config: "KafkaBrokerConfig" + _call_wrapper_class = KafkaHandlerCallWrapper def __init__( self, @@ -68,6 +73,61 @@ def __init__( self.consumer = None + @overload + def __call__( + self, + func: Callable[P_HandlerParams, T_HandlerReturn], + *, + filter: "Filter[Any]" = default_filter, + parser: Optional["CustomCallable"] = None, + decoder: Optional["CustomCallable"] = None, + dependencies: Iterable["Dependant"] = (), + ) -> "KafkaHandlerCallWrapper[P_HandlerParams, T_HandlerReturn]": ... + + @overload + def __call__( + self, + func: None = None, + *, + filter: "Filter[Any]" = default_filter, + parser: Optional["CustomCallable"] = None, + decoder: Optional["CustomCallable"] = None, + dependencies: Iterable["Dependant"] = (), + ) -> Callable[ + [Callable[P_HandlerParams, T_HandlerReturn]], + "KafkaHandlerCallWrapper[P_HandlerParams, T_HandlerReturn]", + ]: ... + + @override + def __call__( + self, + func: Callable[P_HandlerParams, T_HandlerReturn] | None = None, + *, + filter: "Filter[Any]" = default_filter, + parser: Optional["CustomCallable"] = None, + decoder: Optional["CustomCallable"] = None, + dependencies: Iterable["Dependant"] = (), + ) -> Union[ + "KafkaHandlerCallWrapper[P_HandlerParams, T_HandlerReturn]", + Callable[ + [Callable[P_HandlerParams, T_HandlerReturn]], + "KafkaHandlerCallWrapper[P_HandlerParams, T_HandlerReturn]", + ], + ]: + # The base builds the wrapper from `_call_wrapper_class`; this only narrows the name + return cast( + "KafkaHandlerCallWrapper[P_HandlerParams, T_HandlerReturn] | Callable[" + "[Callable[P_HandlerParams, T_HandlerReturn]], " + "KafkaHandlerCallWrapper[P_HandlerParams, T_HandlerReturn]]", + super().__call__( + func, + filter=filter, + parser=parser, + decoder=decoder, + dependencies=dependencies, + ), + ) + @property def pattern(self) -> Address | None: """The pattern this Subscriber was declared with, and its Broker address.""" diff --git a/tests/brokers/confluent/test_test_client.py b/tests/brokers/confluent/test_test_client.py index 26dbcc75607..bd4a3e066a4 100644 --- a/tests/brokers/confluent/test_test_client.py +++ b/tests/brokers/confluent/test_test_client.py @@ -158,6 +158,41 @@ async def m(msg) -> None: ... with pytest.raises(SetupError, match="received a batch"): await m.assert_called_once_with(headers={"key": "value"}) + with pytest.raises(SetupError, match=r"received a batch.*key, partition"): + await m.assert_called_once_with(key=b"k", partition=0) + + async def test_assertions_take_the_kafka_fields(self, queue: str) -> None: + broker = self.get_broker() + + @broker.subscriber(queue) + async def handle(msg) -> None: ... + + async with self.patch_broker(broker) as br: + await br.publish("hello", queue, key=b"k", partition=1) + + await handle.assert_called_once_with("hello", key=b"k", partition=1) + await handle.assert_called_with(key=b"k") + await handle.assert_any_call(partition=1) + + with pytest.raises(AssertionError, match="key: expected b'other', got b'k'"): + await handle.assert_called_once_with("hello", key=b"other") + + async def test_publisher_assertions_take_the_kafka_fields(self, queue: str) -> None: + broker = self.get_broker() + + publisher = broker.publisher(queue + "2") + + @broker.subscriber(queue) + async def handle(msg) -> None: + await publisher.publish("response", key=b"k", partition=1) + + async with self.patch_broker(broker) as br: + await br.publish("hello", queue) + + await publisher.assert_called_once_with("response", key=b"k", partition=1) + await publisher.assert_called_with(key=b"k") + await publisher.assert_any_call(partition=1) + async def test_batch_publisher_mock(self, queue: str) -> None: broker = self.get_broker() diff --git a/tests/brokers/kafka/test_test_client.py b/tests/brokers/kafka/test_test_client.py index c8a1585403b..04a72eef965 100644 --- a/tests/brokers/kafka/test_test_client.py +++ b/tests/brokers/kafka/test_test_client.py @@ -210,6 +210,60 @@ async def m(msg) -> None: ... with pytest.raises(SetupError, match="received a batch"): await m.assert_called_once_with(headers={"key": "value"}) + with pytest.raises(SetupError, match=r"received a batch.*key, partition"): + await m.assert_called_once_with(key=b"k", partition=0) + + async def test_assertions_take_the_kafka_fields(self, queue: str) -> None: + broker = self.get_broker() + + @broker.subscriber(queue) + async def handle(msg) -> None: ... + + async with self.patch_broker(broker) as br: + await br.publish("hello", queue, key=b"k", partition=1) + + await handle.assert_called_once_with("hello", key=b"k", partition=1) + await handle.assert_called_with(key=b"k") + await handle.assert_any_call(partition=1) + + with pytest.raises(AssertionError, match="key: expected b'other', got b'k'"): + await handle.assert_called_once_with("hello", key=b"other") + + async def test_publisher_assertions_take_the_kafka_fields(self, queue: str) -> None: + broker = self.get_broker() + + publisher = broker.publisher(queue + "2") + + @broker.subscriber(queue) + async def handle(msg) -> None: + await publisher.publish("response", key=b"k", partition=1) + + async with self.patch_broker(broker) as br: + await br.publish("hello", queue) + + await publisher.assert_called_once_with("response", key=b"k", partition=1) + await publisher.assert_called_with(key=b"k") + await publisher.assert_any_call(partition=1) + + async def test_kafka_fields_refuse_another_brokers_message(self, queue: str) -> None: + from faststream.rabbit import RabbitBroker, TestRabbitBroker + + broker = self.get_broker() + rabbit = RabbitBroker() + + # The first decorator decides the wrapper class: Kafka's here + @rabbit.subscriber(queue) + @broker.subscriber(queue) + async def handle(msg) -> None: ... + + async with self.patch_broker(broker), TestRabbitBroker(rabbit): + await rabbit.publish("hello", queue) + + await handle.assert_called_once_with("hello") + + with pytest.raises(SetupError, match="`key` is a Kafka field"): + await handle.assert_called_once_with("hello", key=b"k") + async def test_batch_publisher_mock( self, queue: str, diff --git a/tests/docs/getting_started/subscription/test_advanced_testing.py b/tests/docs/getting_started/subscription/test_advanced_testing.py index 9b4980660ac..9accdc8061b 100644 --- a/tests/docs/getting_started/subscription/test_advanced_testing.py +++ b/tests/docs/getting_started/subscription/test_advanced_testing.py @@ -16,6 +16,7 @@ async def test_handle_kafka() -> None: from docs.docs_src.getting_started.subscription.kafka.advanced_testing import ( test_handle, + test_kafka_fields, test_message_context, test_several_messages, ) @@ -23,6 +24,7 @@ async def test_handle_kafka() -> None: await test_handle() await test_message_context() await test_several_messages() + await test_kafka_fields() @pytest.mark.confluent() @@ -31,6 +33,7 @@ async def test_handle_kafka() -> None: async def test_handle_confluent() -> None: from docs.docs_src.getting_started.subscription.confluent.advanced_testing import ( test_handle, + test_kafka_fields, test_message_context, test_several_messages, ) @@ -38,6 +41,7 @@ async def test_handle_confluent() -> None: await test_handle() await test_message_context() await test_several_messages() + await test_kafka_fields() @pytest.mark.rabbit() diff --git a/tests/mypy/confluent.py b/tests/mypy/confluent.py index dda7b79add8..e2d107048bf 100644 --- a/tests/mypy/confluent.py +++ b/tests/mypy/confluent.py @@ -1,9 +1,12 @@ import asyncio +from collections.abc import Coroutine +from typing import Any from confluent_kafka import Message from typing_extensions import assert_type from faststream.confluent import KafkaBroker, KafkaMessage, KafkaRouter, TestKafkaBroker +from faststream.confluent.call_wrapper import KafkaHandlerCallWrapper from faststream.confluent.fastapi import KafkaRouter as FastAPIRouter from faststream.confluent.publisher.usecase import ( BatchPublisher, @@ -125,6 +128,28 @@ def check_subscriber_instance_type( assert_type(sub3, ConcurrentDefaultSubscriber) +async def check_call_assertions_take_the_kafka_fields( + broker: KafkaBroker | FastAPIRouter | KafkaRouter, +) -> None: + @broker.subscriber("test") + async def handle() -> None: ... + + assert_type(handle, KafkaHandlerCallWrapper[[], Coroutine[Any, Any, None]]) + await handle.assert_called_once_with(None, key=b"k", partition=0) + await handle.assert_called_with(key=b"k") + await handle.assert_any_call(partition=0) + + publisher = broker.publisher("test") + + @publisher + async def published() -> None: ... + + assert_type(published, KafkaHandlerCallWrapper[[], Coroutine[Any, Any, None]]) + await publisher.assert_called_once_with(None, key=b"k", partition=0) + await publisher.assert_called_with(key=b"k") + await publisher.assert_any_call(partition=0) + + KafkaBroker(routers=[KafkaRouter()]) KafkaBroker().include_router(KafkaRouter()) KafkaBroker().include_routers(KafkaRouter()) diff --git a/tests/mypy/kafka.py b/tests/mypy/kafka.py index bd72bf13132..2dc43ba6fd4 100644 --- a/tests/mypy/kafka.py +++ b/tests/mypy/kafka.py @@ -1,5 +1,5 @@ import asyncio -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Coroutine from typing import Any import prometheus_client @@ -15,6 +15,7 @@ RecordMetadata, TestKafkaBroker, ) +from faststream.kafka.call_wrapper import KafkaHandlerCallWrapper from faststream.kafka.fastapi import KafkaRouter as FastAPIRouter from faststream.kafka.opentelemetry import KafkaTelemetryMiddleware from faststream.kafka.prometheus import KafkaPrometheusMiddleware @@ -420,6 +421,28 @@ def check_publisher_instance_type( assert_type(pub2, BatchPublisher) +async def check_call_assertions_take_the_kafka_fields( + broker: KafkaBroker | FastAPIRouter | KafkaRouter, +) -> None: + @broker.subscriber("test") + async def handle() -> None: ... + + assert_type(handle, KafkaHandlerCallWrapper[[], Coroutine[Any, Any, None]]) + await handle.assert_called_once_with(None, key=b"k", partition=0) + await handle.assert_called_with(key=b"k") + await handle.assert_any_call(partition=0) + + publisher = broker.publisher("test") + + @publisher + async def published() -> None: ... + + assert_type(published, KafkaHandlerCallWrapper[[], Coroutine[Any, Any, None]]) + await publisher.assert_called_once_with(None, key=b"k", partition=0) + await publisher.assert_called_with(key=b"k") + await publisher.assert_any_call(partition=0) + + def fake_bool() -> bool: return True From 5f57fb035ceb48fd2ffac41b7667410ef66a10f6 Mon Sep 17 00:00:00 2001 From: Lancetnik Date: Wed, 9 Sep 2026 11:33:24 +0300 Subject: [PATCH 2/8] docs: the Kafka testing snippets assert on `key` The AIOKafka and Confluent tabs gain a `test_kafka_fields` snippet, the context example no longer reaches for the key by path, and the testing page says each broker's handlers and publishers take that broker's own fields. Co-Authored-By: Claude Fable 5.1 --- .../en/getting-started/subscription/test.md | 20 +++++++++++++++---- .../confluent/advanced_testing.py | 17 ++++++++++++++++ .../subscription/kafka/advanced_testing.py | 20 +++++++++++++++++-- 3 files changed, 51 insertions(+), 6 deletions(-) diff --git a/docs/docs/en/getting-started/subscription/test.md b/docs/docs/en/getting-started/subscription/test.md index dc69a745ad8..6b6df8825a5 100644 --- a/docs/docs/en/getting-started/subscription/test.md +++ b/docs/docs/en/getting-started/subscription/test.md @@ -308,7 +308,7 @@ A handler that saw several messages answers with two more methods, named as in ` === "AIOKafka" ```python linenums="1" hl_lines="24-27 29-32" - {!> docs_src/getting_started/subscription/kafka/advanced_testing.py [ln:3,5-8,64-90] !} + {!> docs_src/getting_started/subscription/kafka/advanced_testing.py [ln:3,5-8,63-89] !} ``` === "Confluent" @@ -336,11 +336,11 @@ A handler that saw several messages answers with two more methods, named as in ` {!> docs_src/getting_started/subscription/mqtt/advanced_testing.py [ln:3,5-8,63-89] !} ``` -To check only a part of the body, put a [dirty-equals](https://dirty-equals.helpmanual.io/){.external-link target="_blank"} matcher in its place. Anything the fields above do not cover, such as the Kafka message key, lives in the context: check it through `context` by the same path you would give to `Context()`, walking attributes and dict keys from a context name. +To check only a part of the body, put a [dirty-equals](https://dirty-equals.helpmanual.io/){.external-link target="_blank"} matcher in its place. Anything the fields above do not cover lives in the context: check it through `context` by the same path you would give to `Context()`, walking attributes and dict keys from a context name. === "AIOKafka" - ```python linenums="1" hl_lines="2 15 19-20" - {!> docs_src/getting_started/subscription/kafka/advanced_testing.py [ln:3-8,48-62] !} + ```python linenums="1" hl_lines="2 18-19" + {!> docs_src/getting_started/subscription/kafka/advanced_testing.py [ln:3-8,48-61] !} ``` === "Confluent" @@ -371,6 +371,18 @@ To check only a part of the body, put a [dirty-equals](https://dirty-equals.help !!! note A context path reads attributes and keys, it never calls. Where a raw message answers with methods, as the Confluent one does, reach for what **FastStream** put in the context beside it, such as `log_context`. +Beside these fields, the handlers and publishers of a broker take that broker's own, under the names its `publish()` uses for them. A Kafka handler or publisher takes `key` and `partition`, compared exactly as the client delivered them: + +=== "AIOKafka" + ```python linenums="1" hl_lines="14 18-21" + {!> docs_src/getting_started/subscription/kafka/advanced_testing.py [ln:3,5-8,91-106] !} + ``` + +=== "Confluent" + ```python linenums="1" hl_lines="14 18-21" + {!> docs_src/getting_started/subscription/confluent/advanced_testing.py [ln:3,5-8,91-106] !} + ``` + !!! note Both `handle.mock` and the three assertion methods exist only inside the test broker. Outside of it they raise a `SetupError` instead of answering for a handler nobody has called. diff --git a/docs/docs_src/getting_started/subscription/confluent/advanced_testing.py b/docs/docs_src/getting_started/subscription/confluent/advanced_testing.py index 646ced8217f..e36195fe88d 100644 --- a/docs/docs_src/getting_started/subscription/confluent/advanced_testing.py +++ b/docs/docs_src/getting_started/subscription/confluent/advanced_testing.py @@ -87,3 +87,20 @@ async def test_several_messages() -> None: Data(name="John", user_id=1), correlation_id="first", ) + + +@pytest.mark.asyncio +async def test_kafka_fields() -> None: + async with TestKafkaBroker(broker) as br: + await br.publish( + Data(name="John", user_id=1), + topic="test-topic", + headers={"trace-id": "42"}, + key=b"user-1", + ) + + # `key` and `partition` are named as `publish()` names them + await handle.assert_called_once_with( + Data(name="John", user_id=1), + key=b"user-1", + ) diff --git a/docs/docs_src/getting_started/subscription/kafka/advanced_testing.py b/docs/docs_src/getting_started/subscription/kafka/advanced_testing.py index 418bd26adb7..4ee5138a8ba 100644 --- a/docs/docs_src/getting_started/subscription/kafka/advanced_testing.py +++ b/docs/docs_src/getting_started/subscription/kafka/advanced_testing.py @@ -53,12 +53,11 @@ async def test_message_context() -> None: Data(name="John", user_id=1), topic="test-topic", headers={"trace-id": "42"}, - key=b"user-1", ) await handle.assert_called_once_with( IsPartialDict(name="John"), - context={"message.raw_message.key": b"user-1"}, + context={"log_context.topic": "test-topic"}, ) @@ -88,3 +87,20 @@ async def test_several_messages() -> None: Data(name="John", user_id=1), correlation_id="first", ) + + +@pytest.mark.asyncio +async def test_kafka_fields() -> None: + async with TestKafkaBroker(broker) as br: + await br.publish( + Data(name="John", user_id=1), + topic="test-topic", + headers={"trace-id": "42"}, + key=b"user-1", + ) + + # `key` and `partition` are named as `publish()` names them + await handle.assert_called_once_with( + Data(name="John", user_id=1), + key=b"user-1", + ) From 7a8ec509078654da5693fb24850dd1ef17869004 Mon Sep 17 00:00:00 2001 From: Lancetnik Date: Wed, 9 Sep 2026 11:38:45 +0300 Subject: [PATCH 3/8] refactor(testing): a Call assertion carries its Broker fields with their reader `ExpectedCall.broker_fields` is one `BrokerFields`: the values asked for, already stripped of the ones left EMPTY, and the reader that answers them. The recorder no longer filters EMPTY twice or guards a reader that may be missing, and the Kafka mixin's helper is named for what it returns. Co-Authored-By: Claude Fable 5.1 --- faststream/_internal/kafka/calls.py | 24 +++++++++++++++++------- faststream/_internal/testing/calls.py | 23 ++++++++++++++--------- 2 files changed, 31 insertions(+), 16 deletions(-) diff --git a/faststream/_internal/kafka/calls.py b/faststream/_internal/kafka/calls.py index e8f42b40f50..bf8a986e3e8 100644 --- a/faststream/_internal/kafka/calls.py +++ b/faststream/_internal/kafka/calls.py @@ -2,7 +2,11 @@ from typing import TYPE_CHECKING, Any from faststream._internal.constants import EMPTY -from faststream._internal.testing.calls import CallAssertions, ExpectedCall +from faststream._internal.testing.calls import ( + BrokerFields, + CallAssertions, + ExpectedCall, +) if TYPE_CHECKING: from faststream.message import StreamMessage @@ -41,7 +45,7 @@ async def assert_called_once_with( partition: The exact partition the record was read from. """ await self._assert_called_once_with( - self._describe( + self._expected_call( body, headers=headers, correlation_id=correlation_id, @@ -82,7 +86,7 @@ async def assert_called_with( partition: The exact partition the record was read from. """ await self._assert_called_with( - self._describe( + self._expected_call( body, headers=headers, correlation_id=correlation_id, @@ -123,7 +127,7 @@ async def assert_any_call( partition: The exact partition the record was read from. """ await self._assert_any_call( - self._describe( + self._expected_call( body, headers=headers, correlation_id=correlation_id, @@ -141,7 +145,7 @@ def _read_field(name: str, message: "StreamMessage[Any]") -> Any: """Take a Kafka field off the raw message; each Kafka package knows its client's.""" raise NotImplementedError - def _describe( + def _expected_call( self, body: Any = EMPTY, /, @@ -163,6 +167,12 @@ def _describe( content_type=content_type, path=path, context=context, - broker_fields={"key": key, "partition": partition}, - read_field=self._read_field, + broker_fields=BrokerFields( + { + name: value + for name, value in (("key", key), ("partition", partition)) + if value is not EMPTY + }, + self._read_field, + ), ) diff --git a/faststream/_internal/testing/calls.py b/faststream/_internal/testing/calls.py index c7ae2ecc431..6edeeb85cb0 100644 --- a/faststream/_internal/testing/calls.py +++ b/faststream/_internal/testing/calls.py @@ -261,10 +261,9 @@ async def _mismatches( actual = EMPTY checks.compare(f"context[{key!r}]", value, actual) - if expected.read_field is not None: - for name, value in expected.broker_fields.items(): - if value is not EMPTY: - checks.compare(name, value, expected.read_field(name, call.message)) + if (fields := expected.broker_fields) is not None: + for name, value in fields.values.items(): + checks.compare(name, value, fields.read(name, call.message)) return checks.lines @@ -306,10 +305,7 @@ class ExpectedCall: content_type: Any path: Any context: Mapping[str, Any] - # A broker's own fields under the names its `publish()` takes, and the reader - # its wrapper hands over to take each one off the raw message - broker_fields: Mapping[str, Any] = field(default_factory=dict) - read_field: FieldReader | None = None + broker_fields: "BrokerFields | None" = None def message_fields(self) -> list[str]: """The names asked of the message beside its body.""" @@ -323,10 +319,19 @@ def message_fields(self) -> list[str]: ) return [ *(name for name in named if getattr(self, name) is not EMPTY), - *(name for name, value in self.broker_fields.items() if value is not EMPTY), + *(self.broker_fields.values if self.broker_fields else ()), ] +@dataclass(slots=True) +class BrokerFields: + """A broker's own fields a Call assertion asks for, and the reader that answers them.""" + + # Under the names the broker's `publish()` takes; a field not asked for is absent + values: Mapping[str, Any] + read: FieldReader + + @dataclass(slots=True) class RecordedCall: message: "StreamMessage[Any]" From fbcba343c27e3ee20cba9c538cf960259d7ac664 Mon Sep 17 00:00:00 2001 From: Lancetnik Date: Wed, 9 Sep 2026 11:38:48 +0300 Subject: [PATCH 4/8] test(confluent): the reader refuses another broker's message The confluent reader's `isinstance` guard had no test going red on its own. Co-Authored-By: Claude Fable 5.1 --- tests/brokers/confluent/test_test_client.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/brokers/confluent/test_test_client.py b/tests/brokers/confluent/test_test_client.py index bd4a3e066a4..a3da396d714 100644 --- a/tests/brokers/confluent/test_test_client.py +++ b/tests/brokers/confluent/test_test_client.py @@ -193,6 +193,25 @@ async def handle(msg) -> None: await publisher.assert_called_with(key=b"k") await publisher.assert_any_call(partition=1) + async def test_kafka_fields_refuse_another_brokers_message(self, queue: str) -> None: + from faststream.rabbit import RabbitBroker, TestRabbitBroker + + broker = self.get_broker() + rabbit = RabbitBroker() + + # The first decorator decides the wrapper class: Kafka's here + @rabbit.subscriber(queue) + @broker.subscriber(queue) + async def handle(msg) -> None: ... + + async with self.patch_broker(broker), TestRabbitBroker(rabbit): + await rabbit.publish("hello", queue) + + await handle.assert_called_once_with("hello") + + with pytest.raises(SetupError, match="`key` is a Kafka field"): + await handle.assert_called_once_with("hello", key=b"k") + async def test_batch_publisher_mock(self, queue: str) -> None: broker = self.get_broker() From 0831aa17fcf870b4e8b304d87e6af89a098c8e86 Mon Sep 17 00:00:00 2001 From: Lancetnik Date: Wed, 9 Sep 2026 11:57:56 +0300 Subject: [PATCH 5/8] test(mypy): the wrapper type is asserted on a sync handler pyright spells an `async def`'s return as `CoroutineType[...]` where mypy says `Coroutine[...]`, so an `assert_type` over an async handler cannot satisfy both checkers; a sync handler reads the same in both. Co-Authored-By: Claude Fable 5.1 --- tests/mypy/confluent.py | 11 +++++------ tests/mypy/kafka.py | 11 ++++++----- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/tests/mypy/confluent.py b/tests/mypy/confluent.py index e2d107048bf..32a25932da7 100644 --- a/tests/mypy/confluent.py +++ b/tests/mypy/confluent.py @@ -1,6 +1,4 @@ import asyncio -from collections.abc import Coroutine -from typing import Any from confluent_kafka import Message from typing_extensions import assert_type @@ -131,10 +129,11 @@ def check_subscriber_instance_type( async def check_call_assertions_take_the_kafka_fields( broker: KafkaBroker | FastAPIRouter | KafkaRouter, ) -> None: + # A sync handler: mypy and pyright spell an `async def`'s return type differently @broker.subscriber("test") - async def handle() -> None: ... + def handle() -> None: ... - assert_type(handle, KafkaHandlerCallWrapper[[], Coroutine[Any, Any, None]]) + assert_type(handle, KafkaHandlerCallWrapper[[], None]) await handle.assert_called_once_with(None, key=b"k", partition=0) await handle.assert_called_with(key=b"k") await handle.assert_any_call(partition=0) @@ -142,9 +141,9 @@ async def handle() -> None: ... publisher = broker.publisher("test") @publisher - async def published() -> None: ... + def published() -> None: ... - assert_type(published, KafkaHandlerCallWrapper[[], Coroutine[Any, Any, None]]) + assert_type(published, KafkaHandlerCallWrapper[[], None]) await publisher.assert_called_once_with(None, key=b"k", partition=0) await publisher.assert_called_with(key=b"k") await publisher.assert_any_call(partition=0) diff --git a/tests/mypy/kafka.py b/tests/mypy/kafka.py index 2dc43ba6fd4..ebe1e134654 100644 --- a/tests/mypy/kafka.py +++ b/tests/mypy/kafka.py @@ -1,5 +1,5 @@ import asyncio -from collections.abc import Awaitable, Callable, Coroutine +from collections.abc import Awaitable, Callable from typing import Any import prometheus_client @@ -424,10 +424,11 @@ def check_publisher_instance_type( async def check_call_assertions_take_the_kafka_fields( broker: KafkaBroker | FastAPIRouter | KafkaRouter, ) -> None: + # A sync handler: mypy and pyright spell an `async def`'s return type differently @broker.subscriber("test") - async def handle() -> None: ... + def handle() -> None: ... - assert_type(handle, KafkaHandlerCallWrapper[[], Coroutine[Any, Any, None]]) + assert_type(handle, KafkaHandlerCallWrapper[[], None]) await handle.assert_called_once_with(None, key=b"k", partition=0) await handle.assert_called_with(key=b"k") await handle.assert_any_call(partition=0) @@ -435,9 +436,9 @@ async def handle() -> None: ... publisher = broker.publisher("test") @publisher - async def published() -> None: ... + def published() -> None: ... - assert_type(published, KafkaHandlerCallWrapper[[], Coroutine[Any, Any, None]]) + assert_type(published, KafkaHandlerCallWrapper[[], None]) await publisher.assert_called_once_with(None, key=b"k", partition=0) await publisher.assert_called_with(key=b"k") await publisher.assert_any_call(partition=0) From 0c6e0f020ad53184402b6d02794666e5c9e2554e Mon Sep 17 00:00:00 2001 From: Lancetnik Date: Wed, 9 Sep 2026 12:26:22 +0300 Subject: [PATCH 6/8] refactor(testing): the Kafka reader is a ClassVar built by `field_reader` Both Kafka packages built the same reader by hand: an isinstance guard, the same refusal text, and one line that differs (an attribute on aiokafka, a method on confluent). `field_reader(record, read)` in `_internal/kafka` builds it once, so each `call_wrapper` module opens with its class and binds the reader inline. The mixin declares the reader as a `ClassVar` and looks it up on the class, replacing a staticmethod stub that only raised; a `LogicPublisher` takes the reader from its wrapper class instead of restating the import. The refusal names the message class a shared handler got (`RabbitMessage`), not the raw client type. `ExpectedCall.message_fields()` introspects the dataclass again instead of listing its fields by hand. The shared-handler tests carry `require_aiopika`, as the other cross-broker tests do. Co-Authored-By: Claude Fable 5.1 --- faststream/_internal/kafka/__init__.py | 3 ++ faststream/_internal/kafka/calls.py | 40 ++++++++++++++++----- faststream/_internal/testing/calls.py | 11 ++---- faststream/confluent/call_wrapper.py | 24 ++----------- faststream/confluent/publisher/usecase.py | 6 ++-- faststream/kafka/call_wrapper.py | 22 ++---------- faststream/kafka/publisher/usecase.py | 6 ++-- tests/brokers/confluent/test_test_client.py | 2 ++ tests/brokers/kafka/test_test_client.py | 2 ++ tests/mypy/confluent.py | 2 ++ tests/mypy/kafka.py | 2 ++ 11 files changed, 56 insertions(+), 64 deletions(-) diff --git a/faststream/_internal/kafka/__init__.py b/faststream/_internal/kafka/__init__.py index 632fe0127f9..ea105cb6e5a 100644 --- a/faststream/_internal/kafka/__init__.py +++ b/faststream/_internal/kafka/__init__.py @@ -5,6 +5,7 @@ rather than in one of them or in a cross-broker module. """ +from .calls import KafkaCallAssertions, field_reader from .keys import ( extract_per_message_keys_and_bodies, key_for_index, @@ -12,7 +13,9 @@ ) __all__ = ( + "KafkaCallAssertions", "extract_per_message_keys_and_bodies", + "field_reader", "key_for_index", "realign_keys", ) diff --git a/faststream/_internal/kafka/calls.py b/faststream/_internal/kafka/calls.py index bf8a986e3e8..113195ae28a 100644 --- a/faststream/_internal/kafka/calls.py +++ b/faststream/_internal/kafka/calls.py @@ -1,12 +1,14 @@ -from collections.abc import Mapping -from typing import TYPE_CHECKING, Any +from collections.abc import Callable, Mapping +from typing import TYPE_CHECKING, Any, ClassVar from faststream._internal.constants import EMPTY from faststream._internal.testing.calls import ( BrokerFields, CallAssertions, ExpectedCall, + FieldReader, ) +from faststream.exceptions import SetupError if TYPE_CHECKING: from faststream.message import StreamMessage @@ -17,6 +19,9 @@ class KafkaCallAssertions(CallAssertions): __slots__ = () + # Each Kafka package binds the reader of its client's record, see `field_reader` + _read_field: ClassVar[FieldReader] + async def assert_called_once_with( self, body: Any = EMPTY, @@ -140,11 +145,6 @@ async def assert_any_call( ) ) - @staticmethod - def _read_field(name: str, message: "StreamMessage[Any]") -> Any: - """Take a Kafka field off the raw message; each Kafka package knows its client's.""" - raise NotImplementedError - def _expected_call( self, body: Any = EMPTY, @@ -173,6 +173,30 @@ def _expected_call( for name, value in (("key", key), ("partition", partition)) if value is not EMPTY }, - self._read_field, + type(self)._read_field, ), ) + + +def field_reader( + record: type, + read: Callable[[Any, str], Any], +) -> FieldReader: + """The reader of a Kafka client's record: `read(raw, name)` once the record is its. + + Args: + record: The client's record class; a raw message of another class is refused. + read: Takes the field off the record, under the name `publish()` gives it. + """ + + def read_field(name: str, message: "StreamMessage[Any]") -> Any: + raw = message.raw_message + if not isinstance(raw, record): + msg = ( + f"`{name}` is a Kafka field, and this " + f"`{type(message).__name__}` did not come from Kafka." + ) + raise SetupError(msg) + return read(raw, name) + + return read_field diff --git a/faststream/_internal/testing/calls.py b/faststream/_internal/testing/calls.py index 6edeeb85cb0..93f5da1f509 100644 --- a/faststream/_internal/testing/calls.py +++ b/faststream/_internal/testing/calls.py @@ -1,6 +1,6 @@ from collections.abc import Callable, Mapping from copy import copy -from dataclasses import dataclass, field +from dataclasses import dataclass, field, fields from typing import TYPE_CHECKING, Any, TypeAlias from unittest.mock import MagicMock @@ -309,14 +309,7 @@ class ExpectedCall: def message_fields(self) -> list[str]: """The names asked of the message beside its body.""" - named = ( - "headers", - "correlation_id", - "reply_to", - "content_type", - "path", - "context", - ) + named = (f.name for f in fields(self) if f.name not in {"body", "broker_fields"}) return [ *(name for name in named if getattr(self, name) is not EMPTY), *(self.broker_fields.values if self.broker_fields else ()), diff --git a/faststream/confluent/call_wrapper.py b/faststream/confluent/call_wrapper.py index c8003255d4a..07d51c95dd6 100644 --- a/faststream/confluent/call_wrapper.py +++ b/faststream/confluent/call_wrapper.py @@ -1,27 +1,8 @@ -from typing import TYPE_CHECKING, Any - from confluent_kafka import Message from faststream._internal.endpoint.call_wrapper import HandlerCallWrapper -from faststream._internal.kafka.calls import KafkaCallAssertions +from faststream._internal.kafka import KafkaCallAssertions, field_reader from faststream._internal.types import P_HandlerParams, T_HandlerReturn -from faststream.exceptions import SetupError - -if TYPE_CHECKING: - from faststream.message import StreamMessage - - -def read_field(name: str, message: "StreamMessage[Any]") -> Any: - """Take a Kafka field off a confluent message, under the name `publish()` gives it.""" - raw = message.raw_message - if not isinstance(raw, Message): - msg = ( - f"`{name}` is a Kafka field, and this message did not come from Kafka: " - f"its raw message is a `{type(raw).__name__}`." - ) - raise SetupError(msg) - # The confluent client answers with methods where aiokafka has attributes - return getattr(raw, name)() class KafkaHandlerCallWrapper( @@ -32,4 +13,5 @@ class KafkaHandlerCallWrapper( __slots__ = () - _read_field = staticmethod(read_field) + # The confluent client answers with methods where aiokafka has attributes + _read_field = field_reader(Message, lambda raw, name: getattr(raw, name)()) diff --git a/faststream/confluent/publisher/usecase.py b/faststream/confluent/publisher/usecase.py index 59335bfb68c..278fc4d2a2b 100644 --- a/faststream/confluent/publisher/usecase.py +++ b/faststream/confluent/publisher/usecase.py @@ -9,9 +9,9 @@ PublisherSpecification, PublisherUsecase, ) -from faststream._internal.kafka.calls import KafkaCallAssertions +from faststream._internal.kafka import KafkaCallAssertions from faststream._internal.types import P_HandlerParams, T_HandlerReturn -from faststream.confluent.call_wrapper import KafkaHandlerCallWrapper, read_field +from faststream.confluent.call_wrapper import KafkaHandlerCallWrapper from faststream.confluent.response import KafkaPublishCommand from faststream.response.publish_type import PublishType @@ -29,7 +29,7 @@ class LogicPublisher(KafkaCallAssertions, PublisherUsecase): """A class to publish messages to a Kafka topic.""" _call_wrapper_class = KafkaHandlerCallWrapper - _read_field = staticmethod(read_field) + _read_field = KafkaHandlerCallWrapper._read_field def __init__( self, diff --git a/faststream/kafka/call_wrapper.py b/faststream/kafka/call_wrapper.py index 945c97eaa4e..1d54ff03a3f 100644 --- a/faststream/kafka/call_wrapper.py +++ b/faststream/kafka/call_wrapper.py @@ -1,26 +1,8 @@ -from typing import TYPE_CHECKING, Any - from aiokafka import ConsumerRecord from faststream._internal.endpoint.call_wrapper import HandlerCallWrapper -from faststream._internal.kafka.calls import KafkaCallAssertions +from faststream._internal.kafka import KafkaCallAssertions, field_reader from faststream._internal.types import P_HandlerParams, T_HandlerReturn -from faststream.exceptions import SetupError - -if TYPE_CHECKING: - from faststream.message import StreamMessage - - -def read_field(name: str, message: "StreamMessage[Any]") -> Any: - """Take a Kafka field off an aiokafka record, under the name `publish()` gives it.""" - raw = message.raw_message - if not isinstance(raw, ConsumerRecord): - msg = ( - f"`{name}` is a Kafka field, and this message did not come from Kafka: " - f"its raw message is a `{type(raw).__name__}`." - ) - raise SetupError(msg) - return getattr(raw, name) class KafkaHandlerCallWrapper( @@ -31,4 +13,4 @@ class KafkaHandlerCallWrapper( __slots__ = () - _read_field = staticmethod(read_field) + _read_field = field_reader(ConsumerRecord, getattr) diff --git a/faststream/kafka/publisher/usecase.py b/faststream/kafka/publisher/usecase.py index 7427f5f83ea..b9604062b9b 100644 --- a/faststream/kafka/publisher/usecase.py +++ b/faststream/kafka/publisher/usecase.py @@ -4,9 +4,9 @@ from typing_extensions import override from faststream._internal.endpoint.publisher import PublisherUsecase -from faststream._internal.kafka.calls import KafkaCallAssertions +from faststream._internal.kafka import KafkaCallAssertions from faststream._internal.types import P_HandlerParams, T_HandlerReturn -from faststream.kafka.call_wrapper import KafkaHandlerCallWrapper, read_field +from faststream.kafka.call_wrapper import KafkaHandlerCallWrapper from faststream.kafka.response import KafkaPublishCommand from faststream.response.publish_type import PublishType @@ -29,7 +29,7 @@ class LogicPublisher(KafkaCallAssertions, PublisherUsecase): """A class to publish messages to a Kafka topic.""" _call_wrapper_class = KafkaHandlerCallWrapper - _read_field = staticmethod(read_field) + _read_field = KafkaHandlerCallWrapper._read_field def __init__( self, diff --git a/tests/brokers/confluent/test_test_client.py b/tests/brokers/confluent/test_test_client.py index a3da396d714..ee38b828f08 100644 --- a/tests/brokers/confluent/test_test_client.py +++ b/tests/brokers/confluent/test_test_client.py @@ -10,6 +10,7 @@ from faststream.confluent.testing import FakeProducer from faststream.exceptions import SetupError from tests.brokers.base.testclient import BrokerTestclientTestcase +from tests.marks import require_aiopika from tests.tools import spy_decorator from .basic import ConfluentMemoryTestcaseConfig @@ -193,6 +194,7 @@ async def handle(msg) -> None: await publisher.assert_called_with(key=b"k") await publisher.assert_any_call(partition=1) + @require_aiopika async def test_kafka_fields_refuse_another_brokers_message(self, queue: str) -> None: from faststream.rabbit import RabbitBroker, TestRabbitBroker diff --git a/tests/brokers/kafka/test_test_client.py b/tests/brokers/kafka/test_test_client.py index 04a72eef965..e6b4398ac0b 100644 --- a/tests/brokers/kafka/test_test_client.py +++ b/tests/brokers/kafka/test_test_client.py @@ -11,6 +11,7 @@ from faststream.kafka.message import FAKE_CONSUMER from faststream.kafka.testing import FakeProducer from tests.brokers.base.testclient import BrokerTestclientTestcase +from tests.marks import require_aiopika from tests.tools import spy_decorator from .basic import KafkaMemoryTestcaseConfig @@ -245,6 +246,7 @@ async def handle(msg) -> None: await publisher.assert_called_with(key=b"k") await publisher.assert_any_call(partition=1) + @require_aiopika async def test_kafka_fields_refuse_another_brokers_message(self, queue: str) -> None: from faststream.rabbit import RabbitBroker, TestRabbitBroker diff --git a/tests/mypy/confluent.py b/tests/mypy/confluent.py index 32a25932da7..4936ffc32d9 100644 --- a/tests/mypy/confluent.py +++ b/tests/mypy/confluent.py @@ -138,6 +138,8 @@ def handle() -> None: ... await handle.assert_called_with(key=b"k") await handle.assert_any_call(partition=0) + # The publisher's own type is pinned in `check_publisher_instance_type`; here its + # methods take the two fields, which only the Kafka mixin gives them publisher = broker.publisher("test") @publisher diff --git a/tests/mypy/kafka.py b/tests/mypy/kafka.py index ebe1e134654..3432ad19b4f 100644 --- a/tests/mypy/kafka.py +++ b/tests/mypy/kafka.py @@ -433,6 +433,8 @@ def handle() -> None: ... await handle.assert_called_with(key=b"k") await handle.assert_any_call(partition=0) + # The publisher's own type is pinned in `check_publisher_instance_type`; here its + # methods take the two fields, which only the Kafka mixin gives them publisher = broker.publisher("test") @publisher From dbeb985681f97bb7a21534c860c8d3e737bd242d Mon Sep 17 00:00:00 2001 From: Lancetnik Date: Wed, 23 Sep 2026 21:22:18 +0300 Subject: [PATCH 7/8] test(kafka): publish straight from the publisher in its assertion test The subscriber that forwarded the message added nothing: the publisher's mock records the call either way, and the test still goes red when the reader returns the wrong field. Co-Authored-By: Claude Opus 5.5 (1M context) --- tests/brokers/confluent/test_test_client.py | 8 ++------ tests/brokers/kafka/test_test_client.py | 8 ++------ 2 files changed, 4 insertions(+), 12 deletions(-) diff --git a/tests/brokers/confluent/test_test_client.py b/tests/brokers/confluent/test_test_client.py index ee38b828f08..4c897e00664 100644 --- a/tests/brokers/confluent/test_test_client.py +++ b/tests/brokers/confluent/test_test_client.py @@ -181,15 +181,11 @@ async def handle(msg) -> None: ... async def test_publisher_assertions_take_the_kafka_fields(self, queue: str) -> None: broker = self.get_broker() - publisher = broker.publisher(queue + "2") + publisher = broker.publisher(queue) - @broker.subscriber(queue) - async def handle(msg) -> None: + async with self.patch_broker(broker): await publisher.publish("response", key=b"k", partition=1) - async with self.patch_broker(broker) as br: - await br.publish("hello", queue) - await publisher.assert_called_once_with("response", key=b"k", partition=1) await publisher.assert_called_with(key=b"k") await publisher.assert_any_call(partition=1) diff --git a/tests/brokers/kafka/test_test_client.py b/tests/brokers/kafka/test_test_client.py index e6b4398ac0b..50f39d0540d 100644 --- a/tests/brokers/kafka/test_test_client.py +++ b/tests/brokers/kafka/test_test_client.py @@ -233,15 +233,11 @@ async def handle(msg) -> None: ... async def test_publisher_assertions_take_the_kafka_fields(self, queue: str) -> None: broker = self.get_broker() - publisher = broker.publisher(queue + "2") + publisher = broker.publisher(queue) - @broker.subscriber(queue) - async def handle(msg) -> None: + async with self.patch_broker(broker): await publisher.publish("response", key=b"k", partition=1) - async with self.patch_broker(broker) as br: - await br.publish("hello", queue) - await publisher.assert_called_once_with("response", key=b"k", partition=1) await publisher.assert_called_with(key=b"k") await publisher.assert_any_call(partition=1) From ce27d18187a3e1fdab175c96265377dc57682a51 Mon Sep 17 00:00:00 2001 From: Lancetnik Date: Wed, 23 Sep 2026 23:25:44 +0300 Subject: [PATCH 8/8] fix(testing): keep the Kafka fields green on the updated main - The Kafka and confluent `__call__` overrides take `dependencies` as a `Sequence`, as `SubscriberUsecase.__call__` does since #3192. - The tests type their handlers (#3176 put `tests/brokers` under strict mypy). The shared-handler test subscribes RabbitMQ by a call rather than a second decorator, so the name keeps the Kafka wrapper's type; the RabbitMQ import stays inside the test (`# noqa: PLC0415`) because the test is `@require_aiopika`. - `MockConfluentMessage` answers `len()` as the client's message does: its `Message` base read a C-level value the mock never sets, so every mock message was falsy and `assert self.msg` in `tests/brokers/confluent/test_requests.py` failed. Co-Authored-By: Claude Opus 5.5 (1M context) --- faststream/confluent/subscriber/usecase.py | 8 ++++---- faststream/confluent/testing.py | 4 ++++ faststream/kafka/subscriber/usecase.py | 8 ++++---- tests/brokers/confluent/test_test_client.py | 11 ++++++----- tests/brokers/kafka/test_test_client.py | 11 ++++++----- 5 files changed, 24 insertions(+), 18 deletions(-) diff --git a/faststream/confluent/subscriber/usecase.py b/faststream/confluent/subscriber/usecase.py index a8bc77dd683..47fdeb6be8b 100644 --- a/faststream/confluent/subscriber/usecase.py +++ b/faststream/confluent/subscriber/usecase.py @@ -1,6 +1,6 @@ import logging from abc import abstractmethod -from collections.abc import AsyncIterator, Callable, Iterable, Sequence +from collections.abc import AsyncIterator, Callable, Sequence from typing import ( TYPE_CHECKING, Any, @@ -84,7 +84,7 @@ def __call__( filter: "Filter[Any]" = default_filter, parser: Optional["CustomCallable"] = None, decoder: Optional["CustomCallable"] = None, - dependencies: Iterable["Dependant"] = (), + dependencies: Sequence["Dependant"] = (), ) -> "KafkaHandlerCallWrapper[P_HandlerParams, T_HandlerReturn]": ... @overload @@ -95,7 +95,7 @@ def __call__( filter: "Filter[Any]" = default_filter, parser: Optional["CustomCallable"] = None, decoder: Optional["CustomCallable"] = None, - dependencies: Iterable["Dependant"] = (), + dependencies: Sequence["Dependant"] = (), ) -> Callable[ [Callable[P_HandlerParams, T_HandlerReturn]], "KafkaHandlerCallWrapper[P_HandlerParams, T_HandlerReturn]", @@ -109,7 +109,7 @@ def __call__( filter: "Filter[Any]" = default_filter, parser: Optional["CustomCallable"] = None, decoder: Optional["CustomCallable"] = None, - dependencies: Iterable["Dependant"] = (), + dependencies: Sequence["Dependant"] = (), ) -> Union[ "KafkaHandlerCallWrapper[P_HandlerParams, T_HandlerReturn]", Callable[ diff --git a/faststream/confluent/testing.py b/faststream/confluent/testing.py index 9a2aae574c5..532b386cc2f 100644 --- a/faststream/confluent/testing.py +++ b/faststream/confluent/testing.py @@ -319,6 +319,10 @@ def __init__( def len(self) -> int: return 0 if self._raw_msg is None else len(self._raw_msg) + # `Message.__len__` reads the C-level value this mock never sets + def __len__(self) -> int: + return self.len() + def error(self) -> KafkaError | None: return self._error diff --git a/faststream/kafka/subscriber/usecase.py b/faststream/kafka/subscriber/usecase.py index 420320fb4f3..a942dbfdc06 100644 --- a/faststream/kafka/subscriber/usecase.py +++ b/faststream/kafka/subscriber/usecase.py @@ -1,6 +1,6 @@ import logging from abc import abstractmethod -from collections.abc import AsyncIterator, Callable, Iterable, Sequence +from collections.abc import AsyncIterator, Callable, Sequence from itertools import chain from typing import TYPE_CHECKING, Any, Optional, Union, cast @@ -91,7 +91,7 @@ def __call__( filter: "Filter[Any]" = default_filter, parser: Optional["CustomCallable"] = None, decoder: Optional["CustomCallable"] = None, - dependencies: Iterable["Dependant"] = (), + dependencies: Sequence["Dependant"] = (), ) -> "KafkaHandlerCallWrapper[P_HandlerParams, T_HandlerReturn]": ... @overload @@ -102,7 +102,7 @@ def __call__( filter: "Filter[Any]" = default_filter, parser: Optional["CustomCallable"] = None, decoder: Optional["CustomCallable"] = None, - dependencies: Iterable["Dependant"] = (), + dependencies: Sequence["Dependant"] = (), ) -> Callable[ [Callable[P_HandlerParams, T_HandlerReturn]], "KafkaHandlerCallWrapper[P_HandlerParams, T_HandlerReturn]", @@ -116,7 +116,7 @@ def __call__( filter: "Filter[Any]" = default_filter, parser: Optional["CustomCallable"] = None, decoder: Optional["CustomCallable"] = None, - dependencies: Iterable["Dependant"] = (), + dependencies: Sequence["Dependant"] = (), ) -> Union[ "KafkaHandlerCallWrapper[P_HandlerParams, T_HandlerReturn]", Callable[ diff --git a/tests/brokers/confluent/test_test_client.py b/tests/brokers/confluent/test_test_client.py index b8eff013d41..85666fd024c 100644 --- a/tests/brokers/confluent/test_test_client.py +++ b/tests/brokers/confluent/test_test_client.py @@ -168,7 +168,7 @@ async def test_assertions_take_the_kafka_fields(self, queue: str) -> None: broker = self.get_broker() @broker.subscriber(queue) - async def handle(msg) -> None: ... + async def handle(msg: str) -> None: ... async with self.patch_broker(broker) as br: await br.publish("hello", queue, key=b"k", partition=1) @@ -194,15 +194,16 @@ async def test_publisher_assertions_take_the_kafka_fields(self, queue: str) -> N @require_aiopika async def test_kafka_fields_refuse_another_brokers_message(self, queue: str) -> None: - from faststream.rabbit import RabbitBroker, TestRabbitBroker + from faststream.rabbit import RabbitBroker, TestRabbitBroker # noqa: PLC0415 broker = self.get_broker() rabbit = RabbitBroker() - # The first decorator decides the wrapper class: Kafka's here - @rabbit.subscriber(queue) @broker.subscriber(queue) - async def handle(msg) -> None: ... + async def handle(msg: str) -> None: ... + + # The broker that wraps a handler first decides its wrapper class: Kafka's here + _ = rabbit.subscriber(queue)(handle) async with self.patch_broker(broker), TestRabbitBroker(rabbit): await rabbit.publish("hello", queue) diff --git a/tests/brokers/kafka/test_test_client.py b/tests/brokers/kafka/test_test_client.py index b5637e51c19..733f9620aa4 100644 --- a/tests/brokers/kafka/test_test_client.py +++ b/tests/brokers/kafka/test_test_client.py @@ -220,7 +220,7 @@ async def test_assertions_take_the_kafka_fields(self, queue: str) -> None: broker = self.get_broker() @broker.subscriber(queue) - async def handle(msg) -> None: ... + async def handle(msg: str) -> None: ... async with self.patch_broker(broker) as br: await br.publish("hello", queue, key=b"k", partition=1) @@ -246,15 +246,16 @@ async def test_publisher_assertions_take_the_kafka_fields(self, queue: str) -> N @require_aiopika async def test_kafka_fields_refuse_another_brokers_message(self, queue: str) -> None: - from faststream.rabbit import RabbitBroker, TestRabbitBroker + from faststream.rabbit import RabbitBroker, TestRabbitBroker # noqa: PLC0415 broker = self.get_broker() rabbit = RabbitBroker() - # The first decorator decides the wrapper class: Kafka's here - @rabbit.subscriber(queue) @broker.subscriber(queue) - async def handle(msg) -> None: ... + async def handle(msg: str) -> None: ... + + # The broker that wraps a handler first decides its wrapper class: Kafka's here + _ = rabbit.subscriber(queue)(handle) async with self.patch_broker(broker), TestRabbitBroker(rabbit): await rabbit.publish("hello", queue)