feat(testing): Kafka handlers and publishers assert on key and partition

docs/docs/en/getting-started/subscription/test.md

@@ -311,7 +311,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"
@@ -339,11 +339,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"
@@ -374,6 +374,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.
 

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",
+        )

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",
+        )

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]):

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
 
@@ -18,6 +18,10 @@ class Endpoint:
     # few per application. `Registrator` holds them in a WeakSet, hence `__weakref__`.
     __slots__ = ("__dict__", "__weakref__", "_outer_config")
 
+    # 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
 
@@ -26,6 +30,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

faststream/_internal/kafka/__init__.py

@@ -5,14 +5,17 @@
 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,
     realign_keys,
 )
 
 __all__ = (
+    "KafkaCallAssertions",
     "extract_per_message_keys_and_bodies",
+    "field_reader",
     "key_for_index",
     "realign_keys",
 )

faststream/_internal/kafka/calls.py

@@ -0,0 +1,202 @@
+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
+
+
+class KafkaCallAssertions(CallAssertions):
+    """The Call assertions of a Kafka endpoint: the message fields, `key` and `partition`."""
+
+    __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,
+        /,
+        *,
+        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._expected_call(
+                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._expected_call(
+                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._expected_call(
+                body,
+                headers=headers,
+                correlation_id=correlation_id,
+                reply_to=reply_to,
+                content_type=content_type,
+                path=path,
+                context=context,
+                key=key,
+                partition=partition,
+            )
+        )
+
+    def _expected_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,
+    ) -> 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=BrokerFields(
+                {
+                    name: value
+                    for name, value in (("key", key), ("partition", partition))
+                    if value is not EMPTY
+                },
+                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