Coverage for faststream / kafka / configs / broker.py: 89%

1from collections.abc import Callable (empty)

2from dataclasses import dataclass, field (empty)

3from functools import partial (empty)

4from typing import Any, Optional (empty)

5 

6import aiokafka (empty)

7import aiokafka.admin (empty)

8 

9from faststream.__about__ import SERVICE_NAME (empty)

10from faststream._internal.configs import BrokerConfig (empty)

11from faststream._internal.utils.data import filter_by_dict (empty)

12from faststream.exceptions import IncorrectState (empty)

13from faststream.kafka.publisher.producer import ( (empty)

14 AioKafkaFastProducer, 

15 FakeAioKafkaFastProducer, 

16) 

17from faststream.kafka.schemas.params import ( (empty)

18 AdminClientConnectionParams, 

19 ConsumerConnectionParams, 

20) 

21 

22 

23@dataclass(kw_only=True) (empty)

24class KafkaBrokerConfig(BrokerConfig): (empty)

25 producer: "AioKafkaFastProducer" = field(default_factory=FakeAioKafkaFastProducer) (empty)

26 builder: Callable[..., aiokafka.AIOKafkaConsumer] = lambda: None (empty)

27 

28 client_id: str | None = SERVICE_NAME (empty)

29 

30 _admin_client: Optional["aiokafka.admin.client.AIOKafkaAdminClient"] = None (empty)

31 

32 @property (empty)

33 def admin_client(self) -> "aiokafka.admin.client.AIOKafkaAdminClient": (empty)

34 if self._admin_client is None: 34 ↛ 35line 34 didn't jump to line 35 because the condition on line 34 was never true(empty)

35 msg = "Admin client is not initialized. Call connect() first." 

36 raise IncorrectState(msg) 

37 

38 return self._admin_client (empty)

39 

40 async def connect(self, **connection_kwargs: Any) -> "None": (empty)

41 producer = aiokafka.AIOKafkaProducer(**connection_kwargs) (empty)

42 await self.producer.connect(producer, serializer=self.fd_config._serializer) (empty)

43 

44 admin_options, _ = filter_by_dict( (empty)

45 AdminClientConnectionParams, 

46 connection_kwargs, 

47 ) 

48 

49 self._admin_client = aiokafka.admin.client.AIOKafkaAdminClient(**admin_options) (empty)

50 await self._admin_client.start() (empty)

51 

52 consumer_options, _ = filter_by_dict( (empty)

53 ConsumerConnectionParams, 

54 connection_kwargs, 

55 ) 

56 self.builder = partial(aiokafka.AIOKafkaConsumer, **consumer_options) (empty)

57 

58 async def disconnect(self) -> "None": (empty)

59 if self._admin_client is not None: (empty)

60 await self._admin_client.close() (empty)

61 self._admin_client = None (empty)

62 

63 await self.producer.disconnect() (empty)