Skip to content

Commit 204f1c0

Browse files
petyaslavovapetyaslavova
committed
Fixing with patch + one more test
1 parent 55dcfda commit 204f1c0

File tree

2 files changed

+55
-87
lines changed

2 files changed

+55
-87
lines changed

‎tests/test_asyncio/test_multidb/test_client.py‎

Lines changed: 28 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -36,9 +36,7 @@ async def test_execute_command_against_correct_db_on_successful_initialization(
3636
databases = create_weighted_list(mock_db, mock_db1, mock_db2)
3737
mock_multi_db_config.health_checks = [mock_hc]
3838

39-
with (
40-
patch.object(mock_multi_db_config, "databases", return_value=databases),
41-
):
39+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
4240
mock_db1.client.execute_command = AsyncMock(return_value="OK1")
4341

4442
mock_hc.check_health.return_value = True
@@ -71,9 +69,7 @@ async def test_execute_command_against_correct_db_and_closed_circuit(
7169
databases = create_weighted_list(mock_db, mock_db1, mock_db2)
7270
mock_multi_db_config.health_checks = [mock_hc]
7371

74-
with (
75-
patch.object(mock_multi_db_config, "databases", return_value=databases),
76-
):
72+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
7773
mock_db1.client.execute_command = AsyncMock(return_value="OK1")
7874

7975
mock_hc.check_health.side_effect = [
@@ -88,7 +84,8 @@ async def test_execute_command_against_correct_db_and_closed_circuit(
8884

8985
client = MultiDBClient(mock_multi_db_config)
9086
assert mock_multi_db_config.failover_strategy.set_databases.call_count == 1
91-
assert await client.set("key", "value") == "OK1"
87+
result = await client.set("key", "value")
88+
assert result == "OK1"
9289
assert mock_hc.check_health.call_count == 7
9390

9491
assert mock_db.circuit.state == CBState.CLOSED
@@ -185,9 +182,7 @@ async def mock_check_health(database):
185182
mock_hc.check_health.side_effect = mock_check_health
186183
mock_multi_db_config.health_checks = [mock_hc]
187184

188-
with (
189-
patch.object(mock_multi_db_config, "databases", return_value=databases),
190-
):
185+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
191186
mock_db.client.execute_command.return_value = "OK"
192187
mock_db1.client.execute_command.return_value = "OK1"
193188
mock_db2.client.execute_command.return_value = "OK2"
@@ -213,7 +208,7 @@ async def mock_check_health(database):
213208

214209
# Wait for circuit breaker state to actually reflect the unhealthy status
215210
# (instead of just sleeping)
216-
max_retries = 10
211+
max_retries = 20
217212
for _ in range(max_retries):
218213
if cb2.state == CBState.OPEN: # Circuit is open (unhealthy)
219214
break
@@ -266,9 +261,7 @@ async def mock_check_health(database):
266261
mock_hc.check_health.side_effect = mock_check_health
267262
mock_multi_db_config.health_checks = [mock_hc]
268263

269-
with (
270-
patch.object(mock_multi_db_config, "databases", return_value=databases),
271-
):
264+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
272265
mock_db.client.execute_command.return_value = "OK"
273266
mock_db1.client.execute_command.return_value = "OK1"
274267
mock_db2.client.execute_command.return_value = "OK2"
@@ -280,7 +273,7 @@ async def mock_check_health(database):
280273
assert await client.set("key", "value") == "OK1"
281274
await error_event.wait()
282275
# Wait for circuit breaker to actually open (not just the event)
283-
max_retries = 10
276+
max_retries = 20
284277
for _ in range(max_retries):
285278
if mock_db1.circuit.state == CBState.OPEN: # Circuit is open
286279
break
@@ -328,9 +321,7 @@ async def mock_check_health(database):
328321
mock_hc.check_health.side_effect = mock_check_health
329322
mock_multi_db_config.health_checks = [mock_hc]
330323

331-
with (
332-
patch.object(mock_multi_db_config, "databases", return_value=databases),
333-
):
324+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
334325
mock_db.client.execute_command.return_value = "OK"
335326
mock_db1.client.execute_command.return_value = "OK1"
336327
mock_db2.client.execute_command.return_value = "OK2"
@@ -341,6 +332,15 @@ async def mock_check_health(database):
341332
async with MultiDBClient(mock_multi_db_config) as client:
342333
assert await client.set("key", "value") == "OK1"
343334
await error_event.wait()
335+
# Wait for circuit breaker state to actually reflect the unhealthy status
336+
# (instead of just sleeping)
337+
max_retries = 20
338+
for _ in range(max_retries):
339+
if (
340+
mock_db1.circuit.state == CBState.OPEN
341+
): # Circuit is open (unhealthy)
342+
break
343+
await asyncio.sleep(0.01)
344344
assert await client.set("key", "value") == "OK2"
345345
await asyncio.sleep(0.5)
346346
assert await client.set("key", "value") == "OK2"
@@ -364,9 +364,7 @@ async def test_execute_command_throws_exception_on_failed_initialization(
364364
databases = create_weighted_list(mock_db, mock_db1, mock_db2)
365365
mock_multi_db_config.health_checks = [mock_hc]
366366

367-
with (
368-
patch.object(mock_multi_db_config, "databases", return_value=databases),
369-
):
367+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
370368
mock_hc.check_health.return_value = False
371369

372370
client = MultiDBClient(mock_multi_db_config)
@@ -398,9 +396,7 @@ async def test_add_database_throws_exception_on_same_database(
398396
databases = create_weighted_list(mock_db, mock_db1, mock_db2)
399397
mock_multi_db_config.health_checks = [mock_hc]
400398

401-
with (
402-
patch.object(mock_multi_db_config, "databases", return_value=databases),
403-
):
399+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
404400
mock_hc.check_health.return_value = False
405401

406402
client = MultiDBClient(mock_multi_db_config)
@@ -429,9 +425,7 @@ async def test_add_database_makes_new_database_active(
429425
databases = create_weighted_list(mock_db, mock_db2)
430426
mock_multi_db_config.health_checks = [mock_hc]
431427

432-
with (
433-
patch.object(mock_multi_db_config, "databases", return_value=databases),
434-
):
428+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
435429
mock_db1.client.execute_command.return_value = "OK1"
436430
mock_db2.client.execute_command.return_value = "OK2"
437431

@@ -467,9 +461,7 @@ async def test_remove_highest_weighted_database(
467461
databases = create_weighted_list(mock_db, mock_db1, mock_db2)
468462
mock_multi_db_config.health_checks = [mock_hc]
469463

470-
with (
471-
patch.object(mock_multi_db_config, "databases", return_value=databases),
472-
):
464+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
473465
mock_db1.client.execute_command.return_value = "OK1"
474466
mock_db2.client.execute_command.return_value = "OK2"
475467

@@ -503,9 +495,7 @@ async def test_update_database_weight_to_be_highest(
503495
databases = create_weighted_list(mock_db, mock_db1, mock_db2)
504496
mock_multi_db_config.health_checks = [mock_hc]
505497

506-
with (
507-
patch.object(mock_multi_db_config, "databases", return_value=databases),
508-
):
498+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
509499
mock_db1.client.execute_command.return_value = "OK1"
510500
mock_db2.client.execute_command.return_value = "OK2"
511501

@@ -541,9 +531,7 @@ async def test_add_new_failure_detector(
541531
databases = create_weighted_list(mock_db, mock_db1, mock_db2)
542532
mock_multi_db_config.health_checks = [mock_hc]
543533

544-
with (
545-
patch.object(mock_multi_db_config, "databases", return_value=databases),
546-
):
534+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
547535
mock_db1.client.execute_command.return_value = "OK1"
548536
mock_multi_db_config.event_dispatcher = EventDispatcher()
549537
mock_fd = mock_multi_db_config.failure_detectors[0]
@@ -562,7 +550,7 @@ async def test_add_new_failure_detector(
562550
assert mock_hc.check_health.call_count == 9
563551

564552
# Simulate failing command events that lead to a failure detection
565-
for i in range(5):
553+
for _ in range(5):
566554
await mock_multi_db_config.event_dispatcher.dispatch_async(
567555
command_fail_event
568556
)
@@ -573,7 +561,7 @@ async def test_add_new_failure_detector(
573561
client.add_failure_detector(another_fd)
574562

575563
# Simulate failing command events that lead to a failure detection
576-
for i in range(5):
564+
for _ in range(5):
577565
await mock_multi_db_config.event_dispatcher.dispatch_async(
578566
command_fail_event
579567
)
@@ -600,9 +588,7 @@ async def test_add_new_health_check(
600588
databases = create_weighted_list(mock_db, mock_db1, mock_db2)
601589
mock_multi_db_config.health_checks = [mock_hc]
602590

603-
with (
604-
patch.object(mock_multi_db_config, "databases", return_value=databases),
605-
):
591+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
606592
mock_db1.client.execute_command.return_value = "OK1"
607593

608594
mock_hc.check_health.return_value = True
@@ -640,9 +626,7 @@ async def test_set_active_database(
640626
databases = create_weighted_list(mock_db, mock_db1, mock_db2)
641627
mock_multi_db_config.health_checks = [mock_hc]
642628

643-
with (
644-
patch.object(mock_multi_db_config, "databases", return_value=databases),
645-
):
629+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
646630
mock_db1.client.execute_command.return_value = "OK1"
647631
mock_db.client.execute_command.return_value = "OK"
648632

‎tests/test_multidb/test_client.py‎

Lines changed: 27 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import threading
22
from time import sleep
3-
from unittest.mock import patch, Mock
3+
from unittest.mock import MagicMock, patch, Mock
44

55
import pybreaker
66
import pytest
@@ -36,9 +36,7 @@ def test_execute_command_against_correct_db_on_successful_initialization(
3636
databases = create_weighted_list(mock_db, mock_db1, mock_db2)
3737
mock_multi_db_config.health_checks = [mock_hc]
3838

39-
with (
40-
patch.object(mock_multi_db_config, "databases", return_value=databases),
41-
):
39+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
4240
mock_db1.client.execute_command.return_value = "OK1"
4341

4442
mock_hc.check_health.return_value = True
@@ -70,10 +68,8 @@ def test_execute_command_against_correct_db_and_closed_circuit(
7068
databases = create_weighted_list(mock_db, mock_db1, mock_db2)
7169
mock_multi_db_config.health_checks = [mock_hc]
7270

73-
with (
74-
patch.object(mock_multi_db_config, "databases", return_value=databases),
75-
):
76-
mock_db1.client.execute_command.return_value = "OK1"
71+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
72+
mock_db1.client.execute_command = MagicMock(return_value="OK1")
7773

7874
mock_hc.check_health.side_effect = [
7975
False,
@@ -87,7 +83,8 @@ def test_execute_command_against_correct_db_and_closed_circuit(
8783

8884
client = MultiDBClient(mock_multi_db_config)
8985
assert mock_multi_db_config.failover_strategy.set_databases.call_count == 1
90-
assert client.set("key", "value") == "OK1"
86+
result = client.set("key", "value")
87+
assert result == "OK1"
9188
assert mock_hc.check_health.call_count == 7
9289

9390
assert mock_db.circuit.state == CBState.CLOSED
@@ -183,9 +180,7 @@ def mock_check_health(database):
183180
mock_hc.check_health.side_effect = mock_check_health
184181
mock_multi_db_config.health_checks = [mock_hc]
185182

186-
with (
187-
patch.object(mock_multi_db_config, "databases", return_value=databases),
188-
):
183+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
189184
mock_multi_db_config.health_check_interval = 0.1
190185
mock_multi_db_config.failover_strategy = WeightBasedFailoverStrategy()
191186
mock_db.client.execute_command.return_value = "OK"
@@ -255,9 +250,7 @@ def mock_check_health(database):
255250
mock_hc.check_health.side_effect = mock_check_health
256251
mock_multi_db_config.health_checks = [mock_hc]
257252

258-
with (
259-
patch.object(mock_multi_db_config, "databases", return_value=databases),
260-
):
253+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
261254
mock_db.client.execute_command.return_value = "OK"
262255
mock_db1.client.execute_command.return_value = "OK1"
263256
mock_db2.client.execute_command.return_value = "OK2"
@@ -270,7 +263,7 @@ def mock_check_health(database):
270263
error_event.wait(timeout=0.5)
271264

272265
# Wait for circuit breaker to actually open (not just the event)
273-
max_retries = 10
266+
max_retries = 20
274267
for _ in range(max_retries):
275268
if mock_db1.circuit.state == CBState.OPEN: # Circuit is open
276269
break
@@ -317,9 +310,7 @@ def mock_check_health(database):
317310
mock_hc.check_health.side_effect = mock_check_health
318311
mock_multi_db_config.health_checks = [mock_hc]
319312

320-
with (
321-
patch.object(mock_multi_db_config, "databases", return_value=databases),
322-
):
313+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
323314
mock_db.client.execute_command.return_value = "OK"
324315
mock_db1.client.execute_command.return_value = "OK1"
325316
mock_db2.client.execute_command.return_value = "OK2"
@@ -330,6 +321,15 @@ def mock_check_health(database):
330321
client = MultiDBClient(mock_multi_db_config)
331322
assert client.set("key", "value") == "OK1"
332323
error_event.wait(timeout=0.5)
324+
# Wait for circuit breaker state to actually reflect the unhealthy status
325+
# (instead of just sleeping)
326+
max_retries = 20
327+
for _ in range(max_retries):
328+
if (
329+
mock_db1.circuit.state == CBState.OPEN
330+
): # Circuit is open (unhealthy)
331+
break
332+
sleep(0.01)
333333
assert client.set("key", "value") == "OK2"
334334
sleep(0.5)
335335
assert client.set("key", "value") == "OK2"
@@ -352,9 +352,7 @@ def test_execute_command_throws_exception_on_failed_initialization(
352352
databases = create_weighted_list(mock_db, mock_db1, mock_db2)
353353
mock_multi_db_config.health_checks = [mock_hc]
354354

355-
with (
356-
patch.object(mock_multi_db_config, "databases", return_value=databases),
357-
):
355+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
358356
mock_hc.check_health.return_value = False
359357

360358
client = MultiDBClient(mock_multi_db_config)
@@ -386,9 +384,7 @@ def test_add_database_throws_exception_on_same_database(
386384
databases = create_weighted_list(mock_db, mock_db1, mock_db2)
387385
mock_multi_db_config.health_checks = [mock_hc]
388386

389-
with (
390-
patch.object(mock_multi_db_config, "databases", return_value=databases),
391-
):
387+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
392388
mock_hc.check_health.return_value = False
393389

394390
client = MultiDBClient(mock_multi_db_config)
@@ -416,9 +412,7 @@ def test_add_database_makes_new_database_active(
416412
databases = create_weighted_list(mock_db, mock_db2)
417413
mock_multi_db_config.health_checks = [mock_hc]
418414

419-
with (
420-
patch.object(mock_multi_db_config, "databases", return_value=databases),
421-
):
415+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
422416
mock_db1.client.execute_command.return_value = "OK1"
423417
mock_db2.client.execute_command.return_value = "OK2"
424418

@@ -453,9 +447,7 @@ def test_remove_highest_weighted_database(
453447
databases = create_weighted_list(mock_db, mock_db1, mock_db2)
454448
mock_multi_db_config.health_checks = [mock_hc]
455449

456-
with (
457-
patch.object(mock_multi_db_config, "databases", return_value=databases),
458-
):
450+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
459451
mock_db1.client.execute_command.return_value = "OK1"
460452
mock_db2.client.execute_command.return_value = "OK2"
461453

@@ -489,9 +481,7 @@ def test_update_database_weight_to_be_highest(
489481
databases = create_weighted_list(mock_db, mock_db1, mock_db2)
490482
mock_multi_db_config.health_checks = [mock_hc]
491483

492-
with (
493-
patch.object(mock_multi_db_config, "databases", return_value=databases),
494-
):
484+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
495485
mock_db1.client.execute_command.return_value = "OK1"
496486
mock_db2.client.execute_command.return_value = "OK2"
497487

@@ -526,9 +516,7 @@ def test_add_new_failure_detector(
526516
databases = create_weighted_list(mock_db, mock_db1, mock_db2)
527517
mock_multi_db_config.health_checks = [mock_hc]
528518

529-
with (
530-
patch.object(mock_multi_db_config, "databases", return_value=databases),
531-
):
519+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
532520
mock_db1.client.execute_command.return_value = "OK1"
533521
mock_multi_db_config.event_dispatcher = EventDispatcher()
534522
mock_fd = mock_multi_db_config.failure_detectors[0]
@@ -580,9 +568,7 @@ def test_add_new_health_check(
580568
databases = create_weighted_list(mock_db, mock_db1, mock_db2)
581569
mock_multi_db_config.health_checks = [mock_hc]
582570

583-
with (
584-
patch.object(mock_multi_db_config, "databases", return_value=databases),
585-
):
571+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
586572
mock_db1.client.execute_command.return_value = "OK1"
587573

588574
mock_hc.check_health.return_value = True
@@ -619,9 +605,7 @@ def test_set_active_database(
619605
databases = create_weighted_list(mock_db, mock_db1, mock_db2)
620606
mock_multi_db_config.health_checks = [mock_hc]
621607

622-
with (
623-
patch.object(mock_multi_db_config, "databases", return_value=databases),
624-
):
608+
with patch.object(mock_multi_db_config, "databases", return_value=databases):
625609
mock_db1.client.execute_command.return_value = "OK1"
626610
mock_db.client.execute_command.return_value = "OK"
627611

0 commit comments

Comments
 (0)