diff --git a/api/routers/generation.py b/api/routers/generation.py index 7ffcf239..936eb2c9 100644 --- a/api/routers/generation.py +++ b/api/routers/generation.py @@ -9,7 +9,11 @@ from services.generators.base import smooth_progress, GenerationCancelled import re as _re -from services.generator_registry import generator_registry, WORKSPACE_DIR +# Import the module (not the name) so WORKSPACE_DIR is read at call time: the +# settings endpoint rebinds it when the user relocates the workspace, and a +# binding captured at import would keep writing output to the old directory. +import services.generator_registry as registry +from services.generator_registry import generator_registry from schemas.generation import JobStatus router = APIRouter(tags=["generation"]) @@ -164,7 +168,7 @@ def progress_cb(pct: int, step: str = "") -> None: return # Direct output to the collection subfolder - coll_dir = WORKSPACE_DIR / collection + coll_dir = registry.WORKSPACE_DIR / collection coll_dir.mkdir(parents=True, exist_ok=True) gen.outputs_dir = coll_dir @@ -185,7 +189,7 @@ def progress_cb(pct: int, step: str = "") -> None: job.progress = 100 _completed_at[job_id] = time.monotonic() try: - rel = output_path.relative_to(WORKSPACE_DIR) + rel = output_path.relative_to(registry.WORKSPACE_DIR) job.output_url = f"/workspace/{rel.as_posix()}" except ValueError: job.output_url = f"/workspace/{collection}/{output_path.name}" diff --git a/api/tests/test_generation_router.py b/api/tests/test_generation_router.py new file mode 100644 index 00000000..e6f77617 --- /dev/null +++ b/api/tests/test_generation_router.py @@ -0,0 +1,86 @@ +import asyncio +import tempfile +import threading +import unittest +from pathlib import Path + +import services.generator_registry as registry +import routers.generation as generation +from schemas.generation import JobStatus + + +class _FakeGenerator: + """Writes its output into whatever directory generation assigns it.""" + + def __init__(self) -> None: + self.outputs_dir: Path | None = None + + def generate(self, image_bytes, params, progress_cb, cancel_event=None) -> Path: + out = Path(self.outputs_dir) / "model.glb" + out.write_bytes(b"glb") + return out + + +class _FakeRegistry: + def __init__(self, gen: _FakeGenerator) -> None: + self._gen = gen + + def active_status(self) -> dict: + # Report loaded so _run_generation skips the download/load thread. + return {"loaded": True, "name": "fake", "downloaded": True} + + def get_active(self) -> _FakeGenerator: + return self._gen + + +class RunGenerationWorkspaceTests(unittest.TestCase): + """A generation started after the workspace path is relocated at runtime + (POST /settings/paths) must file its output under the *current* workspace, + not the one captured when the module was imported.""" + + def setUp(self) -> None: + self._prev_registry = generation.generator_registry + self._prev_ws = registry.WORKSPACE_DIR + self._tmp = tempfile.TemporaryDirectory() + # The user relocated the workspace: the registry global now points here. + registry.WORKSPACE_DIR = Path(self._tmp.name) / "new_workspace" + # Keep the test hermetic against the module's import-time binding: if the + # stale name still exists (before the fix) redirect it into the temp tree + # so the assertion — not a stray write to the real workspace — is what + # catches the bug. + self._had_stale = hasattr(generation, "WORKSPACE_DIR") + if self._had_stale: + generation.WORKSPACE_DIR = Path(self._tmp.name) / "old_workspace" + + def tearDown(self) -> None: + generation.generator_registry = self._prev_registry + registry.WORKSPACE_DIR = self._prev_ws + if self._had_stale: + generation.WORKSPACE_DIR = self._prev_ws + for store in ( + generation._jobs, + generation._cancel_events, + generation._cancelled, + generation._completed_at, + ): + store.clear() + self._tmp.cleanup() + + def _run(self, collection: str) -> tuple[_FakeGenerator, JobStatus]: + gen = _FakeGenerator() + generation.generator_registry = _FakeRegistry(gen) + job_id = "job-test" + generation._jobs[job_id] = JobStatus(job_id=job_id, status="pending", progress=0) + generation._cancel_events[job_id] = threading.Event() + asyncio.run(generation._run_generation(job_id, b"img", {}, collection)) + return gen, generation._jobs[job_id] + + def test_output_lands_under_the_current_workspace(self) -> None: + gen, job = self._run("MyColl") + self.assertEqual(Path(gen.outputs_dir), registry.WORKSPACE_DIR / "MyColl") + self.assertEqual(job.status, "done") + self.assertEqual(job.output_url, "/workspace/MyColl/model.glb") + + +if __name__ == "__main__": + unittest.main()