Merge pull request #25 from alphabetc1/ci/add_more_testcase

[WIP] ci: add more testcase

Author: zhaochenyang20

.github/workflows/pr-test.yml

@@ -0,0 +1,25 @@
+name: PR Test
+
+on:
+  pull_request:
+    branches: [main]
+
+jobs:
+  pr-test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.12"
+
+      - name: Install package and test deps
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install -e .
+          python -m pip install pytest
+
+      - name: Run CPU-only tests
+        run: pytest tests/unit -v

development.md

@@ -10,7 +10,11 @@ Run tests:
 
 ```bash
 pip install pytest
+# CPU-only tests (unit + fake e2e)
 pytest tests/unit -v
+
+# Real E2E tests (GPU required, longer runtime)
+pytest tests/e2e/test_e2e_sglang.py -v -s
 ```
 
 ## Benchmark Scripts

pyproject.toml

@@ -33,4 +33,4 @@ package-dir = { "" = "src" }
 where = ["src"]
 
 [tool.pytest.ini_options]
-testpaths = ["tests/unit"]
+testpaths = ["tests/unit", "tests/e2e"]

src/sglang_diffusion_routing/cli/main.py

@@ -4,7 +4,6 @@
 from __future__ import annotations
 
 import argparse
-import asyncio
 import sys
 
 from sglang_diffusion_routing import DiffusionRouter
@@ -26,18 +25,9 @@ def _run_router_server(
         worker_urls if worker_urls is not None else args.worker_urls or []
     )
     router = DiffusionRouter(args, verbose=args.verbose)
-    refresh_tasks = []
     for url in worker_urls:
         normalized_url = router.normalize_worker_url(url)
         router.register_worker(normalized_url)
-        refresh_tasks.append(router.refresh_worker_video_support(normalized_url))
-
-    if refresh_tasks:
-
-        async def _refresh_all_worker_video_support() -> None:
-            await asyncio.gather(*refresh_tasks)
-
-        asyncio.run(_refresh_all_worker_video_support())
 
     print(f"{log_prefix} starting router on {args.host}:{args.port}", flush=True)
     print(

src/sglang_diffusion_routing/router/diffusion_router.py

@@ -69,6 +69,16 @@ def _setup_routes(self) -> None:
         )
 
     async def _start_background_health_check(self) -> None:
+        # Probe video capability for pre-registered workers in the running event loop.
+        unknown_workers = [
+            url for url, support in self.worker_video_support.items() if support is None
+        ]
+        if unknown_workers:
+            await asyncio.gather(
+                *(self.refresh_worker_video_support(url) for url in unknown_workers),
+                return_exceptions=True,
+            )
+
         if self._health_task is None or self._health_task.done():
             self._health_task = asyncio.create_task(self._health_check_loop())
 
@@ -383,6 +393,12 @@ async def generate(self, request: Request):
 
     async def generate_video(self, request: Request):
         """Route video generation to /v1/videos."""
+        if not self.worker_request_counts:
+            return JSONResponse(
+                status_code=503,
+                content={"error": "No workers registered in the pool"},
+            )
+
         candidate_workers = [
             worker_url
             for worker_url, support in self.worker_video_support.items()

tests/conftest.py

@@ -0,0 +1,11 @@
+"""Pytest configuration: force local src import precedence."""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+src_str = str(Path(__file__).resolve().parent.parent / "src")
+while src_str in sys.path:
+    sys.path.remove(src_str)
+sys.path.insert(0, src_str)