@@ -34,12 +34,10 @@ async def test_execute_command_against_correct_db_on_successful_initialization(
3434 self , mock_multi_db_config , mock_db , mock_db1 , mock_db2 , mock_hc
3535 ):
3636 databases = create_weighted_list (mock_db , mock_db1 , mock_db2 )
37+ mock_multi_db_config .health_checks = [mock_hc ]
3738
3839 with (
3940 patch .object (mock_multi_db_config , "databases" , return_value = databases ),
40- patch .object (
41- mock_multi_db_config , "default_health_checks" , return_value = [mock_hc ]
42- ),
4341 ):
4442 mock_db1 .client .execute_command = AsyncMock (return_value = "OK1" )
4543
@@ -71,12 +69,10 @@ async def test_execute_command_against_correct_db_and_closed_circuit(
7169 self , mock_multi_db_config , mock_db , mock_db1 , mock_db2 , mock_hc
7270 ):
7371 databases = create_weighted_list (mock_db , mock_db1 , mock_db2 )
72+ mock_multi_db_config .health_checks = [mock_hc ]
7473
7574 with (
7675 patch .object (mock_multi_db_config , "databases" , return_value = databases ),
77- patch .object (
78- mock_multi_db_config , "default_health_checks" , return_value = [mock_hc ]
79- ),
8076 ):
8177 mock_db1 .client .execute_command = AsyncMock (return_value = "OK1" )
8278
@@ -187,14 +183,10 @@ async def mock_check_health(database):
187183 return True
188184
189185 mock_hc .check_health .side_effect = mock_check_health
186+ mock_multi_db_config .health_checks = [mock_hc ]
190187
191188 with (
192189 patch .object (mock_multi_db_config , "databases" , return_value = databases ),
193- patch .object (
194- mock_multi_db_config ,
195- "default_health_checks" ,
196- return_value = [mock_hc ],
197- ),
198190 ):
199191 mock_db .client .execute_command .return_value = "OK"
200192 mock_db1 .client .execute_command .return_value = "OK1"
@@ -264,14 +256,10 @@ async def mock_check_health(database):
264256 return True
265257
266258 mock_hc .check_health .side_effect = mock_check_health
259+ mock_multi_db_config .health_checks = [mock_hc ]
267260
268261 with (
269262 patch .object (mock_multi_db_config , "databases" , return_value = databases ),
270- patch .object (
271- mock_multi_db_config ,
272- "default_health_checks" ,
273- return_value = [mock_hc ],
274- ),
275263 ):
276264 mock_db .client .execute_command .return_value = "OK"
277265 mock_db1 .client .execute_command .return_value = "OK1"
@@ -287,6 +275,60 @@ async def mock_check_health(database):
287275 await asyncio .sleep (0.5 )
288276 assert await client .set ("key" , "value" ) == "OK1"
289277
278+ @pytest .mark .asyncio
279+ @pytest .mark .parametrize (
280+ "mock_multi_db_config,mock_db, mock_db1, mock_db2" ,
281+ [
282+ (
283+ {"health_check_probes" : 1 },
284+ {"weight" : 0.2 , "circuit" : {"state" : CBState .CLOSED }},
285+ {"weight" : 0.7 , "circuit" : {"state" : CBState .CLOSED }},
286+ {"weight" : 0.5 , "circuit" : {"state" : CBState .CLOSED }},
287+ ),
288+ ],
289+ indirect = True ,
290+ )
291+ async def test_execute_command_do_not_auto_fallback_to_highest_weight_db (
292+ self , mock_multi_db_config , mock_db , mock_db1 , mock_db2 , mock_hc
293+ ):
294+ databases = create_weighted_list (mock_db , mock_db1 , mock_db2 )
295+ db1_counter = 0
296+ error_event = asyncio .Event ()
297+ check = False
298+
299+ async def mock_check_health (database ):
300+ nonlocal db1_counter , check
301+
302+ if database == mock_db1 and not check :
303+ db1_counter += 1
304+
305+ if db1_counter > 1 :
306+ error_event .set ()
307+ check = True
308+ return False
309+
310+ return True
311+
312+ mock_hc .check_health .side_effect = mock_check_health
313+ mock_multi_db_config .health_checks = [mock_hc ]
314+
315+ with (
316+ patch .object (mock_multi_db_config , "databases" , return_value = databases ),
317+ ):
318+ mock_db .client .execute_command .return_value = "OK"
319+ mock_db1 .client .execute_command .return_value = "OK1"
320+ mock_db2 .client .execute_command .return_value = "OK2"
321+ mock_multi_db_config .health_check_interval = 0.1
322+ mock_multi_db_config .auto_fallback_interval = - 1
323+ mock_multi_db_config .failover_strategy = WeightBasedFailoverStrategy ()
324+
325+ async with MultiDBClient (mock_multi_db_config ) as client :
326+ assert await client .set ("key" , "value" ) == "OK1"
327+ await error_event .wait ()
328+ assert await client .set ("key" , "value" ) == "OK2"
329+ await asyncio .sleep (0.5 )
330+ assert await client .set ("key" , "value" ) == "OK2"
331+
290332 @pytest .mark .asyncio
291333 @pytest .mark .parametrize (
292334 "mock_multi_db_config,mock_db, mock_db1, mock_db2" ,
@@ -304,12 +346,10 @@ async def test_execute_command_throws_exception_on_failed_initialization(
304346 self , mock_multi_db_config , mock_db , mock_db1 , mock_db2 , mock_hc
305347 ):
306348 databases = create_weighted_list (mock_db , mock_db1 , mock_db2 )
349+ mock_multi_db_config .health_checks = [mock_hc ]
307350
308351 with (
309352 patch .object (mock_multi_db_config , "databases" , return_value = databases ),
310- patch .object (
311- mock_multi_db_config , "default_health_checks" , return_value = [mock_hc ]
312- ),
313353 ):
314354 mock_hc .check_health .return_value = False
315355
@@ -340,12 +380,10 @@ async def test_add_database_throws_exception_on_same_database(
340380 self , mock_multi_db_config , mock_db , mock_db1 , mock_db2 , mock_hc
341381 ):
342382 databases = create_weighted_list (mock_db , mock_db1 , mock_db2 )
383+ mock_multi_db_config .health_checks = [mock_hc ]
343384
344385 with (
345386 patch .object (mock_multi_db_config , "databases" , return_value = databases ),
346- patch .object (
347- mock_multi_db_config , "default_health_checks" , return_value = [mock_hc ]
348- ),
349387 ):
350388 mock_hc .check_health .return_value = False
351389
@@ -373,12 +411,10 @@ async def test_add_database_makes_new_database_active(
373411 self , mock_multi_db_config , mock_db , mock_db1 , mock_db2 , mock_hc
374412 ):
375413 databases = create_weighted_list (mock_db , mock_db2 )
414+ mock_multi_db_config .health_checks = [mock_hc ]
376415
377416 with (
378417 patch .object (mock_multi_db_config , "databases" , return_value = databases ),
379- patch .object (
380- mock_multi_db_config , "default_health_checks" , return_value = [mock_hc ]
381- ),
382418 ):
383419 mock_db1 .client .execute_command .return_value = "OK1"
384420 mock_db2 .client .execute_command .return_value = "OK2"
@@ -413,12 +449,10 @@ async def test_remove_highest_weighted_database(
413449 self , mock_multi_db_config , mock_db , mock_db1 , mock_db2 , mock_hc
414450 ):
415451 databases = create_weighted_list (mock_db , mock_db1 , mock_db2 )
452+ mock_multi_db_config .health_checks = [mock_hc ]
416453
417454 with (
418455 patch .object (mock_multi_db_config , "databases" , return_value = databases ),
419- patch .object (
420- mock_multi_db_config , "default_health_checks" , return_value = [mock_hc ]
421- ),
422456 ):
423457 mock_db1 .client .execute_command .return_value = "OK1"
424458 mock_db2 .client .execute_command .return_value = "OK2"
@@ -451,12 +485,10 @@ async def test_update_database_weight_to_be_highest(
451485 self , mock_multi_db_config , mock_db , mock_db1 , mock_db2 , mock_hc
452486 ):
453487 databases = create_weighted_list (mock_db , mock_db1 , mock_db2 )
488+ mock_multi_db_config .health_checks = [mock_hc ]
454489
455490 with (
456491 patch .object (mock_multi_db_config , "databases" , return_value = databases ),
457- patch .object (
458- mock_multi_db_config , "default_health_checks" , return_value = [mock_hc ]
459- ),
460492 ):
461493 mock_db1 .client .execute_command .return_value = "OK1"
462494 mock_db2 .client .execute_command .return_value = "OK2"
@@ -491,12 +523,10 @@ async def test_add_new_failure_detector(
491523 self , mock_multi_db_config , mock_db , mock_db1 , mock_db2 , mock_hc
492524 ):
493525 databases = create_weighted_list (mock_db , mock_db1 , mock_db2 )
526+ mock_multi_db_config .health_checks = [mock_hc ]
494527
495528 with (
496529 patch .object (mock_multi_db_config , "databases" , return_value = databases ),
497- patch .object (
498- mock_multi_db_config , "default_health_checks" , return_value = [mock_hc ]
499- ),
500530 ):
501531 mock_db1 .client .execute_command .return_value = "OK1"
502532 mock_multi_db_config .event_dispatcher = EventDispatcher ()
@@ -552,12 +582,10 @@ async def test_add_new_health_check(
552582 self , mock_multi_db_config , mock_db , mock_db1 , mock_db2 , mock_hc
553583 ):
554584 databases = create_weighted_list (mock_db , mock_db1 , mock_db2 )
585+ mock_multi_db_config .health_checks = [mock_hc ]
555586
556587 with (
557588 patch .object (mock_multi_db_config , "databases" , return_value = databases ),
558- patch .object (
559- mock_multi_db_config , "default_health_checks" , return_value = [mock_hc ]
560- ),
561589 ):
562590 mock_db1 .client .execute_command .return_value = "OK1"
563591
@@ -594,12 +622,10 @@ async def test_set_active_database(
594622 self , mock_multi_db_config , mock_db , mock_db1 , mock_db2 , mock_hc
595623 ):
596624 databases = create_weighted_list (mock_db , mock_db1 , mock_db2 )
625+ mock_multi_db_config .health_checks = [mock_hc ]
597626
598627 with (
599628 patch .object (mock_multi_db_config , "databases" , return_value = databases ),
600- patch .object (
601- mock_multi_db_config , "default_health_checks" , return_value = [mock_hc ]
602- ),
603629 ):
604630 mock_db1 .client .execute_command .return_value = "OK1"
605631 mock_db .client .execute_command .return_value = "OK"
0 commit comments