Fixing flaky tests - part 2 (#3838)

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/helpers.py

@@ -0,0 +1,47 @@
+import logging
+from time import sleep
+from typing import Callable
+
+
+def wait_for_condition(
+    predicate: Callable[[], bool],
+    timeout: float = 0.2,
+    check_interval: float = 0.01,
+    error_message: str = "Timeout waiting for condition",
+) -> None:
+    """
+    Poll a condition until it becomes True or timeout is reached.
+
+    Args:
+        predicate: A callable that returns True when the condition is met
+        timeout: Maximum time to wait in seconds (default: 0.2s = 20 * 0.01s)
+        check_interval: Time to sleep between checks in seconds (default: 0.01s)
+        error_message: Error message to raise if timeout occurs
+
+    Raises:
+        AssertionError: If the condition is not met within the timeout period
+
+    Example:
+        # Wait for circuit breaker to open
+        wait_for_condition(
+            lambda: cb2.state == CBState.OPEN,
+            timeout=0.2,
+            error_message="Timeout waiting for cb2 to open"
+        )
+
+        # Wait for failover strategy to select a specific database
+        wait_for_condition(
+            lambda: client.command_executor.active_database is mock_db,
+            timeout=0.2,
+            error_message="Timeout waiting for active database to change"
+        )
+    """
+    max_retries = int(timeout / check_interval)
+
+    for attempt in range(max_retries):
+        if predicate():
+            logging.debug(f"Condition met after {attempt} attempts")
+            return
+        sleep(check_interval)
+
+    raise AssertionError(error_message)

tests/test_asyncio/helpers.py

@@ -0,0 +1,47 @@
+import asyncio
+import logging
+from typing import Callable
+
+
+async def wait_for_condition(
+    predicate: Callable[[], bool],
+    timeout: float = 0.2,
+    check_interval: float = 0.01,
+    error_message: str = "Timeout waiting for condition",
+) -> None:
+    """
+    Poll a condition until it becomes True or timeout is reached.
+
+    Args:
+        predicate: A callable that returns True when the condition is met
+        timeout: Maximum time to wait in seconds (default: 0.2s = 20 * 0.01s)
+        check_interval: Time to sleep between checks in seconds (default: 0.01s)
+        error_message: Error message to raise if timeout occurs
+
+    Raises:
+        AssertionError: If the condition is not met within the timeout period
+
+    Example:
+        # Wait for circuit breaker to open
+        await wait_for_condition(
+            lambda: cb2.state == CBState.OPEN,
+            timeout=0.2,
+            error_message="Timeout waiting for cb2 to open"
+        )
+
+        # Wait for failover strategy to select a specific database
+        await wait_for_condition(
+            lambda: client.command_executor.active_database is mock_db,
+            timeout=0.2,
+            error_message="Timeout waiting for active database to change"
+        )
+    """
+    max_retries = int(timeout / check_interval)
+
+    for attempt in range(max_retries):
+        if predicate():
+            logging.debug(f"Condition met after {attempt} attempts")
+            return
+        await asyncio.sleep(check_interval)
+
+    raise AssertionError(error_message)