Tests: better automation and filtering for cloud tests

Author: Acly

tests/conftest.py

@@ -1,9 +1,14 @@
+import asyncio
+import aiohttp
 import sys
+import psutil
 import pytest
+import os
 import shutil
 import subprocess
 import dotenv
 from pathlib import Path
+from typing import Any
 from PyQt5.QtCore import QCoreApplication
 
 sys.path.append(str(Path(__file__).parent.parent))
@@ -15,7 +20,7 @@
 
 def pytest_addoption(parser):
     parser.addoption("--test-install", action="store_true")
-    parser.addoption("--pod-process", action="store_true")
+    parser.addoption("--cloud", action="store_true")
     parser.addoption("--ci", action="store_true")
     parser.addoption("--benchmark", action="store_true")
 
@@ -95,3 +100,110 @@ def local_download_server():
 
 if has_local_cloud:
     dotenv.load_dotenv(root_dir / "service" / "web" / ".env.local")
+
+
+class CloudService:
+    def __init__(self, loop: QtTestApp, enabled=True):
+        self.loop = loop
+        self.dir = root_dir / "service"
+        self.log_dir = result_dir / "logs"
+        self.log_dir.mkdir(exist_ok=True)
+        self.url = os.environ["TEST_SERVICE_URL"]
+        self.coord_proc: asyncio.subprocess.Process | None = None
+        self.coord_log = None
+        self.worker_proc: asyncio.subprocess.Process | None = None
+        self.worker_task: asyncio.Task | None = None
+        self.worker_log = None
+        self.enabled = enabled
+
+    async def serve(self, process: asyncio.subprocess.Process, log_file):
+        try:
+            async for line in util.ensure(process.stdout):
+                print(line.decode("utf-8"), end="", file=log_file, flush=True)
+        except asyncio.CancelledError:
+            pass
+
+    async def launch_coordinator(self):
+        assert self.coord_proc is None, "Coordinator already running"
+        self.coord_log = open(self.log_dir / "api.log", "w", encoding="utf-8")
+        npm = shutil.which("npm")
+        assert npm is not None, "npm not found in PATH"
+        args = [npm, "run", "dev"]
+        self.coord_proc = await asyncio.create_subprocess_exec(
+            *args,
+            cwd=self.dir / "api",
+            stdout=self.coord_log,
+            stderr=asyncio.subprocess.STDOUT,
+        )
+
+    async def launch_worker(self):
+        assert self.worker_proc is None, "Worker already running"
+        self.worker_log = open(self.log_dir / "worker.log", "w", encoding="utf-8")
+        workerpy = str(self.dir / "pod" / "worker.py")
+        config = str(self.dir / "pod" / "_var" / "worker.json")
+        args = ["-u", "-Xutf8", workerpy, config]
+        self.worker_proc = await asyncio.create_subprocess_exec(
+            sys.executable,
+            *args,
+            cwd=self.dir / "pod",
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+        )
+        assert self.worker_proc.stdout is not None
+        async for line in self.worker_proc.stdout:
+            text = line.decode("utf-8")
+            print(text[:80], end="", file=self.worker_log, flush=True)
+            if "Uvicorn running" in text:
+                break
+
+        self.worker_task = asyncio.create_task(self.serve(self.worker_proc, self.worker_log))
+
+    async def start(self):
+        if not self.enabled or not has_local_cloud:
+            return
+        try:
+            await self.launch_coordinator()
+            await self.launch_worker()
+        except Exception as e:
+            await self.stop()
+            raise e
+
+    async def stop(self):
+        if self.worker_task:
+            self.worker_task.cancel()
+            await self.worker_task
+        if self.worker_proc:
+            self.worker_proc.terminate()
+            await self.worker_proc.wait()
+        if self.coord_proc:
+            children = psutil.Process(self.coord_proc.pid).children(recursive=True)
+            for child in children:
+                child.terminate()
+            self.coord_proc.terminate()
+            await self.coord_proc.wait()
+
+    async def create_user(self, username: str) -> dict[str, Any]:
+        assert self.enabled, "Cloud service is not enabled"
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                f"{self.url}/admin/user/create",
+                json={"name": username},
+            ) as response:
+                response.raise_for_status()
+                result = await response.json()
+                if "error" in result:
+                    raise Exception(result["error"])
+                return result
+
+    def __enter__(self):
+        self.loop.run(self.start())
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        self.loop.run(self.stop())
+
+
+@pytest.fixture(scope="session")
+def cloud_service(qtapp, pytestconfig):
+    with CloudService(qtapp, pytestconfig.getoption("--cloud")) as service:
+        yield service

tests/test_image_transfer.py

@@ -1,11 +1,11 @@
 from base64 import b64decode
 from PIL import Image
 from datetime import datetime
-import os
 import pytest
 
 from ai_diffusion.image import ImageCollection, Image as ImageWrapper
 from ai_diffusion.cloud_client import CloudClient
+from tests.conftest import CloudService
 from .config import root_dir, test_dir
 
 if (root_dir / "service" / "pod" / "lib").exists():
@@ -56,16 +56,18 @@ async def main():
         qtapp.run(main())
 
     @pytest.mark.parametrize("mode", ["b64", "transfer"])
-    def test_receive(qtapp, mode: str):
+    def test_receive(qtapp, cloud_service: CloudService, mode: str):
+        if not cloud_service.enabled:
+            pytest.skip("Cloud service not running")
+
         max_b64_size = max_b64_size_config[mode]
         images = [ImageWrapper.load(test_dir / "images" / f) for f in ("cat.webp", "pegonia.webp")]
         bytes, offsets = ImageCollection(images).to_bytes()
         input = {"image_data": {"bytes": bytes, "offsets": offsets}}
 
         async def main():
-            url = os.environ["TEST_SERVICE_URL"]
-            token = os.environ.get("TEST_SERVICE_TOKEN", "")
-            client = await CloudClient.connect(url, token)
+            user = await cloud_service.create_user("image-transfer-test")
+            client = await CloudClient.connect(cloud_service.url, user["token"])
             await client.send_images(input, max_inline_size=max_b64_size)
 
             if mode == "transfer":

tests/test_service.py

@@ -1,10 +1,5 @@
 from pathlib import Path
 import pytest
-import subprocess
-import os
-import sys
-import asyncio
-import dotenv
 
 from ai_diffusion.api import WorkflowInput, WorkflowKind, ControlInput, ImageInput, CheckpointInput
 from ai_diffusion.api import SamplingInput, ConditioningInput, ExtentInput, RegionInput
@@ -13,52 +8,8 @@
 from ai_diffusion.image import Extent, Image, Bounds
 from ai_diffusion.resources import ControlMode, Arch
 from ai_diffusion.util import ensure
-from .conftest import has_local_cloud
-from .config import root_dir, test_dir, result_dir
-
-pod_main = root_dir / "service" / "pod" / "pod.py"
-run_dir = test_dir / "pod"
-
-
-@pytest.fixture(scope="module")
-def pod_server(qtapp, pytestconfig):
-    async def serve(process: asyncio.subprocess.Process):
-        try:
-            async for line in ensure(process.stdout):
-                print(line.decode("utf-8"), end="")
-        except asyncio.CancelledError:
-            process.terminate()
-            await process.wait()
-
-    async def start():
-        env = os.environ.copy()
-        args = ["-u", "-Xutf8", str(pod_main), "--rp_serve_api"]
-        process = await asyncio.create_subprocess_exec(
-            sys.executable,
-            *args,
-            env=env,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.STDOUT,
-        )
-        async for line in ensure(process.stdout):
-            text = line.decode("utf-8")
-            print(text[:80], end="")
-            if "Uvicorn running" in text:
-                break
-
-        return process, asyncio.create_task(serve(process))
-
-    async def stop(process, task):
-        process.terminate()
-        task.cancel()
-        await process.communicate()
-
-    if not pytestconfig.getoption("--pod-process") or pytestconfig.getoption("--ci"):
-        yield None  # For using local docker image or deployed serverless endpoint
-    else:
-        process, task = qtapp.run(start())
-        yield process
-        qtapp.run(stop(process, task))
+from .conftest import CloudService
+from .config import test_dir, result_dir
 
 
 async def receive_images(client: Client, work: WorkflowInput):
@@ -74,16 +25,18 @@ async def receive_images(client: Client, work: WorkflowInput):
     assert False, "Connection closed without receiving images"
 
 
+async def connect_cloud(service: CloudService):
+    user = await service.create_user("workflow-tester")
+    return await CloudClient.connect(service.url, user["token"])
+
+
 @pytest.fixture()
-def cloud_client(pytestconfig, qtapp, pod_server):
+def cloud_client(pytestconfig, qtapp, cloud_service: CloudService):
     if pytestconfig.getoption("--ci"):
         pytest.skip("Diffusion is disabled on CI")
-    if not has_local_cloud:
-        pytest.skip("Local cloud service not found")
-    dotenv.load_dotenv(root_dir / "service" / "web" / ".env.local")
-    url = os.environ["TEST_SERVICE_URL"]
-    token = os.environ["TEST_SERVICE_TOKEN"]
-    return qtapp.run(CloudClient.connect(url, token))
+    if not cloud_service.enabled:
+        pytest.skip("Cloud service not running")
+    return qtapp.run(connect_cloud(cloud_service))
 
 
 def run_and_save(

tests/test_updates.py

@@ -6,7 +6,7 @@
 
 from ai_diffusion.platform_tools import ZipFile
 from ai_diffusion.updates import AutoUpdate, UpdateState
-from .conftest import has_local_cloud
+from .conftest import CloudService
 
 
 class SignalObserver:
@@ -27,15 +27,14 @@ def http_session(service_url: str):
     return ClientSession(service_url, headers=headers)
 
 
-def test_auto_update(qtapp, tmp_path: Path):
-    if not has_local_cloud:
-        pytest.skip("No local cloud service found")
-    qtapp.run(run_auto_update_test(tmp_path))
+def test_auto_update(qtapp, cloud_service: CloudService, tmp_path: Path):
+    if not cloud_service.enabled:
+        pytest.skip("Cloud service not running")
+    qtapp.run(run_auto_update_test(cloud_service, tmp_path))
 
 
-async def run_auto_update_test(tmp_path: Path):
-    service_url = os.environ["TEST_SERVICE_URL"]
-    async with http_session(service_url) as session:
+async def run_auto_update_test(service: CloudService, tmp_path: Path):
+    async with http_session(service.url) as session:
         last_version = new_version = "666.6.6"
 
         # Get the latest plugin version (set from previous test)
@@ -56,7 +55,7 @@ async def run_auto_update_test(tmp_path: Path):
         updater = AutoUpdate(
             current_version=last_version,
             plugin_dir=install_dir,
-            api_url=service_url,
+            api_url=service.url,
         )
         assert updater.state is UpdateState.unknown
 
@@ -105,11 +104,10 @@ async def run_auto_update_test(tmp_path: Path):
         assert install_test_file.read_text() == "if you're feeling orange, try flying a kite"
 
 
-async def test_authorization():
-    if not has_local_cloud:
-        pytest.skip("No local cloud service found")
-    service_url = os.environ["TEST_SERVICE_URL"]
-    async with ClientSession(service_url) as session:
+async def test_authorization(cloud_service: CloudService):
+    if not cloud_service.enabled:
+        pytest.skip("Cloud service not running")
+    async with ClientSession(cloud_service.url) as session:
         # Version check is public
         async with session.get("/plugin/latest?version=1.2.3") as response:
             assert response.status == 200

tests/test_workflow.py

@@ -1,8 +1,6 @@
 import itertools
 import pytest
-import dotenv
 import json
-import os
 from datetime import datetime
 from pathlib import Path
 from typing import Any
@@ -25,6 +23,7 @@
 from ai_diffusion.workflow import detect_inpaint
 from . import config
 from .config import root_dir, test_dir, image_dir, result_dir, reference_dir, default_checkpoint
+from .conftest import CloudService
 
 service_available = (root_dir / "service" / "web" / ".env.local").exists()
 client_params = ["local", "cloud"] if service_available else ["local"]
@@ -38,29 +37,22 @@ async def connect_local():
     return client
 
 
-async def connect_cloud():
-    dotenv.load_dotenv(root_dir / "service" / "web" / ".env.local")
-    url = os.environ["TEST_SERVICE_URL"]
-    token = os.environ.get("TEST_SERVICE_TOKEN", "")
-    if not token:
-        client = CloudClient(url)
-        sign_in = client.sign_in()
-        auth_url = await anext(sign_in)
-        print("\nSign-in required:", auth_url)
-        token = await anext(sign_in)
-        print("\nToken received:", token, "\n")
-    return await CloudClient.connect(url, token)
+async def connect_cloud(service: CloudService):
+    user = await service.create_user("workflow-tester")
+    return await CloudClient.connect(service.url, user["token"])
 
 
 @pytest.fixture(params=client_params)
-def client(pytestconfig, request, qtapp):
+def client(pytestconfig, request, qtapp, cloud_service: CloudService):
     if pytestconfig.getoption("--ci"):
         pytest.skip("Diffusion is disabled on CI")
 
     if request.param == "local":
         client = qtapp.run(connect_local())
     else:
-        client = qtapp.run(connect_cloud())
+        if not cloud_service.enabled:
+            pytest.skip("Cloud service not running")
+        client = qtapp.run(connect_cloud(cloud_service))
     files.loras.update([File.remote(m) for m in client.models.loras], FileSource.remote)
 
     yield client