test: address review findings in polling helpers

Author: vdusek

tests/integration/_utils.py

@@ -2,7 +2,6 @@
 
 import asyncio
 import inspect
-import logging
 import secrets
 import string
 import time
@@ -15,8 +14,6 @@
 if TYPE_CHECKING:
     from collections.abc import Awaitable, Callable
 
-logger = logging.getLogger(__name__)
-
 # Environment variable names for test configuration
 TOKEN_ENV_VAR = 'APIFY_TEST_USER_API_TOKEN'
 TOKEN_ENV_VAR_2 = 'APIFY_TEST_USER_2_API_TOKEN'
@@ -104,7 +101,7 @@ async def maybe_await(value: Awaitable[T] | T) -> T:
     """
     if inspect.isawaitable(value):
         return await cast('Awaitable[T]', value)
-    return cast('T', value)
+    return value
 
 
 async def maybe_sleep(seconds: float, *, is_async: bool) -> None:
@@ -115,53 +112,6 @@ async def maybe_sleep(seconds: float, *, is_async: bool) -> None:
         time.sleep(seconds)  # noqa: ASYNC251
 
 
-@overload
-async def call_with_exp_backoff(
-    fn: Callable[[], Awaitable[T]],
-    condition: Callable[[T], bool] = ...,
-    *,
-    max_retries: int = ...,
-    base_delay: float = ...,
-) -> T: ...
-@overload
-async def call_with_exp_backoff(
-    fn: Callable[[], T],
-    condition: Callable[[T], bool] = ...,
-    *,
-    max_retries: int = ...,
-    base_delay: float = ...,
-) -> T: ...
-async def call_with_exp_backoff(
-    fn: Callable[[], Awaitable[T] | T],
-    condition: Callable[[T], bool] = bool,
-    *,
-    max_retries: int = 5,
-    base_delay: float = 1.0,
-) -> T:
-    """Call `fn`, retrying with exponential backoff until `condition(result)` is True.
-
-    Calls `fn` and checks whether `condition` holds for its result. If it does not, `fn` is retried up to
-    `max_retries` times, sleeping `base_delay * 2 ** attempt` seconds before each retry. The last result is
-    returned regardless of whether the condition was ever satisfied, so the caller can run its own assertion.
-
-    This is useful for eventually-consistent APIs where a freshly created resource may take a moment to become
-    visible. The default condition checks for a truthy result. Pass `max_retries=0` to call `fn` exactly once.
-
-    Unlike `poll_until_condition`, the delay between attempts grows exponentially rather than staying constant.
-    """
-    result = await maybe_await(fn())
-    for attempt in range(max_retries):
-        if condition(result):
-            return result
-        delay = base_delay * 2**attempt
-        logger.info(
-            'Condition not met for %r, retrying in %ss (attempt %d/%d).', result, delay, attempt + 1, max_retries
-        )
-        await asyncio.sleep(delay)
-        result = await maybe_await(fn())
-    return result
-
-
 @overload
 async def poll_until_condition(
     fn: Callable[[], Awaitable[T]],
@@ -192,8 +142,7 @@ async def poll_until_condition(
     assertion. The default condition checks for a truthy result.
 
     Use this instead of a fixed `asyncio.sleep` when waiting for eventually-consistent state (e.g. a freshly
-    created resource appearing in a listing) that may take a variable amount of time to propagate. Unlike
-    `call_with_exp_backoff`, the interval between polls stays constant.
+    created resource appearing in a listing) that may take a variable amount of time to propagate.
     """
     deadline = time.monotonic() + timeout
     result = await maybe_await(fn())
@@ -228,8 +177,8 @@ async def collect_iterate_until_present(
         iterator_factory: No-arg callable returning a fresh iterator on each call.
         expected_ids: IDs that must all appear in the collected items.
         item_type: Asserted to match the runtime type of each yielded item.
-        is_async: Whether the iterator is async.
-        max_attempts: Maximum number of polling rounds.
+        is_async: Whether the iterator is async (and so are sleeps).
+        max_attempts: Maximum number of polling rounds, guaranteed regardless of how long each drain takes.
         interval: Seconds to sleep between attempts.
 
     Returns:
@@ -251,12 +200,16 @@ async def drain() -> list[_HasIdT]:
                 collected.append(item)
         return collected
 
-    return await poll_until_condition(
-        drain,
-        lambda collected: expected_ids.issubset(item.id for item in collected),
-        timeout=max_attempts * interval,
-        poll_interval=interval,
-    )
+    # Loop on attempt count rather than `poll_until_condition`'s wall-clock deadline: each drain performs
+    # paginated HTTP calls, and charging that time against a deadline would shrink the number of retries
+    # under load — exactly when the eventual-consistency tolerance is needed most.
+    collected = await drain()
+    for _ in range(max_attempts - 1):
+        if expected_ids.issubset(item.id for item in collected):
+            break
+        await maybe_sleep(interval, is_async=is_async)
+        collected = await drain()
+    return collected
 
 
 # ============================================================================

tests/integration/test_request_queue.py

@@ -567,10 +567,13 @@ async def test_request_queue_unlock_requests(client: ApifyClient | ApifyClientAs
         # is acknowledged before all three lock writes are visible to subsequent reads, so unlocking
         # immediately can see fewer locks than were just acquired. Locked requests are excluded from the
         # queue head, so poll `list_head` until none of the locked requests reappear there, rather than
-        # guessing a fixed sleep.
+        # guessing a fixed sleep. Note this is a best-effort mitigation: the head read and the unlock may
+        # be served by different replicas, so convergence of one does not guarantee the other — but it
+        # narrows the race window enough in practice.
         async def all_locks_visible() -> bool:
             head = await maybe_await(rq_client.list_head(limit=5))
-            return isinstance(head, RequestQueueHead) and locked_ids.isdisjoint(item.id for item in head.items)
+            assert isinstance(head, RequestQueueHead)
+            return locked_ids.isdisjoint(item.id for item in head.items)
 
         await poll_until_condition(all_locks_visible, timeout=30, poll_interval=1)