fix: monotonic wait_for_task and raise errors in get_task_status

Author: Thordata

src/thordata/async_client.py

@@ -564,6 +564,11 @@ async def create_scraper_task_advanced(self, config: ScraperTaskConfig) -> str:
     async def get_task_status(self, task_id: str) -> str:
         """
         Check async task status.
+
+        Raises:
+            ThordataConfigError: If public credentials are missing.
+            ThordataAPIError: If API returns a non-200 code in JSON payload.
+            ThordataNetworkError: If network/HTTP request fails.
         """
         self._require_public_credentials()
         session = self._get_session()
@@ -577,17 +582,50 @@ async def get_task_status(self, task_id: str) -> str:
             async with session.post(
                 self._status_url, data=payload, headers=headers
             ) as response:
+                response.raise_for_status()
                 data = await response.json()
 
-                if data.get("code") == 200 and data.get("data"):
-                    for item in data["data"]:
+                if isinstance(data, dict):
+                    code = data.get("code")
+                    if code is not None and code != 200:
+                        msg = extract_error_message(data)
+                        raise_for_code(
+                            f"Task status API Error: {msg}",
+                            code=code,
+                            payload=data,
+                        )
+
+                    items = data.get("data") or []
+                    for item in items:
                         if str(item.get("task_id")) == str(task_id):
                             return item.get("status", "unknown")
 
-                return "unknown"
+                    return "unknown"
+
+                raise ThordataNetworkError(
+                    f"Unexpected task status response type: {type(data).__name__}",
+                    original_error=None,
+                )
 
-        except Exception as e:
-            logger.error(f"Async status check failed: {e}")
+        except asyncio.TimeoutError as e:
+            raise ThordataTimeoutError(
+                f"Async status check timed out: {e}", original_error=e
+            )
+        except aiohttp.ClientError as e:
+            raise ThordataNetworkError(
+                f"Async status check failed: {e}", original_error=e
+            )
+
+    async def safe_get_task_status(self, task_id: str) -> str:
+        """
+        Backward-compatible status check.
+
+        Returns:
+            Status string, or "error" on any exception.
+        """
+        try:
+            return await self.get_task_status(task_id)
+        except Exception:
             return "error"
 
     async def get_task_result(self, task_id: str, file_type: str = "json") -> str:
@@ -632,9 +670,12 @@ async def wait_for_task(
         """
         Wait for a task to complete.
         """
-        elapsed = 0.0
 
-        while elapsed < max_wait:
+        import time
+
+        start = time.monotonic()
+
+        while (time.monotonic() - start) < max_wait:
             status = await self.get_task_status(task_id)
 
             logger.debug(f"Task {task_id} status: {status}")
@@ -652,7 +693,6 @@ async def wait_for_task(
                 return status
 
             await asyncio.sleep(poll_interval)
-            elapsed += poll_interval
 
         raise TimeoutError(f"Task {task_id} did not complete within {max_wait} seconds")
 

src/thordata/client.py

@@ -697,11 +697,13 @@ def get_task_status(self, task_id: str) -> str:
         """
         Check the status of an asynchronous scraping task.
 
-        Args:
-            task_id: The task ID from create_scraper_task.
-
         Returns:
             Status string (e.g., "running", "ready", "failed").
+
+        Raises:
+            ThordataConfigError: If public credentials are missing.
+            ThordataAPIError: If API returns a non-200 code in JSON payload.
+            ThordataNetworkError: If network/HTTP request fails.
         """
         self._require_public_credentials()
 
@@ -718,18 +720,46 @@ def get_task_status(self, task_id: str) -> str:
                 timeout=30,
             )
             response.raise_for_status()
-
             data = response.json()
 
-            if data.get("code") == 200 and data.get("data"):
-                for item in data["data"]:
+            if isinstance(data, dict):
+                code = data.get("code")
+                if code is not None and code != 200:
+                    msg = extract_error_message(data)
+                    raise_for_code(
+                        f"Task status API Error: {msg}",
+                        code=code,
+                        payload=data,
+                    )
+
+                items = data.get("data") or []
+                for item in items:
                     if str(item.get("task_id")) == str(task_id):
                         return item.get("status", "unknown")
 
-            return "unknown"
+                return "unknown"
+
+            # Unexpected payload type
+            raise ThordataNetworkError(
+                f"Unexpected task status response type: {type(data).__name__}",
+                original_error=None,
+            )
+
+        except requests.Timeout as e:
+            raise ThordataTimeoutError(f"Status check timed out: {e}", original_error=e)
+        except requests.RequestException as e:
+            raise ThordataNetworkError(f"Status check failed: {e}", original_error=e)
+
+    def safe_get_task_status(self, task_id: str) -> str:
+        """
+        Backward-compatible status check.
 
-        except Exception as e:
-            logger.error(f"Status check failed: {e}")
+        Returns:
+            Status string, or "error" on any exception.
+        """
+        try:
+            return self.get_task_status(task_id)
+        except Exception:
             return "error"
 
     def get_task_result(self, task_id: str, file_type: str = "json") -> str:
@@ -797,9 +827,9 @@ def wait_for_task(
         """
         import time
 
-        elapsed = 0.0
+        start = time.monotonic()
 
-        while elapsed < max_wait:
+        while (time.monotonic() - start) < max_wait:
             status = self.get_task_status(task_id)
 
             logger.debug(f"Task {task_id} status: {status}")
@@ -817,7 +847,6 @@ def wait_for_task(
                 return status
 
             time.sleep(poll_interval)
-            elapsed += poll_interval
 
         raise TimeoutError(f"Task {task_id} did not complete within {max_wait} seconds")
 

tests/test_task_status_and_wait.py

@@ -0,0 +1,71 @@
+import pytest
+from pytest_httpserver import HTTPServer
+
+from thordata import (
+    AsyncThordataClient,
+    ThordataAuthError,
+    ThordataClient,
+)
+
+
+def test_wait_for_task_timeout_uses_monotonic(monkeypatch) -> None:
+    client = ThordataClient(scraper_token="dummy", public_token="p", public_key="k")
+
+    # Always "running" so it must time out quickly.
+    monkeypatch.setattr(client, "get_task_status", lambda task_id: "running")
+
+    with pytest.raises(TimeoutError):
+        client.wait_for_task("t1", poll_interval=0.01, max_wait=0.05)
+
+
+@pytest.mark.asyncio
+async def test_async_wait_for_task_timeout_uses_monotonic(monkeypatch) -> None:
+    async with AsyncThordataClient(
+        scraper_token="dummy", public_token="p", public_key="k"
+    ) as client:
+
+        async def _always_running(task_id: str) -> str:
+            return "running"
+
+        monkeypatch.setattr(client, "get_task_status", _always_running)
+
+        with pytest.raises(TimeoutError):
+            await client.wait_for_task("t1", poll_interval=0.01, max_wait=0.05)
+
+
+def test_get_task_status_raises_on_non_200_code(httpserver: HTTPServer) -> None:
+    httpserver.expect_request("/tasks-status", method="POST").respond_with_json(
+        {"code": 401, "msg": "Unauthorized"}
+    )
+
+    base_url = httpserver.url_for("/").rstrip("/").replace("localhost", "127.0.0.1")
+
+    client = ThordataClient(
+        scraper_token="dummy",
+        public_token="p",
+        public_key="k",
+        web_scraper_api_base_url=base_url,
+    )
+
+    with pytest.raises(ThordataAuthError):
+        client.get_task_status("t1")
+
+
+@pytest.mark.asyncio
+async def test_async_get_task_status_raises_on_non_200_code(
+    httpserver: HTTPServer,
+) -> None:
+    httpserver.expect_request("/tasks-status", method="POST").respond_with_json(
+        {"code": 401, "msg": "Unauthorized"}
+    )
+
+    base_url = httpserver.url_for("/").rstrip("/").replace("localhost", "127.0.0.1")
+
+    async with AsyncThordataClient(
+        scraper_token="dummy",
+        public_token="p",
+        public_key="k",
+        web_scraper_api_base_url=base_url,
+    ) as client:
+        with pytest.raises(ThordataAuthError):
+            await client.get_task_status("t1")