Change broker and tests

Author: chandr-andr

poetry.lock

@@ -68,18 +68,6 @@ def __init__(  # noqa: WPS211
         """
         super().__init__(result_backend, task_id_generator)
 
-        self._bootstrap_servers: Optional[Union[str, List[str]]] = bootstrap_servers
-        self._kafka_topic: Optional[NewTopic] = kafka_topic
-        self._aiokafka_producer: Optional[AIOKafkaProducer] = aiokafka_producer
-        self._aiokafka_consumer: Optional[AIOKafkaConsumer] = aiokafka_consumer
-        self._kafka_admin_client: Optional[KafkaAdminClient] = kafka_admin_client
-        self._loop: Optional[asyncio.AbstractEventLoop] = loop
-
-        self._default_kafka_topic: str = "taskiq_topic"
-        self._delete_topic_on_shutdown: bool = delete_topic_on_shutdown
-
-        self._delay_kick_tasks: Set[asyncio.Task[None]] = set()
-
         if (aiokafka_producer or aiokafka_consumer) and not bootstrap_servers:
             raise WrongAioKafkaBrokerParametersError(
                 (
@@ -88,52 +76,67 @@ def __init__(  # noqa: WPS211
                 ),
             )
 
-    async def startup(self) -> None:  # noqa: C901
-        """Setup AIOKafkaProducer, AIOKafkaConsumer and kafka topics.
-
-        We will have 2 topics for default and high priority.
+        self._bootstrap_servers: Optional[Union[str, List[str]]] = bootstrap_servers
 
-        Also we need to create AIOKafkaProducer and AIOKafkaConsumer
-        if there are no producer and consumer passed.
-        """
-        await super().startup()
+        self._kafka_topic: NewTopic = kafka_topic or NewTopic(
+            name="taskiq_topic",
+            num_partitions=1,
+            replication_factor=1,
+        )
 
-        if not self._aiokafka_producer:
-            self._aiokafka_producer = AIOKafkaProducer(
+        self._aiokafka_producer: AIOKafkaProducer = (
+            aiokafka_producer
+            or AIOKafkaProducer(
                 bootstrap_servers=self._bootstrap_servers,
                 loop=self._loop,
             )
+        )
 
-        if not self._kafka_topic:
-            self._kafka_topic: NewTopic = NewTopic(  # type: ignore
-                name=self._default_kafka_topic,
-                num_partitions=1,
-                replication_factor=1,
+        self._aiokafka_consumer: AIOKafkaConsumer = (
+            aiokafka_consumer
+            or AIOKafkaConsumer(
+                self._kafka_topic.name,
+                bootstrap_servers=self._bootstrap_servers,
+                loop=self._loop,
             )
+        )
 
-        if not self._kafka_admin_client:
-            self._kafka_admin_client: KafkaAdminClient = (  # type: ignore
-                KafkaAdminClient(
-                    bootstrap_servers=self._bootstrap_servers,
-                    client_id="kafka-python-taskiq",
-                )
+        self._kafka_admin_client: KafkaAdminClient = (
+            kafka_admin_client
+            or KafkaAdminClient(
+                bootstrap_servers=self._bootstrap_servers,
+                client_id="kafka-python-taskiq",
             )
+        )
+        self._loop: Optional[asyncio.AbstractEventLoop] = loop
+
+        self._delete_topic_on_shutdown: bool = delete_topic_on_shutdown
+
+        self._delay_kick_tasks: Set[asyncio.Task[None]] = set()
+
+        self._is_started = False
+
+    async def startup(self) -> None:  # noqa: C901
+        """Setup AIOKafkaProducer, AIOKafkaConsumer and kafka topics.
+
+        We will have 2 topics for default and high priority.
+
+        Also we need to create AIOKafkaProducer and AIOKafkaConsumer
+        if there are no producer and consumer passed.
+        """
+        await super().startup()
 
         if self._kafka_topic.name not in self._kafka_admin_client.list_topics():
             self._kafka_admin_client.create_topics(
                 new_topics=[self._kafka_topic],
                 validate_only=False,
             )
 
-        if not self._aiokafka_consumer:
-            self._aiokafka_consumer = AIOKafkaConsumer(
-                self._kafka_topic.name,
-                bootstrap_servers=self._bootstrap_servers,
-                loop=self._loop,
-            )
-
         await self._aiokafka_producer.start()
-        await self._aiokafka_consumer.start()
+        if self.is_worker_process:
+            await self._aiokafka_consumer.start()
+
+        self._is_started = True
 
     async def shutdown(self) -> None:
         """Close all connections on shutdown."""
@@ -171,34 +174,20 @@ async def kick(self, message: BrokerMessage) -> None:
         :raises ValueError: if startup wasn't called.
         :param message: message to send.
         """
-        if not self._aiokafka_producer:
-            raise ValueError("Specify aiokafka_producer or run startup before kicking.")
+        if not self._is_started:
+            raise ValueError("Please run startup before kicking.")
 
         kafka_message: bytes = pickle.dumps(message)
-        topic_name: str = self._kafka_topic.name  # type: ignore
+        topic_name: str = self._kafka_topic.name
 
-        delay = parse_val(int, message.labels.get("delay"))
-        if delay is None:
-            await self._aiokafka_producer.send(  # type: ignore
-                topic=topic_name,
-                value=kafka_message,
-            )
-        else:
-            delay_kick_task = asyncio.create_task(
-                self._delay_kick(
-                    kafka_message=kafka_message,
-                    topic_name=topic_name,
-                    delay=delay,
-                ),
-            )
-            self._delay_kick_tasks.add(delay_kick_task)
-            delay_kick_task.add_done_callback(
-                self._delay_kick_tasks.discard,
-            )
+        await self._aiokafka_producer.send(
+            topic=topic_name,
+            value=kafka_message,
+        )
 
     async def listen(
         self,
-    ) -> AsyncGenerator[BrokerMessage, None]:
+    ) -> AsyncGenerator[bytes, None]:
         """Listen to topic.
 
         This function starts listen to topic and
@@ -207,36 +196,8 @@ async def listen(
         :yields: parsed broker message.
         :raises ValueError: if no aiokafka_consumer or startup wasn't called.
         """
-        if not self._aiokafka_consumer:
-            raise ValueError("Specify aiokafka_consumer or run startup before kicking.")
-        async for raw_kafka_message in self._aiokafka_consumer:
-            try:
-                broker_message: BrokerMessage = pickle.loads(  # noqa: S301
-                    raw_kafka_message.value,
-                )
-            except (TypeError, ValueError) as exc:
-                logger.warning(
-                    "Cannot parse message from Kafka %s",
-                    exc,
-                    exc_info=True,
-                )
-                continue
-            yield broker_message
+        if not self._is_started:
+            raise ValueError("Please run startup before listening.")
 
-    async def _delay_kick(
-        self,
-        kafka_message: bytes,
-        topic_name: str,
-        delay: float,
-    ) -> None:
-        """Send message to the topic after delay time.
-
-        :param kafka_message: message to send.
-        :param topic_name: name of the topic.
-        :param delay: delay before kicking in seconds.
-        """
-        await asyncio.sleep(delay=delay)
-        await self._aiokafka_producer.send(  # type: ignore
-            topic=topic_name,
-            value=kafka_message,
-        )
+        async for raw_kafka_message in self._aiokafka_consumer:
+            yield raw_kafka_message.value

taskiq_aio_kafka/broker.py

@@ -1,6 +1,5 @@
-class BaseTaskiqKafkaError(Exception):
-    """Base class for all TaskiqKafka errors."""
+from taskiq.exceptions import BrokerError
 
 
-class WrongAioKafkaBrokerParametersError(BaseTaskiqKafkaError):
+class WrongAioKafkaBrokerParametersError(BrokerError):
     """Error if Producer or Consumer is specified but no bootstrap_servers."""

taskiq_aio_kafka/exceptions.py

@@ -103,7 +103,6 @@ async def broker_without_arguments(
     broker = AioKafkaBroker(
         bootstrap_servers=kafka_url,
     )
-    broker._default_kafka_topic = base_topic_name  # noqa: WPS437
     broker.is_worker_process = True
 
     await broker.startup()

tests/conftest.py

@@ -10,7 +10,7 @@
 from taskiq_aio_kafka.broker import AioKafkaBroker
 
 
-async def get_first_task(broker: AioKafkaBroker) -> BrokerMessage:  # type: ignore
+async def get_first_task(broker: AioKafkaBroker) -> bytes:  # type: ignore
     """Get first message from the topic.
 
     :param broker: async message broker.
@@ -36,7 +36,7 @@ async def test_kick_success(broker: AioKafkaBroker) -> None:
     message_to_send = BrokerMessage(
         task_id=task_id,
         task_name=task_name,
-        message="my_msg",
+        message=pickle.dumps("my_msg"),
         labels={
             "label1": "val1",
         },
@@ -91,7 +91,7 @@ async def test_listen(
     """
     task_id: str = uuid4().hex
     task_name: str = uuid4().hex
-    message: str = uuid4().hex
+    message: bytes = pickle.dumps(uuid4().hex)
     labels: Dict[str, str] = {"test_label": "123"}
 
     await test_kafka_producer.send(
@@ -106,9 +106,11 @@ async def test_listen(
         ),
     )
 
-    broker_message: BrokerMessage = await asyncio.wait_for(
-        get_first_task(broker),
-        timeout=1,
+    broker_message: BrokerMessage = pickle.loads(  # noqa: S301
+        await asyncio.wait_for(
+            get_first_task(broker),
+            timeout=1,
+        ),
     )
 
     assert broker_message.message == message
@@ -132,7 +134,7 @@ async def test_delayed_message(
     broker_msg = BrokerMessage(
         task_id="1",
         task_name="name",
-        message="message",
+        message=pickle.dumps("message"),
         labels={"delay": "3"},
     )