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

1import asyncio (empty)

2from unittest.mock import MagicMock, call (empty)

3 

4import anyio (empty)

5import pytest (empty)

6from pydantic import BaseModel (empty)

7 

8from faststream import Context, Depends (empty)

9from faststream.exceptions import StopConsume (empty)

10 

11from .basic import BaseTestcaseConfig (empty)

12 

13 

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

15class BrokerConsumeTestcase(BaseTestcaseConfig): (empty)

16 async def test_consume( (empty)

17 self, 

18 queue: str, 

19 ) -> None: 

20 event = asyncio.Event() (empty)

21 consume_broker = self.get_broker() (empty)

22 

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

24 

25 @consume_broker.subscriber(*args, **kwargs) (empty)

26 def subscriber(m) -> None: (empty)

27 event.set() (empty)

28 

29 async with self.patch_broker(consume_broker) as br: (empty)

30 await br.start() (empty)

31 await asyncio.wait( (empty)

32 ( 

33 asyncio.create_task(br.publish("hello", queue)), 

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

35 ), 

36 timeout=self.timeout, 

37 ) 

38 

39 assert event.is_set() (empty)

40 

41 async def test_consume_from_multi( (empty)

42 self, 

43 queue: str, 

44 mock: MagicMock, 

45 ) -> None: 

46 consume_broker = self.get_broker() (empty)

47 

48 consume = asyncio.Event() (empty)

49 consume2 = asyncio.Event() (empty)

50 

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

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

53 

54 @consume_broker.subscriber(*args, **kwargs) (empty)

55 @consume_broker.subscriber(*args2, **kwargs2) (empty)

56 def subscriber(m) -> None: (empty)

57 mock() (empty)

58 if not consume.is_set(): (empty)

59 consume.set() (empty)

60 else: 

61 consume2.set() (empty)

62 

63 async with self.patch_broker(consume_broker) as br: (empty)

64 await br.start() (empty)

65 await asyncio.wait( (empty)

66 ( 

67 asyncio.create_task(br.publish("hello", queue)), 

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

69 asyncio.create_task(consume.wait()), 

70 asyncio.create_task(consume2.wait()), 

71 ), 

72 timeout=self.timeout, 

73 ) 

74 

75 assert consume.is_set() (empty)

76 assert consume2.is_set() (empty)

77 assert mock.call_count == 2 (empty)

78 

79 async def test_consume_double( (empty)

80 self, 

81 queue: str, 

82 mock: MagicMock, 

83 ) -> None: 

84 consume_broker = self.get_broker() (empty)

85 

86 consume = asyncio.Event() (empty)

87 consume2 = asyncio.Event() (empty)

88 

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

90 

91 @consume_broker.subscriber(*args, **kwargs) (empty)

92 async def handler(m) -> None: (empty)

93 mock() (empty)

94 if not consume.is_set(): (empty)

95 consume.set() (empty)

96 else: 

97 consume2.set() (empty)

98 

99 async with self.patch_broker(consume_broker) as br: (empty)

100 await br.start() (empty)

101 await asyncio.wait( (empty)

102 ( 

103 asyncio.create_task(br.publish("hello", queue)), 

104 asyncio.create_task(br.publish("hello", queue)), 

105 asyncio.create_task(consume.wait()), 

106 asyncio.create_task(consume2.wait()), 

107 ), 

108 timeout=self.timeout, 

109 ) 

110 

111 assert consume2.is_set() (empty)

112 assert consume.is_set() (empty)

113 assert mock.call_count == 2 (empty)

114 

115 async def test_different_consume( (empty)

116 self, 

117 queue: str, 

118 mock: MagicMock, 

119 ) -> None: 

120 consume_broker = self.get_broker() (empty)

121 

122 consume = asyncio.Event() (empty)

123 consume2 = asyncio.Event() (empty)

124 

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

126 

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

128 def handler(m) -> None: (empty)

129 mock.handler() (empty)

130 consume.set() (empty)

131 

132 another_topic = queue + "1" (empty)

133 args, kwargs = self.get_subscriber_params(another_topic) (empty)

134 

135 @consume_broker.subscriber(*args, **kwargs) (empty)

136 def handler2(m) -> None: (empty)

137 mock.handler2() (empty)

138 consume2.set() (empty)

139 

140 async with self.patch_broker(consume_broker) as br: (empty)

141 await br.start() (empty)

142 await asyncio.wait( (empty)

143 ( 

144 asyncio.create_task(br.publish("hello", queue)), 

145 asyncio.create_task(br.publish("hello", another_topic)), 

146 asyncio.create_task(consume.wait()), 

147 asyncio.create_task(consume2.wait()), 

148 ), 

149 timeout=self.timeout, 

150 ) 

151 

152 assert consume.is_set() (empty)

153 assert consume2.is_set() (empty)

154 mock.handler.assert_called_once() (empty)

155 mock.handler2.assert_called_once() (empty)

156 

157 async def test_consume_with_filter( (empty)

158 self, 

159 queue: str, 

160 mock: MagicMock, 

161 ) -> None: 

162 consume_broker = self.get_broker() (empty)

163 

164 consume = asyncio.Event() (empty)

165 consume2 = asyncio.Event() (empty)

166 

167 args, kwargs = self.get_subscriber_params( (empty)

168 queue, 

169 ) 

170 

171 sub = consume_broker.subscriber(*args, **kwargs) (empty)

172 

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

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

175 mock.handler(m) (empty)

176 consume.set() (empty)

177 

178 @sub (empty)

179 async def handler2(m) -> None: (empty)

180 mock.handler2(m) (empty)

181 consume2.set() (empty)

182 

183 async with self.patch_broker(consume_broker) as br: (empty)

184 await br.start() (empty)

185 await asyncio.wait( (empty)

186 ( 

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

188 asyncio.create_task(br.publish("hello", queue)), 

189 asyncio.create_task(consume.wait()), 

190 asyncio.create_task(consume2.wait()), 

191 ), 

192 timeout=self.timeout, 

193 ) 

194 

195 assert consume.is_set() (empty)

196 assert consume2.is_set() (empty)

197 mock.handler.assert_called_once_with({"msg": "hello"}) (empty)

198 mock.handler2.assert_called_once_with("hello") (empty)

199 

200 async def test_consume_validate_false( (empty)

201 self, 

202 queue: str, 

203 mock: MagicMock, 

204 ) -> None: 

205 event = asyncio.Event() (empty)

206 

207 consume_broker = self.get_broker( (empty)

208 apply_types=True, 

209 serializer=None, 

210 ) 

211 

212 class Foo(BaseModel): (empty)

213 x: int (empty)

214 

215 def dependency() -> str: (empty)

216 return "100" (empty)

217 

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

219 

220 @consume_broker.subscriber(*args, **kwargs) (empty)

221 async def handler( (empty)

222 m: Foo, 

223 dep: int = Depends(dependency), 

224 broker=Context(), 

225 ) -> None: 

226 mock(m, dep, broker) (empty)

227 event.set() (empty)

228 

229 async with self.patch_broker(consume_broker) as br: (empty)

230 await br.start() (empty)

231 

232 await asyncio.wait( (empty)

233 ( 

234 asyncio.create_task(br.publish({"x": 1}, queue)), 

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

236 ), 

237 timeout=self.timeout, 

238 ) 

239 

240 assert event.is_set() (empty)

241 mock.assert_called_once_with({"x": 1}, "100", consume_broker) (empty)

242 

243 async def test_dynamic_sub(self, queue: str) -> None: (empty)

244 event = asyncio.Event() (empty)

245 

246 consume_broker = self.get_broker() (empty)

247 

248 async def subscriber(m) -> None: (empty)

249 event.set() (empty)

250 

251 async with self.patch_broker(consume_broker) as br: (empty)

252 await br.start() (empty)

253 

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

255 sub = br.subscriber(*args, **kwargs) (empty)

256 sub(subscriber) (empty)

257 await sub.start() (empty)

258 

259 await br.publish("hello", queue) (empty)

260 

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

262 await event.wait() (empty)

263 

264 await sub.stop() (empty)

265 

266 assert event.is_set() (empty)

267 

268 async def test_get_one_conflicts_with_handler(self, queue) -> None: (empty)

269 broker = self.get_broker(apply_types=True) (empty)

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

271 subscriber = broker.subscriber(*args, **kwargs) (empty)

272 

273 @subscriber (empty)

274 async def t() -> None: ... (empty)

275 

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

277 await br.start() (empty)

278 

279 with pytest.raises(AssertionError): (empty)

280 await subscriber.get_one(timeout=1e-24) (empty)

281 

282 

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

284class BrokerRealConsumeTestcase(BrokerConsumeTestcase): (empty)

285 async def test_get_one( (empty)

286 self, 

287 queue: str, 

288 mock: MagicMock, 

289 ) -> None: 

290 broker = self.get_broker(apply_types=True) (empty)

291 

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

293 subscriber = broker.subscriber(*args, **kwargs) (empty)

294 

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

296 await br.start() (empty)

297 

298 async def consume() -> None: (empty)

299 mock(await subscriber.get_one(timeout=self.timeout)) (empty)

300 

301 async def publish() -> None: (empty)

302 await anyio.sleep(1e-24) (empty)

303 await br.publish("test_message", queue) (empty)

304 

305 await asyncio.wait( (empty)

306 ( 

307 asyncio.create_task(consume()), 

308 asyncio.create_task(publish()), 

309 ), 

310 timeout=self.timeout, 

311 ) 

312 

313 mock.assert_called_once() (empty)

314 message = mock.call_args[0][0] (empty)

315 assert message (empty)

316 assert await message.decode() == "test_message", await message.decode() (empty)

317 

318 async def test_get_one_timeout( (empty)

319 self, 

320 queue: str, 

321 mock: MagicMock, 

322 ) -> None: 

323 broker = self.get_broker(apply_types=True) (empty)

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

325 subscriber = broker.subscriber(*args, **kwargs) (empty)

326 

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

328 await br.start() (empty)

329 

330 mock(await subscriber.get_one(timeout=1e-24)) (empty)

331 mock.assert_called_once_with(None) (empty)

332 

333 @pytest.mark.slow() (empty)

334 async def test_stop_consume_exc( (empty)

335 self, 

336 queue: str, 

337 mock: MagicMock, 

338 ) -> None: 

339 event = asyncio.Event() (empty)

340 

341 consume_broker = self.get_broker() (empty)

342 

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

344 

345 @consume_broker.subscriber(*args, **kwargs) (empty)

346 def subscriber(m): (empty)

347 mock() (empty)

348 event.set() (empty)

349 raise StopConsume (empty)

350 

351 async with self.patch_broker(consume_broker) as br: (empty)

352 await br.start() (empty)

353 await asyncio.wait( (empty)

354 ( 

355 asyncio.create_task(br.publish("hello", queue)), 

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

357 ), 

358 timeout=self.timeout, 

359 ) 

360 await asyncio.sleep(0.5) (empty)

361 await br.publish("hello", queue) (empty)

362 await asyncio.sleep(0.5) (empty)

363 

364 assert event.is_set() (empty)

365 mock.assert_called_once() (empty)

366 

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

368 async def test_iteration( (empty)

369 self, 

370 queue: str, 

371 mock: MagicMock, 

372 ) -> None: 

373 expected_messages = ("test_message_1", "test_message_2") (empty)

374 

375 broker = self.get_broker(apply_types=True) (empty)

376 

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

378 subscriber = broker.subscriber(*args, **kwargs) (empty)

379 

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

381 await br.start() (empty)

382 

383 async def publish_test_message(): (empty)

384 for msg in expected_messages: (empty)

385 await br.publish(msg, queue) (empty)

386 

387 async def consume(): (empty)

388 index_message = 0 (empty)

389 async for msg in subscriber: 389 ↛ exitline 389 didn't return from function 'consume' because the loop on line 389 didn't complete(empty)

390 result_message = await msg.decode() (empty)

391 

392 mock(result_message) (empty)

393 

394 index_message += 1 (empty)

395 if index_message >= len(expected_messages): (empty)

396 break (empty)

397 

398 await asyncio.wait( (empty)

399 ( 

400 asyncio.create_task(consume()), 

401 asyncio.create_task(publish_test_message()), 

402 ), 

403 timeout=self.timeout, 

404 ) 

405 

406 calls = [call(msg) for msg in expected_messages] (empty)

407 mock.assert_has_calls(calls=calls) (empty)