Coverage for tests / brokers / base / parser.py: 99%

1import asyncio (empty)

2from unittest.mock import MagicMock (empty)

3 

4import anyio (empty)

5import pytest (empty)

6 

7from faststream.message.message import StreamMessage (empty)

8 

9from .basic import BaseTestcaseConfig (empty)

10 

11 

12@pytest.mark.asyncio() (empty)

13class LocalCustomParserTestcase(BaseTestcaseConfig): (empty)

14 async def test_local_parser( (empty)

15 self, 

16 mock: MagicMock, 

17 event: asyncio.Event, 

18 queue: str, 

19 ) -> None: 

20 broker = self.get_broker() (empty)

21 

22 async def custom_parser(msg, original): (empty)

23 msg = await original(msg) (empty)

24 mock(msg.body) (empty)

25 return msg (empty)

26 

27 args, kwargs = self.get_subscriber_params(queue, parser=custom_parser) (empty)

28 

29 @broker.subscriber(*args, **kwargs) (empty)

30 async def handle(m) -> None: (empty)

31 event.set() (empty)

32 

33 async with self.patch_broker(broker) as br: (empty)

34 await br.start() (empty)

35 

36 await asyncio.wait( (empty)

37 ( 

38 asyncio.create_task(br.publish(b"hello", queue)), 

39 asyncio.create_task(event.wait()), 

40 ), 

41 timeout=self.timeout, 

42 ) 

43 

44 assert event.is_set() (empty)

45 mock.assert_called_once_with(b"hello") (empty)

46 

47 async def test_local_sync_decoder( (empty)

48 self, 

49 mock: MagicMock, 

50 event: asyncio.Event, 

51 queue: str, 

52 ) -> None: 

53 broker = self.get_broker() (empty)

54 

55 def custom_decoder(msg): (empty)

56 mock(msg.body) (empty)

57 return msg (empty)

58 

59 args, kwargs = self.get_subscriber_params(queue, decoder=custom_decoder) (empty)

60 

61 @broker.subscriber(*args, **kwargs) (empty)

62 async def handle(m) -> None: (empty)

63 event.set() (empty)

64 

65 async with self.patch_broker(broker) as br: (empty)

66 await br.start() (empty)

67 

68 await asyncio.wait( (empty)

69 ( 

70 asyncio.create_task(br.publish(b"hello", queue)), 

71 asyncio.create_task(event.wait()), 

72 ), 

73 timeout=self.timeout, 

74 ) 

75 

76 assert event.is_set() (empty)

77 mock.assert_called_once_with(b"hello") (empty)

78 

79 async def test_global_sync_decoder( (empty)

80 self, 

81 mock: MagicMock, 

82 event: asyncio.Event, 

83 queue: str, 

84 ) -> None: 

85 def custom_decoder(msg): (empty)

86 mock(msg.body) (empty)

87 return msg (empty)

88 

89 broker = self.get_broker(decoder=custom_decoder) (empty)

90 

91 args, kwargs = self.get_subscriber_params(queue) (empty)

92 

93 @broker.subscriber(*args, **kwargs) (empty)

94 async def handle(m) -> None: (empty)

95 event.set() (empty)

96 

97 async with self.patch_broker(broker) as br: (empty)

98 await br.start() (empty)

99 

100 await asyncio.wait( (empty)

101 ( 

102 asyncio.create_task(br.publish(b"hello", queue)), 

103 asyncio.create_task(event.wait()), 

104 ), 

105 timeout=self.timeout, 

106 ) 

107 

108 assert event.is_set() (empty)

109 mock.assert_called_once_with(b"hello") (empty)

110 

111 async def test_local_parser_no_share_between_subscribers( (empty)

112 self, 

113 mock: MagicMock, 

114 queue: str, 

115 ) -> None: 

116 event, event2 = asyncio.Event(), asyncio.Event() (empty)

117 broker = self.get_broker() (empty)

118 

119 async def custom_parser(msg, original): (empty)

120 msg = await original(msg) (empty)

121 mock(msg.body) (empty)

122 return msg (empty)

123 

124 args, kwargs = self.get_subscriber_params(queue, parser=custom_parser) (empty)

125 args2, kwargs2 = self.get_subscriber_params(queue + "1") (empty)

126 

127 @broker.subscriber(*args, **kwargs) (empty)

128 @broker.subscriber(*args2, **kwargs2) (empty)

129 async def handle(m) -> None: (empty)

130 if event.is_set(): (empty)

131 event2.set() (empty)

132 else: 

133 event.set() (empty)

134 

135 async with self.patch_broker(broker) as br: (empty)

136 await br.start() (empty)

137 

138 await asyncio.wait( (empty)

139 ( 

140 asyncio.create_task(br.publish(b"hello", queue)), 

141 asyncio.create_task(br.publish(b"hello", queue + "1")), 

142 asyncio.create_task(event.wait()), 

143 asyncio.create_task(event2.wait()), 

144 ), 

145 timeout=self.timeout, 

146 ) 

147 

148 assert event.is_set() (empty)

149 assert event2.is_set() (empty)

150 mock.assert_called_once_with(b"hello") (empty)

151 

152 async def test_local_parser_no_share_between_handlers( (empty)

153 self, 

154 mock: MagicMock, 

155 queue: str, 

156 ) -> None: 

157 event, event2 = asyncio.Event(), asyncio.Event() (empty)

158 

159 broker = self.get_broker() (empty)

160 

161 args, kwargs = self.get_subscriber_params(queue) (empty)

162 sub = broker.subscriber(*args, **kwargs) (empty)

163 

164 @sub(filter=lambda m: m.content_type == "application/json") (empty)

165 async def handle(m) -> None: (empty)

166 event.set() (empty)

167 

168 async def custom_parser(msg, original): (empty)

169 msg = await original(msg) (empty)

170 mock(msg.body) (empty)

171 return msg (empty)

172 

173 @sub(parser=custom_parser) (empty)

174 async def handle2(m) -> None: (empty)

175 event2.set() (empty)

176 

177 async with self.patch_broker(broker) as br: (empty)

178 await br.start() (empty)

179 

180 await asyncio.wait( (empty)

181 ( 

182 asyncio.create_task(br.publish({"msg": "hello"}, queue)), 

183 asyncio.create_task(br.publish(b"hello", queue)), 

184 asyncio.create_task(event.wait()), 

185 asyncio.create_task(event2.wait()), 

186 ), 

187 timeout=self.timeout, 

188 ) 

189 

190 assert event.is_set() (empty)

191 assert event2.is_set() (empty)

192 assert mock.call_count == 1 (empty)

193 

194 @pytest.mark.connected() (empty)

195 async def test_iterator_respect_decoder( (empty)

196 self, 

197 mock: MagicMock, 

198 queue: str, 

199 ) -> None: 

200 """Fixes https://github.com/ag2ai/faststream/issues/2554.""" 

201 start_event, consumed_event, stopped_event = ( (empty)

202 asyncio.Event(), 

203 asyncio.Event(), 

204 asyncio.Event(), 

205 ) 

206 

207 broker = self.get_broker() (empty)

208 

209 async def custom_decoder(msg, original): (empty)

210 mock() (empty)

211 consumed_event.set() (empty)

212 return await original(msg) (empty)

213 

214 args, kwargs = self.get_subscriber_params(queue, decoder=custom_decoder) (empty)

215 sub = broker.subscriber(*args, **kwargs) (empty)

216 

217 async def iter_messages() -> None: (empty)

218 await sub.start() (empty)

219 start_event.set() (empty)

220 

221 async for m in sub: 221 ↛ 227line 221 didn't jump to line 227 because the loop on line 221 didn't complete(empty)

222 assert not mock.called (empty)

223 msg = await m.decode() (empty)

224 mock.assert_called_once() (empty)

225 break (empty)

226 

227 await sub.stop() (empty)

228 stopped_event.set() (empty)

229 return msg (empty)

230 

231 async with broker: (empty)

232 t = asyncio.create_task(iter_messages()) (empty)

233 

234 with anyio.move_on_after(self.timeout): (empty)

235 await start_event.wait() (empty)

236 await broker.publish(b"hello", queue) (empty)

237 await consumed_event.wait() (empty)

238 await stopped_event.wait() (empty)

239 

240 await t (empty)

241 data = t.result() (empty)

242 assert data == b"hello", data (empty)

243 

244 @pytest.mark.connected() (empty)

245 async def test_get_one_respect_decoder( (empty)

246 self, 

247 queue: str, 

248 event: asyncio.Event, 

249 mock: MagicMock, 

250 ) -> None: 

251 """Fixes https://github.com/ag2ai/faststream/issues/2554.""" 

252 broker = self.get_broker() (empty)

253 

254 async def custom_decoder(msg, original): (empty)

255 mock() (empty)

256 return await original(msg) (empty)

257 

258 args, kwargs = self.get_subscriber_params(queue, decoder=custom_decoder) (empty)

259 sub = broker.subscriber(*args, **kwargs) (empty)

260 

261 async def get_msg() -> StreamMessage: (empty)

262 await sub.start() (empty)

263 event.set() (empty)

264 msg = await sub.get_one(timeout=self.timeout) (empty)

265 await sub.stop() (empty)

266 return msg (empty)

267 

268 async with broker: (empty)

269 task = asyncio.create_task(get_msg()) (empty)

270 

271 with anyio.move_on_after(self.timeout): (empty)

272 await event.wait() (empty)

273 await broker.publish(b"hello", queue) (empty)

274 

275 await task (empty)

276 msg = task.result() (empty)

277 

278 assert not mock.called, mock.call_count (empty)

279 await msg.decode() (empty)

280 mock.assert_called_once() (empty)

281 

282 

283class CustomParserTestcase(LocalCustomParserTestcase): (empty)

284 async def test_global_parser( (empty)

285 self, 

286 mock: MagicMock, 

287 event: asyncio.Event, 

288 queue: str, 

289 ) -> None: 

290 async def custom_parser(msg, original): (empty)

291 msg = await original(msg) (empty)

292 mock(msg.body) (empty)

293 return msg (empty)

294 

295 broker = self.get_broker(parser=custom_parser) (empty)

296 

297 args, kwargs = self.get_subscriber_params(queue) (empty)

298 

299 @broker.subscriber(*args, **kwargs) (empty)

300 async def handle(m) -> None: (empty)

301 event.set() (empty)

302 

303 async with self.patch_broker(broker) as br: (empty)

304 await br.start() (empty)

305 

306 await asyncio.wait( (empty)

307 ( 

308 asyncio.create_task(br.publish(b"hello", queue)), 

309 asyncio.create_task(event.wait()), 

310 ), 

311 timeout=self.timeout, 

312 ) 

313 

314 assert event.is_set() (empty)

315 mock.assert_called_once_with(b"hello") (empty)