Spaces:
Paused
Paused
| import json | |
| import os | |
| from pathlib import Path | |
| from tempfile import TemporaryDirectory | |
| from unittest import TestCase | |
| from unittest.mock import patch | |
| from urllib.parse import quote | |
| from catbox.browser_ui import BrowserUiApp, make_default_browser_ui_server | |
| class RecordingBackend: | |
| def __init__(self, image_ref: str, ready: bool = True) -> None: | |
| self.image_ref = image_ref | |
| self.ready = ready | |
| self.observations = 0 | |
| def readiness(self): | |
| if self.ready: | |
| return {"status": "ready", "modelBackend": "ready"} | |
| return {"status": "starting", "modelBackend": "starting"} | |
| def observe(self, trace_callback=None): | |
| self.observations += 1 | |
| if trace_callback is not None: | |
| trace_callback(self.image_ref) | |
| return { | |
| "status": "generated", | |
| "outcome": "living", | |
| "imageRef": self.image_ref, | |
| "traceRefs": [self.image_ref] if trace_callback is not None else [], | |
| "metadata": { | |
| "seed": 41100, | |
| "ephemeral": True, | |
| "traceFrameCount": 1 if trace_callback is not None else 0, | |
| }, | |
| "revealNote": "A local diffusion model generated this outcome.", | |
| } | |
| class FailingBackend: | |
| def __init__(self) -> None: | |
| self.observations = 0 | |
| def readiness(self): | |
| return {"status": "ready", "modelBackend": "ready"} | |
| def observe(self, trace_callback=None): | |
| self.observations += 1 | |
| return { | |
| "status": "generation_failed", | |
| "error": { | |
| "type": "RuntimeError", | |
| "message": "CUDA ran out of memory", | |
| }, | |
| "metadata": { | |
| "seed": 41100, | |
| "outcome": "living", | |
| "ephemeral": True, | |
| }, | |
| } | |
| class BrowserUiServerTests(TestCase): | |
| def test_browser_page_renders_observation_flow_without_outcome_choice(self): | |
| backend = RecordingBackend("/tmp/generated.png") | |
| app = BrowserUiApp(backend=backend) | |
| response = app.handle("GET", "/") | |
| html = response.body.decode("utf-8") | |
| self.assertEqual(response.status, 200) | |
| self.assertEqual(response.headers["Content-Type"], "text/html; charset=utf-8") | |
| self.assertIn("data-state=\"sealed\"", html) | |
| self.assertIn("data-state=\"starting\"", html) | |
| self.assertIn("Preparing the local diffusion chamber", html) | |
| self.assertIn("id=\"observe-button\"", html) | |
| self.assertIn("Observe", html) | |
| self.assertIn("data-state=\"waiting\"", html) | |
| self.assertIn("Captured Denoising Trace", html) | |
| self.assertIn("id=\"trace-frame\"", html) | |
| self.assertIn("id=\"trace-placeholder\"", html) | |
| self.assertIn("TRACE_POLL_MS", html) | |
| self.assertIn("pollTrace", html) | |
| self.assertIn("/api/trace", html) | |
| self.assertIn("id=\"progressive-waiting-status\"", html) | |
| self.assertIn("hidden", html) | |
| self.assertIn("PROGRESSIVE_WAITING_DELAY_MS", html) | |
| self.assertIn("showProgressiveWaiting", html) | |
| self.assertIn("id=\"generated-outcome\"", html) | |
| self.assertIn("data-loading=\"true\"", html) | |
| self.assertIn("generatedOutcome.onload", html) | |
| self.assertIn("generatedOutcome.onerror", html) | |
| self.assertIn("id=\"reveal-note\"", html) | |
| self.assertIn("id=\"reset-button\"", html) | |
| self.assertIn("data-state=\"generation-failure\"", html) | |
| self.assertIn("id=\"generation-failure-message\"", html) | |
| self.assertIn("id=\"retry-button\"", html) | |
| self.assertIn("id=\"failure-reset-button\"", html) | |
| self.assertIn("retryButton.addEventListener(\"click\", observe)", html) | |
| self.assertIn("failureResetButton.addEventListener(\"click\", resetToSealed)", html) | |
| self.assertIn("generationFailureMessage.textContent = \"\"", html) | |
| self.assertIn("transform: scale(0.97)", html) | |
| self.assertIn("@media (prefers-reduced-motion: reduce)", html) | |
| self.assertIn("@media (hover: hover) and (pointer: fine)", html) | |
| self.assertIn("overflow-y: auto", html) | |
| self.assertIn("grid-template-rows: auto minmax(0, 1fr) auto", html) | |
| self.assertIn(".panel {\n display: none", html) | |
| self.assertIn(".panel.is-active {\n display: grid", html) | |
| self.assertIn("bringActivePanelIntoView(revealedPanel)", html) | |
| self.assertIn(".brand {\n position: relative", html) | |
| self.assertIn(".hud {\n position: relative", html) | |
| self.assertNotIn("name=\"outcome\"", html) | |
| def test_readiness_endpoint_returns_backend_startup_state(self): | |
| backend = RecordingBackend("/tmp/generated.png", ready=False) | |
| app = BrowserUiApp(backend=backend) | |
| response = app.handle("GET", "/api/readiness") | |
| payload = json.loads(response.body.decode("utf-8")) | |
| self.assertEqual(response.status, 200) | |
| self.assertEqual(response.headers["Content-Type"], "application/json") | |
| self.assertEqual(payload, {"status": "starting", "modelBackend": "starting"}) | |
| def test_normal_observation_returns_backend_contract_and_serves_generated_file(self): | |
| with TemporaryDirectory() as runtime_dir: | |
| image_path = Path(runtime_dir) / "generated.png" | |
| image_path.write_bytes(b"generated image") | |
| backend = RecordingBackend(str(image_path)) | |
| app = BrowserUiApp(backend=backend) | |
| observe_response = app.handle("POST", "/api/observe") | |
| payload = json.loads(observe_response.body.decode("utf-8")) | |
| self.assertEqual(observe_response.status, 200) | |
| self.assertEqual(observe_response.headers["Content-Type"], "application/json") | |
| self.assertEqual(backend.observations, 1) | |
| self.assertEqual(payload["status"], "generated") | |
| self.assertEqual(payload["outcome"], "living") | |
| self.assertEqual(payload["imageRef"], str(image_path)) | |
| self.assertEqual(payload["traceRefs"], [str(image_path)]) | |
| self.assertEqual(payload["metadata"]["seed"], 41100) | |
| self.assertEqual(payload["metadata"]["traceFrameCount"], 1) | |
| self.assertEqual(payload["revealNote"], "A local diffusion model generated this outcome.") | |
| trace_response = app.handle("GET", "/api/trace") | |
| trace_payload = json.loads(trace_response.body.decode("utf-8")) | |
| self.assertEqual(trace_response.status, 200) | |
| self.assertEqual(trace_payload, {"traceRefs": [str(image_path)]}) | |
| image_response = app.handle( | |
| "GET", | |
| f"/api/generated-outcome?imageRef={quote(payload['imageRef'])}", | |
| ) | |
| self.assertEqual(image_response.status, 200) | |
| self.assertEqual(image_response.body, b"generated image") | |
| self.assertEqual(image_response.headers["Content-Type"], "image/png") | |
| def test_generation_failure_returns_structured_failure_without_generated_image(self): | |
| backend = FailingBackend() | |
| app = BrowserUiApp(backend=backend) | |
| observe_response = app.handle("POST", "/api/observe") | |
| payload = json.loads(observe_response.body.decode("utf-8")) | |
| self.assertEqual(observe_response.status, 200) | |
| self.assertEqual(observe_response.headers["Content-Type"], "application/json") | |
| self.assertEqual(backend.observations, 1) | |
| self.assertEqual(payload["status"], "generation_failed") | |
| self.assertEqual(payload["error"]["type"], "RuntimeError") | |
| self.assertIn("CUDA", payload["error"]["message"]) | |
| self.assertNotIn("imageRef", payload) | |
| image_response = app.handle( | |
| "GET", | |
| "/api/generated-outcome?imageRef=/tmp/not-generated.png", | |
| ) | |
| self.assertEqual(image_response.status, 404) | |
| def test_default_server_loads_hf_token_from_env_file_before_model_runner(self): | |
| with TemporaryDirectory() as project_dir: | |
| env_file = Path(project_dir) / ".env" | |
| env_file.write_text("HF_TOKEN=hf_test_token\n", encoding="utf-8") | |
| observed_token = [] | |
| class TokenAwareRunner: | |
| def __init__(self, runtime_dir): | |
| observed_token.append(os.environ.get("HF_TOKEN")) | |
| def is_ready(self): | |
| return True | |
| def generate(self, outcome, seed, config=None): | |
| return { | |
| "image_ref": "/tmp/generated.png", | |
| "generation_seconds": 0.1, | |
| "metadata": {}, | |
| } | |
| with patch.dict(os.environ, {}, clear=True): | |
| with patch("catbox.browser_ui.SdTurboImageToImageModelRunner", TokenAwareRunner): | |
| with patch("catbox.browser_ui.CatboxBrowserUiServer") as server_class: | |
| make_default_browser_ui_server( | |
| ("127.0.0.1", 0), | |
| env_file=env_file, | |
| ) | |
| self.assertEqual(server_class.call_count, 1) | |
| self.assertEqual(observed_token, ["hf_test_token"]) | |