Tests: timeout and error cases when using cloud service

Author: Acly

tests/conftest.py

@@ -115,6 +115,8 @@ def __init__(self, loop: QtTestApp, enabled=True):
         self.worker_proc: asyncio.subprocess.Process | None = None
         self.worker_task: asyncio.Task | None = None
         self.worker_log = None
+        self.worker_url = ""
+        self.worker_secret = ""
         self.enabled = has_local_cloud and enabled
 
     async def serve(self, process: asyncio.subprocess.Process, log_file):
@@ -158,11 +160,11 @@ async def launch_worker(self):
         config = self.dir / "pod" / "_var" / "worker.json"
         assert config.exists(), "Worker config not found"
         config_dict = json.loads(config.read_text(encoding="utf-8"))
-        worker_url = config_dict["public_url"]
-        admin_secret = config_dict["admin_secret"]
+        self.worker_url = config_dict["public_url"]
+        self.worker_secret = config_dict["admin_secret"]
         self.worker_log = open(self.log_dir / "worker.log", "w", encoding="utf-8")
-        if await self.check(f"{worker_url}/health", token=admin_secret):
-            print(f"Worker running in external process at {worker_url}", file=self.worker_log)
+        if await self.check(f"{self.worker_url}/health", token=self.worker_secret):
+            print(f"Worker running in external process at {self.worker_url}", file=self.worker_log)
             return
 
         workerpy = str(self.dir / "pod" / "worker.py")
@@ -220,6 +222,17 @@ async def create_user(self, username: str) -> dict[str, Any]:
                     raise Exception(result["error"])
                 return result
 
+    async def set_worker_job_timeout(self, timeout: int):
+        headers = {
+            "Authorization": f"Bearer {self.worker_secret}",
+            "Content-Type": "application/json",
+        }
+        async with aiohttp.ClientSession(headers=headers) as session:
+            async with session.post(
+                f"{self.worker_url}/configure", json={"job_timeout": timeout}
+            ) as response:
+                response.raise_for_status()
+
     def __enter__(self):
         self.loop.run(self.start())
         return self

tests/test_service.py

@@ -1,4 +1,5 @@
 import asyncio
+import aiohttp
 from pathlib import Path
 from timeit import default_timer as timer
 import pytest
@@ -57,16 +58,20 @@ async def runner():
     return results[0]
 
 
-def create_simple_workflow(prompt="fluffy ball", input: Image | None = None):
+def create_simple_workflow(prompt="fluffy ball", input: Image | Extent | None = None):
     start = 0
-    images = ImageInput.from_extent(Extent(512, 512))
-    if input:
+    if isinstance(input, Image):
+        images = ImageInput.from_extent(input.extent)
         images.initial_image = input
         images.hires_image = input
         start = 4
+    elif isinstance(input, Extent):
+        images = ImageInput.from_extent(input)
+    else:
+        images = ImageInput.from_extent(Extent(512, 512))
 
     return WorkflowInput(
-        WorkflowKind.generate if input is None else WorkflowKind.refine,
+        WorkflowKind.generate if images.initial_image is None else WorkflowKind.refine,
         images=images,
         models=CheckpointInput("dreamshaper_8.safetensors"),
         sampling=SamplingInput(
@@ -230,3 +235,38 @@ async def main():
     end_time = timer()
     duration = end_time - start_time
     print(f"Completed 5 x 2 jobs in {duration:.2f} seconds", end=" ")
+
+
+def test_error_workflow(qtapp, cloud_client: CloudClient):
+    workflow = create_simple_workflow()
+    workflow.kind = WorkflowKind.refine  # Error: refine requires an input image
+    with pytest.raises(Exception, match="failed"):
+        run_and_save(qtapp, cloud_client, workflow, "error_workflow")
+
+
+def test_job_timeout(pytestconfig, qtapp, cloud_service: CloudService):
+    if not cloud_service.enabled:
+        pytest.skip("Cloud service not running")
+
+    async def main():
+        user = await cloud_service.create_user("timeout-tester")
+        client = await CloudClient.connect(cloud_service.url, user["token"])
+        big_workflow = create_simple_workflow(input=Extent(2048, 1536))
+
+        await cloud_service.set_worker_job_timeout(5)
+
+        with pytest.raises(Exception, match="timeout"):
+            await receive_images(client, big_workflow)
+
+        # Worker should be restarted and accept new jobs
+        for _attempt in range(5):
+            try:
+                await cloud_service.set_worker_job_timeout(480)
+                break
+            except aiohttp.ClientConnectionError:
+                await asyncio.sleep(2)  # Wait for worker to be back up
+        small_workflow = create_simple_workflow()
+        images = await receive_images(client, small_workflow)
+        assert len(images) == 2
+
+    qtapp.run(main())