Fail fast on polling callback contract violations

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

Author: cursoragent

hyperbrowser/client/polling.py

@@ -15,6 +15,10 @@
 _MAX_OPERATION_NAME_LENGTH = 200
 
 
+class _NonRetryablePollingError(HyperbrowserError):
+    pass
+
+
 def _validate_non_negative_real(value: float, *, field_name: str) -> None:
     if isinstance(value, bool) or not isinstance(value, Real):
         raise HyperbrowserError(f"{field_name} must be a number")
@@ -52,7 +56,7 @@ def _ensure_boolean_terminal_result(result: object, *, operation_name: str) -> b
         result, callback_name="is_terminal_status", operation_name=operation_name
     )
     if not isinstance(result, bool):
-        raise HyperbrowserError(
+        raise _NonRetryablePollingError(
             f"is_terminal_status must return a boolean for {operation_name}"
         )
     return result
@@ -63,15 +67,17 @@ def _ensure_status_string(status: object, *, operation_name: str) -> str:
         status, callback_name="get_status", operation_name=operation_name
     )
     if not isinstance(status, str):
-        raise HyperbrowserError(f"get_status must return a string for {operation_name}")
+        raise _NonRetryablePollingError(
+            f"get_status must return a string for {operation_name}"
+        )
     return status
 
 
 def _ensure_awaitable(
     result: object, *, callback_name: str, operation_name: str
 ) -> Awaitable[object]:
     if not inspect.isawaitable(result):
-        raise HyperbrowserError(
+        raise _NonRetryablePollingError(
             f"{callback_name} must return an awaitable for {operation_name}"
         )
     return result
@@ -83,7 +89,7 @@ def _ensure_non_awaitable(
     if inspect.isawaitable(result):
         if inspect.iscoroutine(result):
             result.close()
-        raise HyperbrowserError(
+        raise _NonRetryablePollingError(
             f"{callback_name} must return a non-awaitable result for {operation_name}"
         )
 
@@ -182,9 +188,6 @@ def poll_until_terminal_status(
     while True:
         try:
             status = get_status()
-            _ensure_non_awaitable(
-                status, callback_name="get_status", operation_name=operation_name
-            )
             failures = 0
         except Exception as exc:
             failures += 1
@@ -227,19 +230,21 @@ def retry_operation(
     while True:
         try:
             operation_result = operation()
-            _ensure_non_awaitable(
-                operation_result,
-                callback_name="operation",
-                operation_name=operation_name,
-            )
-            return operation_result
         except Exception as exc:
             failures += 1
             if failures >= max_attempts:
                 raise HyperbrowserError(
                     f"{operation_name} failed after {max_attempts} attempts: {exc}"
                 ) from exc
             time.sleep(retry_delay_seconds)
+            continue
+
+        _ensure_non_awaitable(
+            operation_result,
+            callback_name="operation",
+            operation_name=operation_name,
+        )
+        return operation_result
 
 
 async def poll_until_terminal_status_async(
@@ -412,6 +417,8 @@ def collect_paginated_results(
             else:
                 stagnation_failures = 0
             should_sleep = current_page_batch < total_page_batches
+        except _NonRetryablePollingError:
+            raise
         except HyperbrowserPollingError:
             raise
         except Exception as exc:
@@ -490,6 +497,8 @@ async def collect_paginated_results_async(
             else:
                 stagnation_failures = 0
             should_sleep = current_page_batch < total_page_batches
+        except _NonRetryablePollingError:
+            raise
         except HyperbrowserPollingError:
             raise
         except Exception as exc:

tests/test_polling.py

@@ -120,18 +120,26 @@ def test_poll_until_terminal_status_rejects_awaitable_status_callback_result():
     async def async_get_status() -> str:
         return "completed"
 
+    attempts = {"count": 0}
+
+    def get_status() -> object:
+        attempts["count"] += 1
+        return async_get_status()
+
     with pytest.raises(
         HyperbrowserError, match="get_status must return a non-awaitable result"
     ):
         poll_until_terminal_status(
             operation_name="sync poll awaitable callback",
-            get_status=lambda: async_get_status(),  # type: ignore[return-value]
+            get_status=get_status,  # type: ignore[arg-type]
             is_terminal_status=lambda value: value == "completed",
             poll_interval_seconds=0.0001,
             max_wait_seconds=1.0,
-            max_status_failures=1,
+            max_status_failures=5,
         )
 
+    assert attempts["count"] == 1
+
 
 def test_retry_operation_retries_and_returns_value():
     attempts = {"count": 0}
@@ -166,16 +174,24 @@ def test_retry_operation_rejects_awaitable_operation_result():
     async def async_operation() -> str:
         return "ok"
 
+    attempts = {"count": 0}
+
+    def operation() -> object:
+        attempts["count"] += 1
+        return async_operation()
+
     with pytest.raises(
         HyperbrowserError, match="operation must return a non-awaitable result"
     ):
         retry_operation(
             operation_name="sync retry awaitable callback",
-            operation=lambda: async_operation(),  # type: ignore[return-value]
-            max_attempts=2,
+            operation=operation,  # type: ignore[arg-type]
+            max_attempts=5,
             retry_delay_seconds=0.0001,
         )
 
+    assert attempts["count"] == 1
+
 
 def test_async_polling_and_retry_helpers():
     async def run() -> None:
@@ -327,22 +343,36 @@ def test_collect_paginated_results_rejects_awaitable_page_callback_result():
     async def async_get_page() -> dict:
         return {"current": 1, "total": 1, "items": []}
 
+    attempts = {"count": 0}
+
+    def get_next_page(page: int) -> object:
+        attempts["count"] += 1
+        return async_get_page()
+
     with pytest.raises(
         HyperbrowserError, match="get_next_page must return a non-awaitable result"
     ):
         collect_paginated_results(
             operation_name="sync paginated awaitable page callback",
-            get_next_page=lambda page: async_get_page(),  # type: ignore[return-value]
+            get_next_page=get_next_page,  # type: ignore[arg-type]
             get_current_page_batch=lambda response: response["current"],
             get_total_page_batches=lambda response: response["total"],
             on_page_success=lambda response: None,
             max_wait_seconds=1.0,
-            max_attempts=2,
+            max_attempts=5,
             retry_delay_seconds=0.0001,
         )
 
+    assert attempts["count"] == 1
+
 
 def test_collect_paginated_results_rejects_awaitable_on_page_success_result():
+    callback_attempts = {"count": 0}
+
+    def on_page_success(response: dict) -> object:
+        callback_attempts["count"] += 1
+        return asyncio.sleep(0)
+
     with pytest.raises(
         HyperbrowserError, match="on_page_success must return a non-awaitable result"
     ):
@@ -351,12 +381,14 @@ def test_collect_paginated_results_rejects_awaitable_on_page_success_result():
             get_next_page=lambda page: {"current": 1, "total": 1, "items": []},
             get_current_page_batch=lambda response: response["current"],
             get_total_page_batches=lambda response: response["total"],
-            on_page_success=lambda response: asyncio.sleep(0),  # type: ignore[return-value]
+            on_page_success=on_page_success,  # type: ignore[arg-type]
             max_wait_seconds=1.0,
-            max_attempts=2,
+            max_attempts=5,
             retry_delay_seconds=0.0001,
         )
 
+    assert callback_attempts["count"] == 1
+
 
 def test_collect_paginated_results_allows_single_page_on_zero_max_wait():
     collected = []
@@ -401,25 +433,38 @@ async def run() -> None:
 
 def test_collect_paginated_results_async_rejects_non_awaitable_page_callback_result():
     async def run() -> None:
+        attempts = {"count": 0}
+
+        def get_next_page(page: int) -> object:
+            attempts["count"] += 1
+            return {"current": 1, "total": 1, "items": []}
+
         with pytest.raises(
             HyperbrowserError, match="get_next_page must return an awaitable"
         ):
             await collect_paginated_results_async(
                 operation_name="async paginated awaitable validation",
-                get_next_page=lambda page: {"current": 1, "total": 1, "items": []},  # type: ignore[return-value]
+                get_next_page=get_next_page,  # type: ignore[arg-type]
                 get_current_page_batch=lambda response: response["current"],
                 get_total_page_batches=lambda response: response["total"],
                 on_page_success=lambda response: None,
                 max_wait_seconds=1.0,
-                max_attempts=2,
+                max_attempts=5,
                 retry_delay_seconds=0.0001,
             )
+        assert attempts["count"] == 1
 
     asyncio.run(run())
 
 
 def test_collect_paginated_results_async_rejects_awaitable_on_page_success_result():
     async def run() -> None:
+        callback_attempts = {"count": 0}
+
+        def on_page_success(response: dict) -> object:
+            callback_attempts["count"] += 1
+            return asyncio.sleep(0)
+
         with pytest.raises(
             HyperbrowserError,
             match="on_page_success must return a non-awaitable result",
@@ -431,11 +476,12 @@ async def run() -> None:
                 ),
                 get_current_page_batch=lambda response: response["current"],
                 get_total_page_batches=lambda response: response["total"],
-                on_page_success=lambda response: asyncio.sleep(0),  # type: ignore[return-value]
+                on_page_success=on_page_success,  # type: ignore[arg-type]
                 max_wait_seconds=1.0,
-                max_attempts=2,
+                max_attempts=5,
                 retry_delay_seconds=0.0001,
             )
+        assert callback_attempts["count"] == 1
 
     asyncio.run(run())