Fail fast on invalid polling callback return types

Co-authored-by: Shri Sukhani <shrisukhani@users.noreply.github.com>

Author: cursoragent

hyperbrowser/client/polling.py

@@ -1,4 +1,5 @@
 import asyncio
+import inspect
 import math
 from numbers import Real
 import time
@@ -60,6 +61,16 @@ def _ensure_status_string(status: object, *, operation_name: str) -> str:
     return status
 
 
+def _ensure_awaitable(
+    result: object, *, callback_name: str, operation_name: str
+) -> Awaitable[object]:
+    if not inspect.isawaitable(result):
+        raise HyperbrowserError(
+            f"{callback_name} must return an awaitable for {operation_name}"
+        )
+    return result
+
+
 def _validate_retry_config(
     *,
     max_attempts: int,
@@ -153,7 +164,7 @@ def poll_until_terminal_status(
 
     while True:
         try:
-            status = _ensure_status_string(get_status(), operation_name=operation_name)
+            status = get_status()
             failures = 0
         except Exception as exc:
             failures += 1
@@ -168,6 +179,7 @@ def poll_until_terminal_status(
             time.sleep(poll_interval_seconds)
             continue
 
+        status = _ensure_status_string(status, operation_name=operation_name)
         if _ensure_boolean_terminal_result(
             is_terminal_status(status), operation_name=operation_name
         ):
@@ -226,9 +238,25 @@ async def poll_until_terminal_status_async(
 
     while True:
         try:
-            status = _ensure_status_string(
-                await get_status(), operation_name=operation_name
-            )
+            status_result = get_status()
+        except Exception as exc:
+            failures += 1
+            if failures >= max_status_failures:
+                raise HyperbrowserPollingError(
+                    f"Failed to poll {operation_name} after {max_status_failures} attempts: {exc}"
+                ) from exc
+            if has_exceeded_max_wait(start_time, max_wait_seconds):
+                raise HyperbrowserTimeoutError(
+                    f"Timed out waiting for {operation_name} after {max_wait_seconds} seconds"
+                )
+            await asyncio.sleep(poll_interval_seconds)
+            continue
+
+        status_awaitable = _ensure_awaitable(
+            status_result, callback_name="get_status", operation_name=operation_name
+        )
+        try:
+            status = await status_awaitable
             failures = 0
         except Exception as exc:
             failures += 1
@@ -243,6 +271,7 @@ async def poll_until_terminal_status_async(
             await asyncio.sleep(poll_interval_seconds)
             continue
 
+        status = _ensure_status_string(status, operation_name=operation_name)
         if _ensure_boolean_terminal_result(
             is_terminal_status(status), operation_name=operation_name
         ):

tests/test_polling.py

@@ -894,6 +894,16 @@ async def validate_async_operation_name() -> None:
                 poll_interval_seconds=0.0,
                 max_wait_seconds=1.0,
             )
+        with pytest.raises(
+            HyperbrowserError, match="get_status must return an awaitable"
+        ):
+            await poll_until_terminal_status_async(
+                operation_name="invalid-status-awaitable-async",
+                get_status=lambda: "completed",  # type: ignore[return-value]
+                is_terminal_status=lambda value: value == "completed",
+                poll_interval_seconds=0.0,
+                max_wait_seconds=1.0,
+            )
         with pytest.raises(
             HyperbrowserError, match="operation_name must be 200 characters or fewer"
         ):