Coverage for faststream / kafka / message.py: 96%

1from typing import Any, Protocol, Union (empty)

2 

3from aiokafka import ( (empty)

4 AIOKafkaConsumer, 

5 ConsumerRecord, 

6 TopicPartition as AIOKafkaTopicPartition, 

7) 

8 

9from faststream.message import AckStatus, StreamMessage (empty)

10 

11 

12class ConsumerProtocol(Protocol): (empty)

13 """A protocol for Kafka consumers.""" 

14 

15 async def commit(self) -> None: ... (empty)

16 

17 def seek( (empty)

18 self, 

19 partition: AIOKafkaTopicPartition, 

20 offset: int, 

21 ) -> None: 

22 pass 

23 

24 

25class FakeConsumer: (empty)

26 """A fake Kafka consumer.""" 

27 

28 async def commit(self) -> None: (empty)

29 pass (empty)

30 

31 def seek( (empty)

32 self, 

33 partition: AIOKafkaTopicPartition, 

34 offset: int, 

35 ) -> None: 

36 pass (empty)

37 

38 

39FAKE_CONSUMER = FakeConsumer() (empty)

40 

41 

42class KafkaRawMessage(ConsumerRecord): # type: ignore[misc] (empty)

43 consumer: AIOKafkaConsumer (empty)

44 

45 

46class KafkaMessage( (empty)

47 StreamMessage[ (empty)

48 Union[ (empty)

49 "ConsumerRecord", (empty)

50 tuple["ConsumerRecord", ...], (empty)

51 ] 

52 ], 

53): 

54 """Represents a Kafka message in the FastStream framework. (empty)

55 

56 This class extends `StreamMessage` and is specialized for handling Kafka ConsumerRecord objects. 

57 """ 

58 

59 def __init__(self, *args: Any, consumer: ConsumerProtocol, **kwargs: Any) -> None: (empty)

60 super().__init__(*args, **kwargs) (empty)

61 self.consumer = consumer (empty)

62 self.committed = AckStatus.ACKED (empty)

63 

64 

65class KafkaAckableMessage(KafkaMessage): (empty)

66 def __init__(self, *args: Any, **kwargs: Any) -> None: (empty)

67 super().__init__(*args, **kwargs) (empty)

68 self.committed = None (empty)

69 

70 async def ack(self) -> None: (empty)

71 if not self.committed: (empty)

72 await self.consumer.commit() (empty)

73 await super().ack() (empty)

74 

75 async def reject(self) -> None: (empty)

76 await self.ack() (empty)

77 

78 async def nack(self) -> None: (empty)

79 if not self.committed: 79 ↛ 93line 79 didn't jump to line 93 because the condition on line 79 was always true(empty)

80 raw_message = ( (empty)

81 self.raw_message[0] 

82 if isinstance(self.raw_message, tuple) 

83 else self.raw_message 

84 ) 

85 topic_partition = AIOKafkaTopicPartition( (empty)

86 raw_message.topic, 

87 raw_message.partition, 

88 ) 

89 self.consumer.seek( (empty)

90 partition=topic_partition, 

91 offset=raw_message.offset, 

92 ) 

93 await super().nack() (empty)