Coverage for tests / brokers / kafka / test_test_client.py: 99%

1import asyncio (empty)

2from unittest.mock import AsyncMock, patch (empty)

3 

4import pytest (empty)

5 

6from faststream import AckPolicy, BaseMiddleware (empty)

7from faststream.kafka import TopicPartition (empty)

8from faststream.kafka.annotations import KafkaMessage (empty)

9from faststream.kafka.message import FAKE_CONSUMER (empty)

10from faststream.kafka.testing import FakeProducer (empty)

11from tests.brokers.base.testclient import BrokerTestclientTestcase (empty)

12from tests.tools import spy_decorator (empty)

13 

14from .basic import KafkaMemoryTestcaseConfig (empty)

15 

16 

17@pytest.mark.kafka() (empty)

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

19class TestTestclient(KafkaMemoryTestcaseConfig, BrokerTestclientTestcase): (empty)

20 async def test_partition_match( (empty)

21 self, 

22 queue: str, 

23 ) -> None: 

24 broker = self.get_broker() (empty)

25 

26 @broker.subscriber(partitions=[TopicPartition(queue, 1)]) (empty)

27 async def m(msg) -> None: (empty)

28 pass (empty)

29 

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

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

32 

33 m.mock.assert_called_once_with("hello") (empty)

34 

35 async def test_partition_match_exect( (empty)

36 self, 

37 queue: str, 

38 ) -> None: 

39 broker = self.get_broker() (empty)

40 

41 @broker.subscriber(partitions=[TopicPartition(queue, 1)]) (empty)

42 async def m(msg) -> None: (empty)

43 pass (empty)

44 

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

46 await br.publish("hello", queue, partition=1) (empty)

47 

48 m.mock.assert_called_once_with("hello") (empty)

49 

50 async def test_partition_mismatch( (empty)

51 self, 

52 queue: str, 

53 ) -> None: 

54 broker = self.get_broker() (empty)

55 

56 @broker.subscriber(partitions=[TopicPartition(queue, 1)]) (empty)

57 async def m(msg) -> None: (empty)

58 pass 

59 

60 @broker.subscriber(queue) (empty)

61 async def m2(msg) -> None: (empty)

62 pass (empty)

63 

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

65 await br.publish("hello", queue, partition=2) (empty)

66 

67 assert not m.mock.called (empty)

68 m2.mock.assert_called_once_with("hello") (empty)

69 

70 async def test_message_nack_seek( (empty)

71 self, 

72 queue: str, 

73 ) -> None: 

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

75 

76 @broker.subscriber(queue, group_id=f"{queue}1", ack_policy=AckPolicy.MANUAL) (empty)

77 async def m(msg: KafkaMessage) -> None: (empty)

78 await msg.nack() (empty)

79 

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

81 with patch.object( (empty)

82 FAKE_CONSUMER, 

83 "seek", 

84 spy_decorator(FAKE_CONSUMER.seek), 

85 ) as mocked: 

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

87 mocked.mock.assert_called_once() (empty)

88 

89 async def test_publisher_autoflush_mock( (empty)

90 self, 

91 queue: str, 

92 ) -> None: 

93 broker = self.get_broker() (empty)

94 

95 publisher = broker.publisher(queue + "1", autoflush=True) (empty)

96 publisher.flush = AsyncMock() (empty)

97 

98 @publisher (empty)

99 @broker.subscriber(queue) (empty)

100 async def m(msg): (empty)

101 return 1 (empty)

102 

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

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

105 

106 m.mock.assert_called_once_with("hello") (empty)

107 publisher.mock.assert_called_once_with(1) (empty)

108 

109 publisher.flush.assert_awaited_once() (empty)

110 

111 async def test_batch_publisher_autoflush_mock( (empty)

112 self, 

113 queue: str, 

114 ) -> None: 

115 broker = self.get_broker() (empty)

116 

117 publisher = broker.publisher(queue + "1", batch=True, autoflush=True) (empty)

118 publisher.flush = AsyncMock() (empty)

119 

120 @publisher (empty)

121 @broker.subscriber(queue) (empty)

122 async def m(msg): (empty)

123 return 1, 2, 3 (empty)

124 

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

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

127 

128 m.mock.assert_called_once_with("hello") (empty)

129 publisher.mock.assert_called_once_with([1, 2, 3]) (empty)

130 

131 publisher.flush.assert_awaited_once() (empty)

132 

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

134 async def test_with_real_testclient( (empty)

135 self, 

136 queue: str, 

137 ) -> None: 

138 event = asyncio.Event() (empty)

139 

140 broker = self.get_broker() (empty)

141 

142 @broker.subscriber(queue) (empty)

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

144 event.set() (empty)

145 

146 async with self.patch_broker(broker, with_real=True) as br: (empty)

147 await asyncio.wait( (empty)

148 ( 

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

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

151 ), 

152 timeout=3, 

153 ) 

154 

155 assert event.is_set() (empty)

156 

157 async def test_batch_pub_by_default_pub( (empty)

158 self, 

159 queue: str, 

160 ) -> None: 

161 broker = self.get_broker() (empty)

162 

163 @broker.subscriber(queue, batch=True) (empty)

164 async def m(msg) -> None: (empty)

165 pass (empty)

166 

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

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

169 m.mock.assert_called_once_with(["hello"]) (empty)

170 

171 async def test_batch_pub_by_pub_batch( (empty)

172 self, 

173 queue: str, 

174 ) -> None: 

175 broker = self.get_broker() (empty)

176 

177 @broker.subscriber(queue, batch=True) (empty)

178 async def m(msg) -> None: (empty)

179 pass (empty)

180 

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

182 await br.publish_batch("hello", topic=queue) (empty)

183 m.mock.assert_called_once_with(["hello"]) (empty)

184 

185 async def test_batch_publisher_mock( (empty)

186 self, 

187 queue: str, 

188 ) -> None: 

189 broker = self.get_broker() (empty)

190 

191 publisher = broker.publisher(queue + "1", batch=True) (empty)

192 

193 @publisher (empty)

194 @broker.subscriber(queue) (empty)

195 async def m(msg): (empty)

196 return 1, 2, 3 (empty)

197 

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

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

200 m.mock.assert_called_once_with("hello") (empty)

201 publisher.mock.assert_called_once_with([1, 2, 3]) (empty)

202 

203 async def test_respect_middleware(self, queue: str) -> None: (empty)

204 routes = [] (empty)

205 

206 class Middleware(BaseMiddleware): (empty)

207 async def on_receive(self) -> None: (empty)

208 routes.append(None) (empty)

209 return await super().on_receive() (empty)

210 

211 broker = self.get_broker(middlewares=(Middleware,)) (empty)

212 

213 @broker.subscriber(queue) (empty)

214 async def h1(msg) -> None: ... (empty)

215 

216 @broker.subscriber(queue + "1") (empty)

217 async def h2(msg) -> None: ... (empty)

218 

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

220 await br.publish("", queue) (empty)

221 await br.publish("", queue + "1") (empty)

222 

223 assert len(routes) == 2 (empty)

224 

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

226 async def test_real_respect_middleware(self, queue: str) -> None: (empty)

227 routes = [] (empty)

228 

229 class Middleware(BaseMiddleware): (empty)

230 async def on_receive(self) -> None: (empty)

231 routes.append(None) (empty)

232 return await super().on_receive() (empty)

233 

234 broker = self.get_broker(middlewares=(Middleware,)) (empty)

235 

236 @broker.subscriber(queue) (empty)

237 async def h1(msg) -> None: ... (empty)

238 

239 @broker.subscriber(queue + "1") (empty)

240 async def h2(msg) -> None: ... (empty)

241 

242 async with self.patch_broker(broker, with_real=True) as br: (empty)

243 await br.publish("", queue) (empty)

244 await br.publish("", queue + "1") (empty)

245 await h1.wait_call(3) (empty)

246 await h2.wait_call(3) (empty)

247 

248 assert len(routes) == 2 (empty)

249 

250 async def test_multiple_subscribers_different_groups( (empty)

251 self, 

252 queue: str, 

253 ) -> None: 

254 test_broker = self.get_broker() (empty)

255 

256 @test_broker.subscriber(queue, group_id="group1") (empty)

257 async def subscriber1(msg) -> None: ... (empty)

258 

259 @test_broker.subscriber(queue, group_id="group2") (empty)

260 async def subscriber2(msg) -> None: ... (empty)

261 

262 async with self.patch_broker(test_broker) as br: (empty)

263 await br.start() (empty)

264 await br.publish("", queue) (empty)

265 

266 assert subscriber1.mock.call_count == 1 (empty)

267 assert subscriber2.mock.call_count == 1 (empty)

268 

269 async def test_multiple_subscribers_same_group(self, queue: str) -> None: (empty)

270 broker = self.get_broker() (empty)

271 

272 @broker.subscriber(queue, group_id="group1") (empty)

273 async def subscriber1(msg) -> None: ... (empty)

274 

275 @broker.subscriber(queue, group_id="group1") (empty)

276 async def subscriber2(msg) -> None: ... (empty)

277 

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

279 await br.start() (empty)

280 await br.publish("", queue) (empty)

281 

282 # we can't guarantee the order of calls 

283 assert {subscriber1.mock.call_count, subscriber2.mock.call_count} == {1, 0} (empty)

284 

285 async def test_multiple_batch_subscriber_with_different_group( (empty)

286 self, 

287 queue: str, 

288 ) -> None: 

289 broker = self.get_broker() (empty)

290 

291 @broker.subscriber(queue, batch=True, group_id="group1") (empty)

292 async def subscriber1(msg) -> None: ... (empty)

293 

294 @broker.subscriber(queue, batch=True, group_id="group2") (empty)

295 async def subscriber2(msg) -> None: ... (empty)

296 

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

298 await br.start() (empty)

299 await br.publish("", queue) (empty)

300 

301 assert subscriber1.mock.call_count == 1 (empty)

302 assert subscriber2.mock.call_count == 1 (empty)

303 

304 async def test_multiple_batch_subscriber_with_same_group( (empty)

305 self, 

306 queue: str, 

307 ) -> None: 

308 broker = self.get_broker() (empty)

309 

310 @broker.subscriber(queue, batch=True, group_id="group1") (empty)

311 async def subscriber1(msg) -> None: ... (empty)

312 

313 @broker.subscriber(queue, batch=True, group_id="group1") (empty)

314 async def subscriber2(msg) -> None: ... (empty)

315 

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

317 await br.start() (empty)

318 await br.publish("", queue) (empty)

319 

320 # we can't guarantee the order of calls 

321 assert {subscriber1.mock.call_count, subscriber2.mock.call_count} == {1, 0} (empty)

322 

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

324 async def test_broker_gets_patched_attrs_within_cm(self) -> None: (empty)

325 await super().test_broker_gets_patched_attrs_within_cm(FakeProducer) (empty)

326 

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

328 async def test_broker_with_real_doesnt_get_patched(self) -> None: (empty)

329 await super().test_broker_with_real_doesnt_get_patched() (empty)

330 

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

332 async def test_broker_with_real_patches_publishers_and_subscribers( (empty)

333 self, 

334 queue: str, 

335 ) -> None: 

336 await super().test_broker_with_real_patches_publishers_and_subscribers(queue) (empty)

337 

338 @pytest.mark.xfail(reason="https://github.com/ag2ai/faststream/issues/2513") (empty)

339 async def test_publisher_without_destination(self) -> None: (empty)

340 """Fixes https://github.com/ag2ai/faststream/issues/2513.""" 

341 broker = self.get_broker() (empty)

342 

343 # use two publishers to check that we don't have conflicts 

344 publisher = broker.publisher(topic="") (empty)

345 another_publisher = broker.publisher(topic="") (empty)

346 

347 async with self.patch_broker(broker): (empty)

348 await publisher.publish(None, topic="new-key") (empty)

349 publisher.mock.assert_called_once() (empty)

350 

351 await another_publisher.publish(None, topic="new-key") 

352 another_publisher.mock.assert_called_once()