Coverage for faststream / redis / testing.py: 98%

1import re (empty)

2from collections.abc import Iterator, Sequence (empty)

3from contextlib import ExitStack, contextmanager (empty)

4from typing import ( (empty)

5 TYPE_CHECKING, 

6 Any, 

7 Optional, 

8 Protocol, 

9 Union, 

10 cast, 

11) 

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

13 

14import anyio (empty)

15from typing_extensions import TypedDict, override (empty)

16 

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

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

19from faststream.exceptions import SetupError, SubscriberNotFound (empty)

20from faststream.message import gen_cor_id (empty)

21from faststream.redis.broker.broker import RedisBroker (empty)

22from faststream.redis.message import ( (empty)

23 BatchListMessage, 

24 BatchStreamMessage, 

25 DefaultListMessage, 

26 DefaultStreamMessage, 

27 PubSubMessage, 

28 bDATA_KEY, 

29) 

30from faststream.redis.parser import MessageFormat, ParserConfig, RedisPubSubParser (empty)

31from faststream.redis.publisher.producer import RedisFastProducer (empty)

32from faststream.redis.response import DestinationType, RedisPublishCommand (empty)

33from faststream.redis.schemas import INCORRECT_SETUP_MSG (empty)

34from faststream.redis.subscriber.usecases.channel_subscriber import ChannelSubscriber (empty)

35from faststream.redis.subscriber.usecases.list_subscriber import _ListHandlerMixin (empty)

36from faststream.redis.subscriber.usecases.stream_subscriber import _StreamHandlerMixin (empty)

37 

38if TYPE_CHECKING: (empty)

39 from fast_depends.library.serializer import SerializerProto 

40 

41 from faststream._internal.basic_types import SendableMessage 

42 from faststream.redis.publisher.usecase import LogicPublisher 

43 from faststream.redis.subscriber.usecases.basic import LogicSubscriber 

44 

45__all__ = ("TestRedisBroker",) (empty)

46 

47 

48class TestRedisBroker(TestBroker[RedisBroker]): (empty)

49 """A class to test Redis brokers.""" 

50 

51 @contextmanager (empty)

52 def _patch_producer(self, broker: RedisBroker) -> Iterator[None]: (empty)

53 with ExitStack() as es: (empty)

54 es.enter_context( (empty)

55 change_producer( 

56 broker.config.broker_config, FakeProducer(broker, broker.config) 

57 ), 

58 ) 

59 

60 for publisher in cast("list[LogicPublisher]", broker.publishers): (empty)

61 es.enter_context( (empty)

62 change_producer(publisher, FakeProducer(broker, publisher.config)), 

63 ) 

64 

65 yield (empty)

66 

67 @staticmethod (empty)

68 def create_publisher_fake_subscriber( (empty)

69 broker: RedisBroker, 

70 publisher: "LogicPublisher", 

71 ) -> tuple["LogicSubscriber", bool]: 

72 sub: LogicSubscriber | None = None (empty)

73 

74 named_property = publisher.subscriber_property(name_only=True) (empty)

75 visitors = (ChannelVisitor(), ListVisitor(), StreamVisitor()) (empty)

76 

77 for handler in broker.subscribers: # pragma: no branch (empty)

78 handler = cast("LogicSubscriber", handler) (empty)

79 for visitor in visitors: (empty)

80 if visitor.visit(**named_property, sub=handler): (empty)

81 sub = handler (empty)

82 break (empty)

83 

84 if sub is None: (empty)

85 is_real = False (empty)

86 sub_options = publisher.subscriber_property(name_only=False) (empty)

87 sub = broker.subscriber(**sub_options, persistent=False) (empty)

88 else: 

89 is_real = True (empty)

90 

91 return sub, is_real (empty)

92 

93 @staticmethod (empty)

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

95 broker: RedisBroker, 

96 *args: Any, 

97 **kwargs: Any, 

98 ) -> AsyncMock: 

99 connection = MagicMock() (empty)

100 

101 pub_sub = AsyncMock() (empty)

102 

103 async def get_msg(*args: Any, timeout: float, **kwargs: Any) -> None: (empty)

104 await anyio.sleep(timeout) (empty)

105 

106 pub_sub.get_message = get_msg (empty)

107 

108 connection.pubsub.side_effect = lambda: pub_sub (empty)

109 connection.aclose = AsyncMock() (empty)

110 

111 connection.xack = AsyncMock() (empty)

112 connection.xdel = AsyncMock() (empty)

113 

114 broker.config.broker_config.connection._client = connection (empty)

115 return connection (empty)

116 

117 

118class FakeProducer(RedisFastProducer): (empty)

119 def __init__(self, broker: RedisBroker, config: ParserConfig) -> None: (empty)

120 self.broker = broker (empty)

121 

122 default = RedisPubSubParser(config) (empty)

123 

124 self._parser = ParserComposition( (empty)

125 broker._parser, 

126 default.parse_message, 

127 ) 

128 self._decoder = ParserComposition( (empty)

129 broker._decoder, 

130 default.decode_message, 

131 ) 

132 

133 @override (empty)

134 async def publish(self, cmd: "RedisPublishCommand") -> int | bytes: (empty)

135 body = build_message( (empty)

136 message=cmd.body, 

137 reply_to=cmd.reply_to, 

138 correlation_id=cmd.correlation_id or gen_cor_id(), 

139 headers=cmd.headers, 

140 message_format=cmd.message_format, 

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

142 ) 

143 

144 destination = _make_destination_kwargs(cmd) (empty)

145 visitors = (ChannelVisitor(), ListVisitor(), StreamVisitor()) (empty)

146 

147 for handler in self.broker.subscribers: # pragma: no branch (empty)

148 handler = cast("LogicSubscriber", handler) (empty)

149 for visitor in visitors: (empty)

150 if visited_ch := visitor.visit(**destination, sub=handler): (empty)

151 msg = visitor.get_message( (empty)

152 visited_ch, 

153 body, 

154 handler, # type: ignore[arg-type] 

155 ) 

156 

157 await self._execute_handler(msg, handler) (empty)

158 

159 return 0 (empty)

160 

161 @override (empty)

162 async def request(self, cmd: "RedisPublishCommand") -> "PubSubMessage": (empty)

163 body = build_message( (empty)

164 message=cmd.body, 

165 correlation_id=cmd.correlation_id or gen_cor_id(), 

166 headers=cmd.headers, 

167 message_format=cmd.message_format, 

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

169 ) 

170 

171 destination = _make_destination_kwargs(cmd) (empty)

172 visitors = (ChannelVisitor(), ListVisitor(), StreamVisitor()) (empty)

173 

174 for handler in self.broker.subscribers: # pragma: no branch (empty)

175 handler = cast("LogicSubscriber", handler) (empty)

176 for visitor in visitors: (empty)

177 if visited_ch := visitor.visit(**destination, sub=handler): (empty)

178 msg = visitor.get_message( (empty)

179 visited_ch, 

180 body, 

181 handler, # type: ignore[arg-type] 

182 ) 

183 

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

185 return await self._execute_handler(msg, handler) (empty)

186 

187 raise SubscriberNotFound (empty)

188 

189 @override (empty)

190 async def publish_batch(self, cmd: "RedisPublishCommand") -> int: (empty)

191 data_to_send = [ (empty)

192 build_message( 

193 m, 

194 correlation_id=cmd.correlation_id or gen_cor_id(), 

195 headers=cmd.headers, 

196 message_format=cmd.message_format, 

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

198 ) 

199 for m in cmd.batch_bodies 

200 ] 

201 

202 visitor = ListVisitor() (empty)

203 for handler in self.broker.subscribers: # pragma: no branch (empty)

204 handler = cast("LogicSubscriber", handler) (empty)

205 if visitor.visit(list=cmd.destination, sub=handler): (empty)

206 casted_handler = cast("_ListHandlerMixin", handler) (empty)

207 

208 if casted_handler.list_sub.batch: 208 ↛ 203line 208 didn't jump to line 203 because the condition on line 208 was always true(empty)

209 msg = visitor.get_message( (empty)

210 channel=cmd.destination, 

211 body=data_to_send, 

212 sub=casted_handler, 

213 ) 

214 

215 await self._execute_handler(msg, handler) (empty)

216 

217 return 0 (empty)

218 

219 async def _execute_handler( (empty)

220 self, 

221 msg: Any, 

222 handler: "LogicSubscriber", 

223 ) -> "PubSubMessage": 

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

225 

226 return PubSubMessage( (empty)

227 type="message", 

228 data=build_message( 

229 message=result.body, 

230 headers=result.headers, 

231 correlation_id=result.correlation_id or "", 

232 message_format=handler.config.message_format, 

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

234 ), 

235 channel="", 

236 pattern=None, 

237 ) 

238 

239 

240def build_message( (empty)

241 message: Union[Sequence["SendableMessage"], "SendableMessage"], 

242 *, 

243 correlation_id: str, 

244 message_format: type["MessageFormat"], 

245 reply_to: str = "", 

246 headers: dict[str, Any] | None = None, 

247 serializer: Optional["SerializerProto"] = None, 

248) -> bytes: 

249 return message_format.encode( (empty)

250 message=message, 

251 reply_to=reply_to, 

252 headers=headers, 

253 correlation_id=correlation_id, 

254 serializer=serializer, 

255 ) 

256 

257 

258class Visitor(Protocol): (empty)

259 def visit( (empty)

260 self, 

261 *, 

262 channel: str | None, (empty)

263 list: str | None, (empty)

264 stream: str | None, (empty)

265 sub: "LogicSubscriber", (empty)

266 ) -> str | None: ... (empty)

267 

268 def get_message(self, channel: str, body: Any, sub: "LogicSubscriber") -> Any: ... (empty)

269 

270 

271class ChannelVisitor(Visitor): (empty)

272 def visit( (empty)

273 self, 

274 *, 

275 sub: "LogicSubscriber", 

276 channel: str | None = None, 

277 list: str | None = None, 

278 stream: str | None = None, 

279 ) -> str | None: 

280 if channel is None or not isinstance(sub, ChannelSubscriber): (empty)

281 return None (empty)

282 

283 sub_channel = sub.channel (empty)

284 

285 if ( (empty)

286 sub_channel.pattern 

287 and bool( 

288 re.match( 

289 sub_channel.name.replace(".", "\\.").replace("*", ".*"), 

290 channel or "", 

291 ), 

292 ) 

293 ) or channel == sub_channel.name: 

294 return channel (empty)

295 

296 return None (empty)

297 

298 def get_message( # type: ignore[override] (empty)

299 self, 

300 channel: str, 

301 body: Any, 

302 sub: "ChannelSubscriber", 

303 ) -> Any: 

304 return PubSubMessage( (empty)

305 type="message", 

306 data=body, 

307 channel=channel, 

308 pattern=sub.channel.pattern.encode() if sub.channel.pattern else None, 

309 ) 

310 

311 

312class ListVisitor(Visitor): (empty)

313 def visit( (empty)

314 self, 

315 *, 

316 sub: "LogicSubscriber", 

317 channel: str | None = None, 

318 list: str | None = None, 

319 stream: str | None = None, 

320 ) -> str | None: 

321 if list is None or not isinstance(sub, _ListHandlerMixin): (empty)

322 return None (empty)

323 

324 if list == sub.list_sub.name: (empty)

325 return list (empty)

326 

327 return None (empty)

328 

329 def get_message( # type: ignore[override] (empty)

330 self, 

331 channel: str, 

332 body: Any, 

333 sub: "_ListHandlerMixin", 

334 ) -> Any: 

335 if sub.list_sub.batch: (empty)

336 return BatchListMessage( (empty)

337 type="blist", 

338 channel=channel, 

339 data=body if isinstance(body, list) else [body], 

340 ) 

341 

342 return DefaultListMessage( (empty)

343 type="list", 

344 channel=channel, 

345 data=body, 

346 ) 

347 

348 

349class StreamVisitor(Visitor): (empty)

350 def visit( (empty)

351 self, 

352 *, 

353 sub: "LogicSubscriber", 

354 channel: str | None = None, 

355 list: str | None = None, 

356 stream: str | None = None, 

357 ) -> str | None: 

358 if stream is None or not isinstance(sub, _StreamHandlerMixin): (empty)

359 return None (empty)

360 

361 if stream == sub.stream_sub.name: (empty)

362 return stream (empty)

363 

364 return None (empty)

365 

366 def get_message( # type: ignore[override] (empty)

367 self, 

368 channel: str, 

369 body: Any, 

370 sub: "_StreamHandlerMixin", 

371 ) -> Any: 

372 if sub.stream_sub.batch: (empty)

373 return BatchStreamMessage( (empty)

374 type="bstream", 

375 channel=channel, 

376 data=[{bDATA_KEY: body}], 

377 message_ids=[], 

378 ) 

379 

380 return DefaultStreamMessage( (empty)

381 type="stream", 

382 channel=channel, 

383 data={bDATA_KEY: body}, 

384 message_ids=[], 

385 ) 

386 

387 

388class _DestinationKwargs(TypedDict, total=False): (empty)

389 channel: str (empty)

390 list: str (empty)

391 stream: str (empty)

392 

393 

394def _make_destination_kwargs(cmd: RedisPublishCommand) -> _DestinationKwargs: (empty)

395 destination: _DestinationKwargs = {} (empty)

396 if cmd.destination_type is DestinationType.Channel: (empty)

397 destination["channel"] = cmd.destination (empty)

398 if cmd.destination_type is DestinationType.List: (empty)

399 destination["list"] = cmd.destination (empty)

400 if cmd.destination_type is DestinationType.Stream: (empty)

401 destination["stream"] = cmd.destination (empty)

402 

403 if len(destination) != 1: 403 ↛ 404line 403 didn't jump to line 404 because the condition on line 403 was never true(empty)

404 raise SetupError(INCORRECT_SETUP_MSG) 

405 

406 return destination (empty)