AsyncAPI 2.6.0 fix empty channels for KafkaSubscriber and ConfluentSu…

…bscriber if partitions provided (#1930)

* AsyncAPI 2.6.0 fix empty channels for KafkaSubscriber if partitions provided

* Remove unused import

* AsyncAPI 2.6.0 fix empty channels for ConfluentSubscriber if partitions provided

* Fix formatting errors

* chore: fix precommit

---------

Co-authored-by: Nikita Pastukhov <[email protected]>

faststream/confluent/subscriber/asyncapi.py

@@ -1,3 +1,4 @@
+from itertools import chain
 from typing import (
     TYPE_CHECKING,
     Dict,
@@ -34,7 +35,10 @@ def get_schema(self) -> Dict[str, Channel]:
         channels = {}
 
         payloads = self.get_payloads()
-        for t in self.topics:
+
+        topics = chain(self.topics, {part.topic for part in self.partitions})
+
+        for t in topics:
             handler_name = self.title_ or f"{t}:{self.call_name}"
 
             channels[handler_name] = Channel(

faststream/kafka/subscriber/asyncapi.py

@@ -1,3 +1,4 @@
+from itertools import chain
 from typing import (
     TYPE_CHECKING,
     Dict,
@@ -35,7 +36,9 @@ def get_schema(self) -> Dict[str, Channel]:
 
         payloads = self.get_payloads()
 
-        for t in self.topics:
+        topics = chain(self.topics, {part.topic for part in self.partitions})
+
+        for t in topics:
             handler_name = self.title_ or f"{t}:{self.call_name}"
 
             channels[handler_name] = Channel(

tests/asyncapi/confluent/test_arguments.py

@@ -1,5 +1,5 @@
 from faststream.asyncapi.generate import get_app_schema
-from faststream.confluent import KafkaBroker
+from faststream.confluent import KafkaBroker, TopicPartition
 from tests.asyncapi.base.arguments import ArgumentsTestcase
 
 
@@ -18,3 +18,39 @@ async def handle(msg): ...
         assert schema["channels"][key]["bindings"] == {
             "kafka": {"bindingVersion": "0.4.0", "topic": "test"}
         }
+
+    def test_subscriber_with_one_topic_partitions(self):
+        broker = self.broker_class()
+
+        part1 = TopicPartition("topic_name", 1)
+        part2 = TopicPartition("topic_name", 2)
+
+        @broker.subscriber(partitions=[part1, part2])
+        async def handle(msg): ...
+
+        schema = get_app_schema(self.build_app(broker)).to_jsonable()
+        key = tuple(schema["channels"].keys())[0]  # noqa: RUF015
+
+        assert schema["channels"][key]["bindings"] == {
+            "kafka": {"bindingVersion": "0.4.0", "topic": "topic_name"}
+        }
+
+    def test_subscriber_with_multi_topics_partitions(self):
+        broker = self.broker_class()
+
+        part1 = TopicPartition("topic_name1", 1)
+        part2 = TopicPartition("topic_name2", 2)
+
+        @broker.subscriber(partitions=[part1, part2])
+        async def handle(msg): ...
+
+        schema = get_app_schema(self.build_app(broker)).to_jsonable()
+        key1 = tuple(schema["channels"].keys())[0]  # noqa: RUF015
+        key2 = tuple(schema["channels"].keys())[1]
+
+        assert sorted(
+            (
+                schema["channels"][key1]["bindings"]["kafka"]["topic"],
+                schema["channels"][key2]["bindings"]["kafka"]["topic"],
+            )
+        ) == sorted(("topic_name1", "topic_name2"))

tests/asyncapi/kafka/test_arguments.py

@@ -1,3 +1,5 @@
+from aiokafka import TopicPartition
+
 from faststream.asyncapi.generate import get_app_schema
 from faststream.kafka import KafkaBroker
 from tests.asyncapi.base.arguments import ArgumentsTestcase
@@ -18,3 +20,39 @@ async def handle(msg): ...
         assert schema["channels"][key]["bindings"] == {
             "kafka": {"bindingVersion": "0.4.0", "topic": "test"}
         }
+
+    def test_subscriber_with_one_topic_partitions(self):
+        broker = self.broker_class()
+
+        part1 = TopicPartition("topic_name", 1)
+        part2 = TopicPartition("topic_name", 2)
+
+        @broker.subscriber(partitions=[part1, part2])
+        async def handle(msg): ...
+
+        schema = get_app_schema(self.build_app(broker)).to_jsonable()
+        key = tuple(schema["channels"].keys())[0]  # noqa: RUF015
+
+        assert schema["channels"][key]["bindings"] == {
+            "kafka": {"bindingVersion": "0.4.0", "topic": "topic_name"}
+        }
+
+    def test_subscriber_with_multi_topics_partitions(self):
+        broker = self.broker_class()
+
+        part1 = TopicPartition("topic_name1", 1)
+        part2 = TopicPartition("topic_name2", 2)
+
+        @broker.subscriber(partitions=[part1, part2])
+        async def handle(msg): ...
+
+        schema = get_app_schema(self.build_app(broker)).to_jsonable()
+        key1 = tuple(schema["channels"].keys())[0]  # noqa: RUF015
+        key2 = tuple(schema["channels"].keys())[1]
+
+        assert sorted(
+            (
+                schema["channels"][key1]["bindings"]["kafka"]["topic"],
+                schema["channels"][key2]["bindings"]["kafka"]["topic"],
+            )
+        ) == sorted(("topic_name1", "topic_name2"))