Coverage for faststream / kafka / testing.py: 88%

1import re (empty)

2from collections.abc import Callable, Generator, Iterable, Iterator (empty)

3from contextlib import ExitStack, contextmanager (empty)

4from datetime import datetime, timezone (empty)

5from typing import TYPE_CHECKING, Any, Optional, cast (empty)

6from unittest.mock import AsyncMock, MagicMock (empty)

7 

8import anyio (empty)

9from aiokafka import ConsumerRecord (empty)

10from typing_extensions import override (empty)

11 

12from faststream._internal.endpoint.utils import ParserComposition (empty)

13from faststream._internal.testing.broker import TestBroker, change_producer (empty)

14from faststream.exceptions import SubscriberNotFound (empty)

15from faststream.kafka import TopicPartition (empty)

16from faststream.kafka.broker import KafkaBroker (empty)

17from faststream.kafka.message import KafkaMessage (empty)

18from faststream.kafka.parser import AioKafkaParser (empty)

19from faststream.kafka.publisher.producer import AioKafkaFastProducer (empty)

20from faststream.kafka.publisher.usecase import BatchPublisher (empty)

21from faststream.kafka.subscriber.usecase import BatchSubscriber (empty)

22from faststream.message import encode_message, gen_cor_id (empty)

23 

24if TYPE_CHECKING: (empty)

25 from fast_depends.library.serializer import SerializerProto 

26 

27 from faststream._internal.basic_types import SendableMessage 

28 from faststream.kafka.publisher.usecase import LogicPublisher 

29 from faststream.kafka.response import KafkaPublishCommand 

30 from faststream.kafka.subscriber.usecase import LogicSubscriber 

31 

32__all__ = ("TestKafkaBroker",) (empty)

33 

34 

35class TestKafkaBroker(TestBroker[KafkaBroker]): (empty)

36 """A class to test Kafka brokers.""" 

37 

38 @contextmanager (empty)

39 def _patch_producer(self, broker: KafkaBroker) -> Iterator[None]: (empty)

40 fake_producer = FakeProducer(broker) (empty)

41 

42 with ExitStack() as es: (empty)

43 es.enter_context( (empty)

44 change_producer(broker.config.broker_config, fake_producer), 

45 ) 

46 yield (empty)

47 

48 @staticmethod (empty)

49 async def _fake_connect( # type: ignore[override] (empty)

50 broker: KafkaBroker, (empty)

51 *args: Any, (empty)

52 **kwargs: Any, (empty)

53 ) -> Callable[..., AsyncMock]: (empty)

54 broker.config.broker_config._admin_client = AsyncMock() (empty)

55 

56 builder = MagicMock(return_value=FakeConsumer()) (empty)

57 broker.config.broker_config.builder = builder (empty)

58 

59 return _fake_connection (empty)

60 

61 @staticmethod (empty)

62 def create_publisher_fake_subscriber( (empty)

63 broker: KafkaBroker, 

64 publisher: "LogicPublisher", 

65 ) -> tuple["LogicSubscriber[Any]", bool]: 

66 sub: LogicSubscriber[Any] | None = None (empty)

67 for handler in broker.subscribers: (empty)

68 handler = cast("LogicSubscriber[Any]", handler) (empty)

69 if _is_handler_matches(handler, publisher.topic, publisher.partition): (empty)

70 sub = handler (empty)

71 break (empty)

72 

73 if sub is None: (empty)

74 is_real = False (empty)

75 

76 topic_name = publisher.topic (empty)

77 

78 if publisher.partition: 78 ↛ 79line 78 didn't jump to line 79 because the condition on line 78 was never true(empty)

79 tp = TopicPartition( 

80 topic=topic_name, 

81 partition=publisher.partition, 

82 ) 

83 sub = broker.subscriber( 

84 partitions=[tp], 

85 batch=isinstance(publisher, BatchPublisher), 

86 persistent=False, 

87 ) 

88 else: 

89 sub = broker.subscriber( (empty)

90 topic_name, 

91 batch=isinstance(publisher, BatchPublisher), 

92 persistent=False, 

93 ) 

94 else: 

95 is_real = True (empty)

96 

97 return sub, is_real (empty)

98 

99 

100class FakeConsumer: (empty)

101 async def start(self) -> None: (empty)

102 pass (empty)

103 

104 async def stop(self) -> None: (empty)

105 pass (empty)

106 

107 def subscribe(self, *args: Any, **kwargs: Any) -> None: (empty)

108 pass (empty)

109 

110 

111class FakeProducer(AioKafkaFastProducer): (empty)

112 """A fake Kafka producer for testing purposes. 

113 

114 This class extends AioKafkaFastProducer and is used to simulate Kafka message publishing during tests. 

115 """ 

116 

117 def __init__(self, broker: KafkaBroker) -> None: (empty)

118 self.broker = broker (empty)

119 

120 default = AioKafkaParser( (empty)

121 msg_class=KafkaMessage, 

122 regex=None, 

123 ) 

124 

125 self._parser = ParserComposition(broker._parser, default.parse_message) (empty)

126 self._decoder = ParserComposition(broker._decoder, default.decode_message) (empty)

127 

128 def __bool__(self) -> bool: (empty)

129 return True 

130 

131 @property (empty)

132 def closed(self) -> bool: (empty)

133 return False 

134 

135 @override (empty)

136 async def publish(self, cmd: "KafkaPublishCommand") -> None: (empty)

137 """Publish a message to the Kafka broker.""" 

138 incoming = build_message( (empty)

139 message=cmd.body, 

140 topic=cmd.destination, 

141 key=cmd.key, 

142 partition=cmd.partition, 

143 timestamp_ms=cmd.timestamp_ms, 

144 headers=cmd.headers, 

145 correlation_id=cmd.correlation_id, 

146 reply_to=cmd.reply_to, 

147 serializer=self.broker.config.fd_config._serializer, 

148 ) 

149 

150 for handler in _find_handler( (empty)

151 cast("list[LogicSubscriber[Any]]", self.broker.subscribers), 

152 cmd.destination, 

153 cmd.partition, 

154 ): 

155 msg_to_send = [incoming] if isinstance(handler, BatchSubscriber) else incoming (empty)

156 

157 await self._execute_handler(msg_to_send, cmd.destination, handler) (empty)

158 

159 @override (empty)

160 async def request(self, cmd: "KafkaPublishCommand") -> "ConsumerRecord": (empty)

161 incoming = build_message( (empty)

162 message=cmd.body, 

163 topic=cmd.destination, 

164 key=cmd.key, 

165 partition=cmd.partition, 

166 timestamp_ms=cmd.timestamp_ms, 

167 headers=cmd.headers, 

168 correlation_id=cmd.correlation_id, 

169 serializer=self.broker.config.fd_config._serializer, 

170 ) 

171 

172 for handler in _find_handler( (empty)

173 cast("list[LogicSubscriber[Any]]", self.broker.subscribers), 

174 cmd.destination, 

175 cmd.partition, 

176 ): 

177 msg_to_send = [incoming] if isinstance(handler, BatchSubscriber) else incoming (empty)

178 

179 with anyio.fail_after(cmd.timeout): (empty)

180 return await self._execute_handler( (empty)

181 msg_to_send, 

182 cmd.destination, 

183 handler, 

184 ) 

185 

186 raise SubscriberNotFound (empty)

187 

188 @override (empty)

189 async def publish_batch( (empty)

190 self, 

191 cmd: "KafkaPublishCommand", 

192 ) -> None: 

193 """Publish a batch of messages to the Kafka broker.""" 

194 for handler in _find_handler( (empty)

195 cast("list[LogicSubscriber[Any]]", self.broker.subscribers), 

196 cmd.destination, 

197 cmd.partition, 

198 ): 

199 messages = ( (empty)

200 build_message( 

201 message=message, 

202 topic=cmd.destination, 

203 partition=cmd.partition, 

204 timestamp_ms=cmd.timestamp_ms, 

205 key=cmd.key_for(message_position), 

206 headers=cmd.headers, 

207 correlation_id=cmd.correlation_id, 

208 reply_to=cmd.reply_to, 

209 serializer=self.broker.config.fd_config._serializer, 

210 ) 

211 for message_position, message in enumerate(cmd.batch_bodies) 

212 ) 

213 

214 if isinstance(handler, BatchSubscriber): 214 ↛ 218line 214 didn't jump to line 218 because the condition on line 214 was always true(empty)

215 await self._execute_handler(list(messages), cmd.destination, handler) (empty)

216 

217 else: 

218 for m in messages: 

219 await self._execute_handler(m, cmd.destination, handler) 

220 

221 async def _execute_handler( (empty)

222 self, 

223 msg: Any, 

224 topic: str, 

225 handler: "LogicSubscriber[Any]", 

226 ) -> "ConsumerRecord": 

227 result = await handler.process_message(msg) (empty)

228 

229 return build_message( (empty)

230 topic=topic, 

231 message=result.body, 

232 headers=result.headers, 

233 correlation_id=result.correlation_id, 

234 serializer=self.broker.config.fd_config._serializer, 

235 ) 

236 

237 

238def build_message( (empty)

239 message: "SendableMessage", 

240 topic: str, 

241 partition: int | None = None, 

242 timestamp_ms: int | None = None, 

243 key: bytes | None = None, 

244 headers: dict[str, str] | None = None, 

245 correlation_id: str | None = None, 

246 *, 

247 reply_to: str = "", 

248 serializer: Optional["SerializerProto"], 

249) -> "ConsumerRecord": 

250 """Build a Kafka ConsumerRecord for a sendable message.""" 

251 msg, content_type = encode_message(message, serializer=serializer) (empty)

252 

253 k = key or b"" (empty)

254 

255 headers = { (empty)

256 "content-type": content_type or "", 

257 "correlation_id": correlation_id or gen_cor_id(), 

258 **(headers or {}), 

259 } 

260 

261 if reply_to: (empty)

262 headers["reply_to"] = headers.get("reply_to", reply_to) (empty)

263 

264 return ConsumerRecord( (empty)

265 value=msg, 

266 topic=topic, 

267 partition=partition or 0, 

268 key=k, 

269 serialized_key_size=len(k), 

270 serialized_value_size=len(msg), 

271 checksum=sum(msg), 

272 offset=0, 

273 headers=[(i, j.encode()) for i, j in headers.items()], 

274 timestamp_type=1, 

275 timestamp=timestamp_ms or int(datetime.now(timezone.utc).timestamp() * 1000), 

276 ) 

277 

278 

279def _fake_connection(*args: Any, **kwargs: Any) -> AsyncMock: (empty)

280 mock = AsyncMock() 

281 mock.subscribe = MagicMock 

282 mock.assign = MagicMock 

283 return mock 

284 

285 

286def _find_handler( (empty)

287 subscribers: Iterable["LogicSubscriber[Any]"], 

288 topic: str, 

289 partition: int | None, 

290) -> Generator["LogicSubscriber[Any]", None, None]: 

291 published_groups = set() (empty)

292 for handler in subscribers: # pragma: no branch (empty)

293 if _is_handler_matches(handler, topic, partition): (empty)

294 if handler.group_id: (empty)

295 if handler.group_id in published_groups: (empty)

296 continue (empty)

297 else: 

298 published_groups.add(handler.group_id) (empty)

299 yield handler (empty)

300 

301 

302def _is_handler_matches( (empty)

303 handler: "LogicSubscriber[Any]", 

304 topic: str, 

305 partition: int | None, 

306) -> bool: 

307 return bool( (empty)

308 any( 

309 p.topic == topic and (partition is None or p.partition == partition) 

310 for p in handler.partitions 

311 ) 

312 or topic in handler.topics 

313 or (handler.pattern and re.match(handler.pattern, topic)), 

314 )