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/asyncio/multidb/healthcheck.py

@@ -81,6 +81,7 @@ async def execute(self, health_checks: List[HealthCheck], database) -> bool:
         for health_check in health_checks:
             for attempt in range(self.health_check_probes):
                 try:
+                    print(f"Checking health of {database}")
                     if not await health_check.check_health(database):
                         return False
                 except Exception as e:

redis/multidb/healthcheck.py

@@ -80,6 +80,7 @@ def execute(self, health_checks: List[HealthCheck], database) -> bool:
         for health_check in health_checks:
             for attempt in range(self.health_check_probes):
                 try:
+                    print(f"Checking health of {database}")
                     if not health_check.check_health(database):
                         return False
                 except Exception as e:

tests/test_asyncio/test_multidb/test_client.py

@@ -66,27 +66,37 @@ 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
 
+            mock_hc.check_health.side_effect = mock_check_health
             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
+            assert mock_hc.check_health.call_count >= 7
 
             assert mock_db.circuit.state == CBState.CLOSED
             assert mock_db1.circuit.state == CBState.CLOSED

tests/test_asyncio/test_multidb/test_pipeline.py

@@ -299,7 +299,10 @@ async def callback(pipe: Pipeline):
                 pipe.get("key1")
 
             assert await client.transaction(callback) == ["OK1", "value1"]
-            assert mock_hc.check_health.call_count == 9
+            # if we assume at least 3 health checks have run per each database
+            # we should have at least 9 total calls
+            assert len(mock_hc.check_health.call_args_list) >= 9
+            assert mock_hc.check_health.call_count >= 9
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
@@ -327,15 +330,13 @@ async def test_execute_transaction_against_correct_db_and_closed_circuit(
         ):
             mock_db1.client.transaction.return_value = ["OK1", "value1"]
 
-            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
+
+            mock_hc.check_health.side_effect = mock_check_health
 
             client = MultiDBClient(mock_multi_db_config)
             assert mock_multi_db_config.failover_strategy.set_databases.call_count == 1
@@ -345,7 +346,7 @@ async def callback(pipe: Pipeline):
                 pipe.get("key1")
 
             assert await client.transaction(callback) == ["OK1", "value1"]
-            assert mock_hc.check_health.call_count == 7
+            assert mock_hc.check_health.call_count >= 7
 
             assert mock_db.circuit.state == CBState.CLOSED
             assert mock_db1.circuit.state == CBState.CLOSED

tests/test_auth/test_token_manager.py

@@ -174,13 +174,13 @@ def test_token_renewal_with_skip_initial(self):
         mock_provider.request_token.side_effect = [
             SimpleToken(
                 "value",
-                (datetime.now(timezone.utc).timestamp() * 1000) + 50,
+                (datetime.now(timezone.utc).timestamp() * 1000) + 1000,
                 (datetime.now(timezone.utc).timestamp() * 1000),
                 {"oid": "test"},
             ),
             SimpleToken(
                 "value",
-                (datetime.now(timezone.utc).timestamp() * 1000) + 150,
+                (datetime.now(timezone.utc).timestamp() * 1000) + 1500,
                 (datetime.now(timezone.utc).timestamp() * 1000),
                 {"oid": "test"},
             ),
@@ -194,12 +194,12 @@ def on_next(token):
         mock_listener.on_next = on_next
 
         retry_policy = RetryPolicy(3, 10)
-        config = TokenManagerConfig(1, 0, 1000, retry_policy)
+        config = TokenManagerConfig(0.5, 0, 1000, retry_policy)
         mgr = TokenManager(mock_provider, config)
         mgr.start(mock_listener, skip_initial=True)
-        # Should be less than a 0.1, or it will be flacky due to
-        # additional token renewal.
-        sleep(0.1)
+        assert len(tokens) == 0
+
+        sleep(0.6)
 
         assert len(tokens) > 0
 
@@ -210,19 +210,19 @@ async def test_async_token_renewal_with_skip_initial(self):
         mock_provider.request_token.side_effect = [
             SimpleToken(
                 "value",
-                (datetime.now(timezone.utc).timestamp() * 1000) + 100,
+                (datetime.now(timezone.utc).timestamp() * 1000) + 1000,
                 (datetime.now(timezone.utc).timestamp() * 1000),
                 {"oid": "test"},
             ),
             SimpleToken(
                 "value",
-                (datetime.now(timezone.utc).timestamp() * 1000) + 120,
+                (datetime.now(timezone.utc).timestamp() * 1000) + 1200,
                 (datetime.now(timezone.utc).timestamp() * 1000),
                 {"oid": "test"},
             ),
             SimpleToken(
                 "value",
-                (datetime.now(timezone.utc).timestamp() * 1000) + 140,
+                (datetime.now(timezone.utc).timestamp() * 1000) + 1400,
                 (datetime.now(timezone.utc).timestamp() * 1000),
                 {"oid": "test"},
             ),
@@ -236,13 +236,12 @@ async def on_next(token):
         mock_listener.on_next = on_next
 
         retry_policy = RetryPolicy(3, 10)
-        config = TokenManagerConfig(1, 0, 1000, retry_policy)
+        config = TokenManagerConfig(0.5, 0, 1000, retry_policy)
         mgr = TokenManager(mock_provider, config)
         await mgr.start_async(mock_listener, skip_initial=True)
-        # Should be less than a 0.1, or it will be flacky
-        # due to additional token renewal.
-        await asyncio.sleep(0.2)
+        assert len(tokens) == 0
 
+        await asyncio.sleep(0.6)
         assert len(tokens) > 0
 
     def test_success_token_renewal_with_retry(self):

tests/test_multidb/test_client.py

@@ -65,27 +65,38 @@ def test_execute_command_against_correct_db_on_successful_initialization(
     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 = MagicMock(
+                return_value="NOT_OK-->Response from unexpected db - mock_db"
+            )
             mock_db1.client.execute_command = MagicMock(return_value="OK1")
+            mock_db2.client.execute_command = MagicMock(
+                return_value="NOT_OK-->Response from unexpected db - mock_db2"
+            )
+
+            def mock_check_health(database):
+                if database == mock_db2:
+                    return False
+                else:
+                    return True
 
-            mock_hc.check_health.side_effect = [
-                False,
-                True,
-                True,
-                True,
-                True,
-                True,
-                True,
-            ]
+            mock_hc.check_health.side_effect = mock_check_health
 
             client = MultiDBClient(mock_multi_db_config)
             assert mock_multi_db_config.failover_strategy.set_databases.call_count == 1
             result = client.set("key", "value")
             assert result == "OK1"
-            assert mock_hc.check_health.call_count == 7
+            assert mock_hc.check_health.call_count >= 7
 
             assert mock_db.circuit.state == CBState.CLOSED
             assert mock_db1.circuit.state == CBState.CLOSED

tests/test_multidb/test_pipeline.py

@@ -89,15 +89,13 @@ def test_execute_pipeline_against_correct_db_and_closed_circuit(
             pipe.execute.return_value = ["OK1", "value1"]
             mock_db1.client.pipeline.return_value = pipe
 
-            mock_hc.check_health.side_effect = [
-                False,
-                True,
-                True,
-                True,
-                True,
-                True,
-                True,
-            ]
+            def mock_check_health(database):
+                if database == mock_db2:
+                    return False
+                else:
+                    return True
+
+            mock_hc.check_health.side_effect = mock_check_health
 
             client = MultiDBClient(mock_multi_db_config)
             assert mock_multi_db_config.failover_strategy.set_databases.call_count == 1
@@ -107,7 +105,7 @@ def test_execute_pipeline_against_correct_db_and_closed_circuit(
                 pipe.get("key1")
 
             assert pipe.execute() == ["OK1", "value1"]
-            assert mock_hc.check_health.call_count == 7
+            assert mock_hc.check_health.call_count >= 7
 
             assert mock_db.circuit.state == CBState.CLOSED
             assert mock_db1.circuit.state == CBState.CLOSED
@@ -298,7 +296,8 @@ def callback(pipe: Pipeline):
                 pipe.get("key1")
 
             assert client.transaction(callback) == ["OK1", "value1"]
-            assert mock_hc.check_health.call_count == 9
+            assert len(mock_hc.check_health.call_args_list) >= 9
+            assert mock_hc.check_health.call_count >= 9
 
     @pytest.mark.parametrize(
         "mock_multi_db_config,mock_db, mock_db1, mock_db2",