Fixing flaky tests - part 2

Author: petyaslavova

dev_requirements.txt

@@ -12,6 +12,7 @@ pytest==8.3.4 ; platform_python_implementation == "PyPy"
 pytest-asyncio>=0.23.0
 pytest-asyncio==1.1.0 ; platform_python_implementation == "PyPy"
 pytest-cov
+coverage<7.11.1
 pytest-cov==6.0.0 ; platform_python_implementation == "PyPy"
 coverage==7.6.12 ; platform_python_implementation == "PyPy"
 pytest-profiling==1.8.1

redis/multidb/client.py

@@ -301,7 +301,13 @@ def _on_circuit_state_change_callback(
             )
 
     def close(self):
-        self.command_executor.active_database.client.close()
+        """
+        Closes the client and all its resources.
+        """
+        if self._bg_scheduler:
+            self._bg_scheduler.stop()
+        if self.command_executor.active_database:
+            self.command_executor.active_database.client.close()
 
 
 def _half_open_circuit(circuit: CircuitBreaker):

tests/test_asyncio/test_multidb/test_client.py

@@ -41,14 +41,16 @@ async def test_execute_command_against_correct_db_on_successful_initialization(
 
             mock_hc.check_health.return_value = True
 
-            client = MultiDBClient(mock_multi_db_config)
-            assert mock_multi_db_config.failover_strategy.set_databases.call_count == 1
-            assert await client.set("key", "value") == "OK1"
-            assert mock_hc.check_health.call_count == 9
+            async with MultiDBClient(mock_multi_db_config) as client:
+                assert (
+                    mock_multi_db_config.failover_strategy.set_databases.call_count == 1
+                )
+                assert await client.set("key", "value") == "OK1"
+                assert len(mock_hc.check_health.call_args_list) == 9
 
-            assert mock_db.circuit.state == CBState.CLOSED
-            assert mock_db1.circuit.state == CBState.CLOSED
-            assert mock_db2.circuit.state == CBState.CLOSED
+                assert mock_db.circuit.state == CBState.CLOSED
+                assert mock_db1.circuit.state == CBState.CLOSED
+                assert mock_db2.circuit.state == CBState.CLOSED
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
@@ -66,31 +68,44 @@ async def test_execute_command_against_correct_db_on_successful_initialization(
     async def test_execute_command_against_correct_db_and_closed_circuit(
         self, mock_multi_db_config, mock_db, mock_db1, mock_db2, mock_hc
     ):
+        """
+        Validates that commands are executed against the correct
+        database when one database becomes unhealthy during initialization.
+        Ensures the client selects the highest-weighted
+        healthy database (mock_db1) and executes commands against it
+        with a CLOSED circuit.
+        """
         databases = create_weighted_list(mock_db, mock_db1, mock_db2)
         mock_multi_db_config.health_checks = [mock_hc]
 
         with patch.object(mock_multi_db_config, "databases", return_value=databases):
+            mock_db.client.execute_command = AsyncMock(
+                return_value="NOT_OK-->Response from unexpected db - mock_db"
+            )
             mock_db1.client.execute_command = AsyncMock(return_value="OK1")
+            mock_db2.client.execute_command = AsyncMock(
+                return_value="NOT_OK-->Response from unexpected db - mock_db2"
+            )
 
-            mock_hc.check_health.side_effect = [
-                False,
-                True,
-                True,
-                True,
-                True,
-                True,
-                True,
-            ]
+            async def mock_check_health(database):
+                if database == mock_db2:
+                    return False
+                else:
+                    return True
 
-            client = MultiDBClient(mock_multi_db_config)
-            assert mock_multi_db_config.failover_strategy.set_databases.call_count == 1
-            result = await client.set("key", "value")
-            assert result == "OK1"
-            assert mock_hc.check_health.call_count == 7
+            mock_hc.check_health.side_effect = mock_check_health
 
-            assert mock_db.circuit.state == CBState.CLOSED
-            assert mock_db1.circuit.state == CBState.CLOSED
-            assert mock_db2.circuit.state == CBState.OPEN
+            async with MultiDBClient(mock_multi_db_config) as client:
+                assert (
+                    mock_multi_db_config.failover_strategy.set_databases.call_count == 1
+                )
+                result = await client.set("key", "value")
+                assert result == "OK1"
+                assert mock_hc.check_health.call_count >= 7
+
+                assert mock_db.circuit.state == CBState.CLOSED
+                assert mock_db1.circuit.state == CBState.CLOSED
+                assert mock_db2.circuit.state == CBState.OPEN
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
@@ -189,40 +204,40 @@ async def mock_check_health(database):
             mock_multi_db_config.health_check_interval = 0.1
             mock_multi_db_config.failover_strategy = WeightBasedFailoverStrategy()
 
-            client = MultiDBClient(mock_multi_db_config)
-            assert await client.set("key", "value") == "OK1"
+            async with MultiDBClient(mock_multi_db_config) as client:
+                assert await client.set("key", "value") == "OK1"
 
-            # Wait for mock_db1 to become unhealthy
-            assert await db1_became_unhealthy.wait(), (
-                "Timeout waiting for mock_db1 to become unhealthy"
-            )
+                # Wait for mock_db1 to become unhealthy
+                assert await db1_became_unhealthy.wait(), (
+                    "Timeout waiting for mock_db1 to become unhealthy"
+                )
 
-            await asyncio.sleep(0.01)
+                await asyncio.sleep(0.01)
 
-            assert await client.set("key", "value") == "OK2"
+                assert await client.set("key", "value") == "OK2"
 
-            # Wait for mock_db2 to become unhealthy
-            assert await db2_became_unhealthy.wait(), (
-                "Timeout waiting for mock_db2 to become unhealthy"
-            )
+                # Wait for mock_db2 to become unhealthy
+                assert await db2_became_unhealthy.wait(), (
+                    "Timeout waiting for mock_db2 to become unhealthy"
+                )
 
-            # Wait for circuit breaker state to actually reflect the unhealthy status
-            # (instead of just sleeping)
-            max_retries = 20
-            for _ in range(max_retries):
-                if cb2.state == CBState.OPEN:  # Circuit is open (unhealthy)
-                    break
-                await asyncio.sleep(0.01)
+                # Wait for circuit breaker state to actually reflect the unhealthy status
+                # (instead of just sleeping)
+                max_retries = 20
+                for _ in range(max_retries):
+                    if cb2.state == CBState.OPEN:  # Circuit is open (unhealthy)
+                        break
+                    await asyncio.sleep(0.01)
 
-            assert await client.set("key", "value") == "OK"
+                assert await client.set("key", "value") == "OK"
 
-            # Wait for mock_db to become unhealthy
-            assert await db_became_unhealthy.wait(), (
-                "Timeout waiting for mock_db to become unhealthy"
-            )
-            await asyncio.sleep(0.01)
+                # Wait for mock_db to become unhealthy
+                assert await db_became_unhealthy.wait(), (
+                    "Timeout waiting for mock_db to become unhealthy"
+                )
+                await asyncio.sleep(0.01)
 
-            assert await client.set("key", "value") == "OK1"
+                assert await client.set("key", "value") == "OK1"
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
@@ -375,7 +390,7 @@ async def test_execute_command_throws_exception_on_failed_initialization(
                 match="Initial connection failed - no active database found",
             ):
                 await client.set("key", "value")
-                assert mock_hc.check_health.call_count == 9
+                assert len(mock_hc.check_health.call_args_list) == 9
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
@@ -404,7 +419,6 @@ async def test_add_database_throws_exception_on_same_database(
 
             with pytest.raises(ValueError, match="Given database already exists"):
                 await client.add_database(mock_db)
-                assert mock_hc.check_health.call_count == 9
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
@@ -431,16 +445,18 @@ async def test_add_database_makes_new_database_active(
 
             mock_hc.check_health.return_value = True
 
-            client = MultiDBClient(mock_multi_db_config)
-            assert mock_multi_db_config.failover_strategy.set_databases.call_count == 1
+            async with MultiDBClient(mock_multi_db_config) as client:
+                assert (
+                    mock_multi_db_config.failover_strategy.set_databases.call_count == 1
+                )
 
-            assert await client.set("key", "value") == "OK2"
-            assert mock_hc.check_health.call_count == 6
+                assert await client.set("key", "value") == "OK2"
+                assert len(mock_hc.check_health.call_args_list) == 6
 
-            await client.add_database(mock_db1)
-            assert mock_hc.check_health.call_count == 9
+                await client.add_database(mock_db1)
+                assert len(mock_hc.check_health.call_args_list) == 9
 
-            assert await client.set("key", "value") == "OK1"
+                assert await client.set("key", "value") == "OK1"
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
@@ -467,14 +483,16 @@ async def test_remove_highest_weighted_database(
 
             mock_hc.check_health.return_value = True
 
-            client = MultiDBClient(mock_multi_db_config)
-            assert mock_multi_db_config.failover_strategy.set_databases.call_count == 1
+            async with MultiDBClient(mock_multi_db_config) as client:
+                assert (
+                    mock_multi_db_config.failover_strategy.set_databases.call_count == 1
+                )
 
-            assert await client.set("key", "value") == "OK1"
-            assert mock_hc.check_health.call_count == 9
+                assert await client.set("key", "value") == "OK1"
+                assert len(mock_hc.check_health.call_args_list) == 9
 
-            await client.remove_database(mock_db1)
-            assert await client.set("key", "value") == "OK2"
+                await client.remove_database(mock_db1)
+                assert await client.set("key", "value") == "OK2"
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
@@ -501,16 +519,18 @@ async def test_update_database_weight_to_be_highest(
 
             mock_hc.check_health.return_value = True
 
-            client = MultiDBClient(mock_multi_db_config)
-            assert mock_multi_db_config.failover_strategy.set_databases.call_count == 1
+            async with MultiDBClient(mock_multi_db_config) as client:
+                assert (
+                    mock_multi_db_config.failover_strategy.set_databases.call_count == 1
+                )
 
-            assert await client.set("key", "value") == "OK1"
-            assert mock_hc.check_health.call_count == 9
+                assert await client.set("key", "value") == "OK1"
+                assert len(mock_hc.check_health.call_args_list) == 9
 
-            await client.update_database_weight(mock_db2, 0.8)
-            assert mock_db2.weight == 0.8
+                await client.update_database_weight(mock_db2, 0.8)
+                assert mock_db2.weight == 0.8
 
-            assert await client.set("key", "value") == "OK2"
+                assert await client.set("key", "value") == "OK2"
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
@@ -544,30 +564,32 @@ async def test_add_new_failure_detector(
 
             mock_hc.check_health.return_value = True
 
-            client = MultiDBClient(mock_multi_db_config)
-            assert mock_multi_db_config.failover_strategy.set_databases.call_count == 1
-            assert await client.set("key", "value") == "OK1"
-            assert mock_hc.check_health.call_count == 9
-
-            # Simulate failing command events that lead to a failure detection
-            for _ in range(5):
-                await mock_multi_db_config.event_dispatcher.dispatch_async(
-                    command_fail_event
+            async with MultiDBClient(mock_multi_db_config) as client:
+                assert (
+                    mock_multi_db_config.failover_strategy.set_databases.call_count == 1
                 )
+                assert await client.set("key", "value") == "OK1"
+                assert len(mock_hc.check_health.call_args_list) == 9
 
-            assert mock_fd.register_failure.call_count == 5
+                # Simulate failing command events that lead to a failure detection
+                for _ in range(5):
+                    await mock_multi_db_config.event_dispatcher.dispatch_async(
+                        command_fail_event
+                    )
 
-            another_fd = Mock(spec=AsyncFailureDetector)
-            client.add_failure_detector(another_fd)
+                assert mock_fd.register_failure.call_count == 5
 
-            # Simulate failing command events that lead to a failure detection
-            for _ in range(5):
-                await mock_multi_db_config.event_dispatcher.dispatch_async(
-                    command_fail_event
-                )
+                another_fd = Mock(spec=AsyncFailureDetector)
+                client.add_failure_detector(another_fd)
+
+                # Simulate failing command events that lead to a failure detection
+                for _ in range(5):
+                    await mock_multi_db_config.event_dispatcher.dispatch_async(
+                        command_fail_event
+                    )
 
-            assert mock_fd.register_failure.call_count == 10
-            assert another_fd.register_failure.call_count == 5
+                assert mock_fd.register_failure.call_count == 10
+                assert another_fd.register_failure.call_count == 5
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
@@ -593,19 +615,21 @@ async def test_add_new_health_check(
 
             mock_hc.check_health.return_value = True
 
-            client = MultiDBClient(mock_multi_db_config)
-            assert mock_multi_db_config.failover_strategy.set_databases.call_count == 1
-            assert await client.set("key", "value") == "OK1"
-            assert mock_hc.check_health.call_count == 9
+            async with MultiDBClient(mock_multi_db_config) as client:
+                assert (
+                    mock_multi_db_config.failover_strategy.set_databases.call_count == 1
+                )
+                assert await client.set("key", "value") == "OK1"
+                assert len(mock_hc.check_health.call_args_list) == 9
 
-            another_hc = Mock(spec=HealthCheck)
-            another_hc.check_health.return_value = True
+                another_hc = Mock(spec=HealthCheck)
+                another_hc.check_health.return_value = True
 
-            await client.add_health_check(another_hc)
-            await client._check_db_health(mock_db1)
+                await client.add_health_check(another_hc)
+                await client._check_db_health(mock_db1)
 
-            assert mock_hc.check_health.call_count == 12
-            assert another_hc.check_health.call_count == 3
+                assert len(mock_hc.check_health.call_args_list) == 12
+                assert len(another_hc.check_health.call_args_list) == 3
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
@@ -632,23 +656,26 @@ async def test_set_active_database(
 
             mock_hc.check_health.return_value = True
 
-            client = MultiDBClient(mock_multi_db_config)
-            assert mock_multi_db_config.failover_strategy.set_databases.call_count == 1
-            assert await client.set("key", "value") == "OK1"
-            assert mock_hc.check_health.call_count == 9
+            async with MultiDBClient(mock_multi_db_config) as client:
+                assert (
+                    mock_multi_db_config.failover_strategy.set_databases.call_count == 1
+                )
+                assert await client.set("key", "value") == "OK1"
+                assert len(mock_hc.check_health.call_args_list) >= 9
+                assert mock_hc.check_health.call_count >= 9
 
-            await client.set_active_database(mock_db)
-            assert await client.set("key", "value") == "OK"
+                await client.set_active_database(mock_db)
+                assert await client.set("key", "value") == "OK"
 
-            with pytest.raises(
-                ValueError, match="Given database is not a member of database list"
-            ):
-                await client.set_active_database(Mock(spec=AsyncDatabase))
+                with pytest.raises(
+                    ValueError, match="Given database is not a member of database list"
+                ):
+                    await client.set_active_database(Mock(spec=AsyncDatabase))
 
-            mock_hc.check_health.return_value = False
+                mock_hc.check_health.return_value = False
 
-            with pytest.raises(
-                NoValidDatabaseException,
-                match="Cannot set active database, database is unhealthy",
-            ):
-                await client.set_active_database(mock_db1)
+                with pytest.raises(
+                    NoValidDatabaseException,
+                    match="Cannot set active database, database is unhealthy",
+                ):
+                    await client.set_active_database(mock_db1)