Fixing flaky tests - part 2

Author: petyaslavova

tests/test_asyncio/test_multidb/test_client.py

@@ -1,4 +1,5 @@
 import asyncio
+import threading
 from unittest.mock import patch, AsyncMock, Mock
 
 import pybreaker
@@ -66,27 +67,63 @@ 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"
+            )
+
+            # Track health check runs across all databases
+            health_check_run = 0
+            run_lock = threading.Lock()
+
+            async def mock_check_health(database):
+                nonlocal health_check_run
+
+                # Increment run counter for each health check call
+                with run_lock:
+                    health_check_run += 1
+                    current_run = health_check_run
+
+                # Run 1 (health_check_run 1): mock_db unhealthy
+                if current_run == 1:
+                    if database == mock_db:
+                        return False
+                    return True
+
+                # Run 2 (health_check_run 2-3): mock_db1 and mock_db2 healthy
+                elif current_run <= 3:
+                    return True
+
+                # Run 3 (health_check_run 4-6): All databases healthy
+                elif current_run <= 6:
+                    return True
 
-            mock_hc.check_health.side_effect = [
-                False,
-                True,
-                True,
-                True,
-                True,
-                True,
-                True,
-            ]
+                # Run 4 (health_check_run 7): mock_db2 unhealthy
+                else:
+                    if database == mock_db2:
+                        return False
+                    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_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,68 @@ 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"
+            )
+
+            # Track health check runs across all databases
+            health_check_run = 0
+            run_lock = threading.Lock()
+
+            def mock_check_health(database):
+                nonlocal health_check_run
+
+                # Increment run counter for each health check call
+                with run_lock:
+                    health_check_run += 1
+                    current_run = health_check_run
+
+                # Run 1 (health_check_run 1): mock_db unhealthy
+                if current_run == 1:
+                    if database == mock_db:
+                        return False
+                    return True
+
+                # Run 2 (health_check_run 2-3): mock_db1 and mock_db2 healthy
+                elif current_run <= 3:
+                    return True
+
+                # Run 3 (health_check_run 4-6): All databases healthy
+                elif current_run <= 6:
+                    return True
+
+                # Run 4 (health_check_run 7): mock_db2 unhealthy
+                elif current_run == 7:
+                    if database == mock_db2:
+                        return False
+                    return True
+
+                # Background health checks (calls 8+): All healthy
+                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