"""Tests for the SGLang-Omni backend package (backends/sglomni).""" import io import json import shutil import urllib.error import sys import tempfile import unittest from contextlib import redirect_stdout from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock, patch from backends import common, envs, probe, servers from backends.sglomni import catalog, constants, models, pythonenv, status from backends.sglomni.catalog import CAPABILITY_CLONE, CAPABILITY_DESIGN, \ CAPABILITY_SPEAKER, ENTRIES, config_path, entry_by_key, entry_by_repo, \ install_tree_families from backends.sglomni.pythonenv import SGLOMNI_ENV class CatalogTests(unittest.TestCase): """The model catalog is the single source of hosting/voice facts.""" def test_unique_keys_and_repos(self): keys = [entry.key for entry in ENTRIES] repos = [entry.repo for entry in ENTRIES] self.assertEqual(len(keys), len(set(keys))) self.assertEqual(len(repos), len(set(repos))) def test_capabilities_are_known(self): for entry in ENTRIES: self.assertIn(entry.capability, (CAPABILITY_SPEAKER, CAPABILITY_CLONE, CAPABILITY_DESIGN)) def test_vendored_config_files_exist(self): for entry in ENTRIES: path = config_path(entry) if entry.config is None: self.assertIsNone(path) else: self.assertTrue(path.is_file(), f"missing {path}") def test_config_declares_the_entry_repo(self): # The vendored yaml pins model_path — it must match the entry's # repo, or the server would host something else than the run # selected (the client's connect check would refuse it). The # files are flat `key: value` documents, parsed by hand here. for entry in ENTRIES: path = config_path(entry) if path is None: continue data = {} for line in path.read_text(encoding="utf-8").splitlines(): line = line.strip() if not line or line.startswith("#") or ":" not in line: continue key, _, value = line.partition(":") data[key.strip()] = value.strip() self.assertEqual(data.get("model_path"), entry.repo, f"stale config for {entry.key}") def test_clone_capability_matches_reference_requirement(self): # Only clone models carry a reference requirement. for entry in ENTRIES: if entry.requires_reference: self.assertEqual(entry.capability, CAPABILITY_CLONE, entry.key) def test_entry_lookup_by_key_and_repo(self): entry = ENTRIES[0] self.assertIs(entry_by_key(entry.key), entry) self.assertIs(entry_by_repo(entry.repo), entry) self.assertIsNone(entry_by_key("nope")) self.assertIsNone(entry_by_repo("nope")) def test_install_tree_covers_every_entry(self): families = install_tree_families(list(ENTRIES)) covered = [option["key"] for family in families for option in family["options"]] self.assertEqual(sorted(covered), sorted(entry.key for entry in ENTRIES)) def test_install_tree_has_no_detail_line(self): # The install screen's status line under the Confirm/Back buttons # would only repeat the catalog keys under the cursor — the tree # carries no detail at all, so no status line is drawn. families = install_tree_families(list(ENTRIES)) self.assertTrue(families) for family in families: self.assertNotIn("detail", family) def test_install_tree_filters_unavailable(self): some = [ENTRIES[0]] families = install_tree_families(some) covered = [option["key"] for family in families for option in family["options"]] self.assertEqual(covered, [ENTRIES[0].key]) class ModelInstallStateTests(unittest.TestCase): """Install state reads the shared HuggingFace hub cache layout.""" def setUp(self): self._tmp = tempfile.TemporaryDirectory() self.cache = Path(self._tmp.name) patcher = patch("backends.common.hf_cache_dir", return_value=self.cache) patcher.start() self.addCleanup(patcher.stop) self.addCleanup(self._tmp.cleanup) def _seed_repo(self, repo): directory = self.cache / ("models--" + repo.replace("/", "--")) (directory / "refs").mkdir(parents=True) (directory / "refs" / "main").write_text("hash\n") (directory / "snapshots" / "abc").mkdir(parents=True) (directory / "snapshots" / "abc" / "weights.safetensors") \ .write_bytes(b"x") return directory def test_model_installed_needs_refs_and_snapshots(self): entry = entry_by_key("higgs_audio_v3_tts") self.assertFalse(models.model_installed(entry)) self._seed_repo(entry.repo) self.assertTrue(models.model_installed(entry)) def test_installed_entries_in_catalog_order(self): first, second = ENTRIES[0], ENTRIES[4] self._seed_repo(second.repo) self._seed_repo(first.repo) keys = models.installed_keys() self.assertEqual(keys, [first.key, second.key]) def test_delete_model_weights_removes_only_targeted_repos(self): entry = ENTRIES[0] other = ENTRIES[1] self._seed_repo(entry.repo) self._seed_repo(other.repo) removed = models.delete_model_weights([entry]) self.assertEqual(removed, 1) self.assertFalse(models.model_installed(entry)) self.assertTrue(models.model_installed(other)) def test_preset_voices_from_voice_embedding(self): # Voxtral-style: preset voices ship as voice_embedding/*.pt in the # downloaded snapshot. entry = entry_by_key("voxtral_tts") directory = self._seed_repo(entry.repo) (directory / "snapshots" / "abc" / "voice_embedding").mkdir() (directory / "snapshots" / "abc" / "voice_embedding" / "casual_male.pt") \ .write_bytes(b"x") (directory / "snapshots" / "abc" / "voice_embedding" / "default.pt") \ .write_bytes(b"x") self.assertEqual(models.preset_voices(entry), ["casual_male", "default"]) def test_preset_voices_from_catalog_table(self): entry = entry_by_key("qwen3_tts_0_6b_customvoice") self.assertTrue(entry.speakers) self.assertEqual(models.preset_voices(entry), list(entry.speakers)) def test_resolve_model_unknown_key_raises(self): with self.assertRaises(RuntimeError) as ctx: models.resolve_model("nope") self.assertIn("Unknown sglang-omni model", str(ctx.exception)) def test_resolve_model_requires_installed_weights(self): with self.assertRaises(RuntimeError) as ctx: models.resolve_model("higgs_audio_v3_tts") self.assertIn("not downloaded", str(ctx.exception)) def test_resolve_model_auto_selects_the_single_install(self): entry = ENTRIES[0] self._seed_repo(entry.repo) self.assertIs(models.resolve_model(None), entry) self.assertIs(models.resolve_model(entry.key), entry) def test_resolve_model_needs_a_pick_with_several_installs(self): self._seed_repo(ENTRIES[0].repo) self._seed_repo(ENTRIES[1].repo) with self.assertRaises(RuntimeError) as ctx: models.resolve_model(None) self.assertIn("--model", str(ctx.exception)) class CompanionPackageTests(unittest.TestCase): """The venv probe and start-time heal for a model's companion packages. Weights can reach the shared HF cache by another route (another backend's install) or outlive a failed companion pip run — which install_model only warns about — so the probe is what keeps a "installed" model from booting into a ModuleNotFoundError. """ def test_extra_import_names(self): self.assertEqual(catalog.extra_import_name("sox"), "sox") self.assertEqual(catalog.extra_import_name("einops"), "einops") self.assertEqual(catalog.extra_import_name("qwen-tts==0.1.1"), "qwen_tts") self.assertEqual( catalog.extra_import_name("descript-audiotools==0.7.2"), "audiotools") self.assertEqual( catalog.extra_import_name("descript-audio-codec==1.0.0"), "dac") self.assertEqual( catalog.extra_import_name("protobuf==6.33.6"), "google.protobuf.runtime_version") def test_missing_companions_reports_absent_modules(self): entry = entry_by_key("qwen3_tts_1_7b_base") with patch.object(envs, "env_exists", return_value=True), \ patch.object(envs, "module_available", side_effect=lambda name, env_dir: name != "qwen_tts"): missing = models.missing_companions(entry) self.assertEqual([spec for spec, _no_deps in missing], ["qwen-tts==0.1.1"]) def test_missing_companions_skips_every_extra_that_imports(self): entry = entry_by_key("qwen3_tts_1_7b_base") with patch.object(envs, "env_exists", return_value=True), \ patch.object(envs, "module_available", return_value=True): self.assertEqual(models.missing_companions(entry), []) def test_missing_companions_no_verdict_without_a_venv(self): # A missing venv cannot be healed here: the start flow fails on # the missing sgl-omni executable instead. entry = entry_by_key("qwen3_tts_1_7b_base") with patch.object(envs, "env_exists", return_value=False): self.assertEqual(models.missing_companions(entry), []) def test_dac_heal_reinstalls_protobuf_after_the_downgrade(self): # descript-audiotools carries a vestigial protobuf<3.20 pin: its # with-deps install downgrades the protobuf 6.x the sglang-omni # stack itself needs. The restore extra must therefore re-enter # missing_companions whenever that downgrade happened — probed via # google.protobuf.runtime_version, which 3.19.6 does not provide. entry = entry_by_key("zonos2") self.assertIn(("protobuf==6.33.6", False), entry.extras) with patch.object(envs, "env_exists", return_value=True), \ patch.object(envs, "module_available", side_effect=lambda name, env_dir: name != "google.protobuf.runtime_version"): missing = models.missing_companions(entry) self.assertEqual(missing, [("protobuf==6.33.6", False)]) def test_dac_protobuf_restore_installs_after_descript_audiotools(self): # The restore only re-pins protobuf if it pip-installs AFTER the # package whose dependency resolution did the downgrading. specs = [spec for spec, _no_deps in entry_by_key("fish_s2_pro").extras] self.assertLess(specs.index("descript-audiotools==0.7.2"), specs.index("protobuf==6.33.6")) def test_install_companions_installs_only_missing_with_no_deps(self): entry = entry_by_key("qwen3_tts_1_7b_base") calls = [] def fake_pip_install(specs, **kwargs): calls.append((list(specs), kwargs.get("extra_args"))) return 0 with patch.object(envs, "env_exists", return_value=True), \ patch.object(envs, "module_available", side_effect=lambda name, env_dir: name not in ("sox", "qwen_tts")), \ patch.object(common, "pip_install", side_effect=fake_pip_install): self.assertEqual(models.install_companions(entry), 0) self.assertEqual([specs for specs, _args in calls], [["sox"], ["qwen-tts==0.1.1"]]) self.assertTrue(all(args == ["--no-deps"] for _specs, args in calls)) def test_install_companions_stops_at_the_first_failure(self): entry = entry_by_key("qwen3_tts_1_7b_base") with patch.object(envs, "env_exists", return_value=True), \ patch.object(envs, "module_available", return_value=False), \ patch.object(common, "pip_install", return_value=23): self.assertEqual(models.install_companions(entry), 23) def test_models_with_no_extras_need_nothing(self): entry = entry_by_key("higgs_audio_v3_tts") self.assertEqual(entry.extras, ()) self.assertEqual(models.missing_companions(entry), []) class PythonEnvTests(unittest.TestCase): """Interpreter selection for the version-pinned venv.""" def test_env_compatible_needs_310_to_312(self): with patch.object(pythonenv, "env_version", return_value=(3, 12)): self.assertTrue(pythonenv.env_compatible()) with patch.object(pythonenv, "env_version", return_value=(3, 13)): self.assertFalse(pythonenv.env_compatible()) with patch.object(pythonenv, "env_version", return_value=None): self.assertFalse(pythonenv.env_compatible()) def test_prepare_env_noop_on_compatible_venv(self): with patch.object(pythonenv, "env_compatible", return_value=True), \ patch.object(envs, "create_env") as mock_create, \ patch.object(envs, "provision_env_with_uv") as mock_uv: self.assertEqual(pythonenv.prepare_env(), 0) mock_create.assert_not_called() mock_uv.assert_not_called() def test_prepare_env_uses_compatible_system_interpreter(self): interpreter = Path("/usr/bin/python3.12") with patch.object(pythonenv, "env_compatible", return_value=False), \ patch.object(envs, "env_exists", return_value=False), \ patch.object(envs, "compatible_interpreter", return_value=interpreter) as mock_find, \ patch.object(envs, "create_env", return_value=0) as mock_create: self.assertEqual(pythonenv.prepare_env(), 0) mock_find.assert_called_once() mock_create.assert_called_once_with(SGLOMNI_ENV, interpreter) def test_prepare_env_falls_back_to_uv(self): with patch.object(pythonenv, "env_compatible", return_value=False), \ patch.object(envs, "env_exists", return_value=False), \ patch.object(envs, "compatible_interpreter", return_value=None), \ patch.object(envs, "ensure_uv", return_value=0) as mock_uv_install, \ patch.object(envs, "provision_env_with_uv", return_value=0) as mock_uv: self.assertEqual(pythonenv.prepare_env(), 0) mock_uv_install.assert_called_once() mock_uv.assert_called_once() def test_prepare_env_reports_uv_failure(self): with patch.object(pythonenv, "env_compatible", return_value=False), \ patch.object(envs, "env_exists", return_value=False), \ patch.object(envs, "compatible_interpreter", return_value=None), \ patch.object(envs, "ensure_uv", return_value=1): self.assertEqual(pythonenv.prepare_env(), 1) class BuildSpecTests(unittest.TestCase): """The managed ServerSpec: sgl-omni serve with model/config/port.""" def test_spec_hosts_the_entry_repo_with_config(self): entry = entry_by_key("qwen3_tts_0_6b_customvoice") spec = status.build_spec(entry) self.assertEqual(spec.name, constants.SERVER_NAME) self.assertEqual(spec.identity, probe.IDENTITY_SGLOMNI) self.assertEqual(spec.start_timeout, constants.SERVER_START_TIMEOUT) script, serve, flag, repo, cfg_flag, cfg, port_flag, port = spec.argv self.assertEqual(serve, "serve") self.assertEqual(flag, "--model-path") self.assertEqual(repo, entry.repo) self.assertEqual(cfg_flag, "--config") self.assertEqual(Path(cfg), config_path(entry)) self.assertEqual(port_flag, "--port") self.assertIn(port, str(spec.url)) def test_spec_launches_the_higgs_config(self): entry = entry_by_key("higgs_audio_v3_tts") spec = status.build_spec(entry) self.assertIn("--config", spec.argv) self.assertEqual(Path(spec.argv[spec.argv.index("--config") + 1]), catalog.config_path(entry)) class GpuCapabilityTests(unittest.TestCase): """The nvidia-smi-backed GPU facts are best-effort and cached.""" def setUp(self): gpu_module = status.gpu gpu_module._query.cache_clear() self.addCleanup(gpu_module._query.cache_clear) def _nvidia_smi(self, *, stdout="", returncode=0, installed=True): def fake_run(argv, **_kwargs): if not installed: raise FileNotFoundError("nvidia-smi") return SimpleNamespace(returncode=returncode, stdout=stdout) return fake_run def test_parses_name_and_compute_capability(self): with patch.object(status.gpu.shutil, "which", return_value="/x"), \ patch.object(status.gpu.subprocess, "run", side_effect=self._nvidia_smi( stdout="NVIDIA GeForce RTX 3090, 8.6\n")): self.assertEqual(status.gpu.compute_capability(), (8, 6)) self.assertEqual(status.gpu.describe(), "NVIDIA GeForce RTX 3090 (compute capability 8.6)") def test_none_when_nvidia_smi_missing(self): with patch.object(status.gpu.shutil, "which", return_value=None): self.assertIsNone(status.gpu.compute_capability()) self.assertIsNone(status.gpu.describe()) def test_none_when_the_query_fails_or_is_garbage(self): for kwargs in (dict(returncode=1), dict(stdout=""), dict(stdout="name only\n")): with self.subTest(stdout=kwargs.get("stdout")): with patch.object(status.gpu.shutil, "which", return_value="/x"), \ patch.object(status.gpu.subprocess, "run", side_effect=self._nvidia_smi(**kwargs)): self.assertIsNone(status.gpu.compute_capability()) class HiggsConfigTests(unittest.TestCase): """The vendored Higgs config: VRAM headroom and a raised frame cap.""" def test_config_declares_the_repo_budget_and_frame_cap(self): entry = entry_by_key("higgs_audio_v3_tts") path = catalog.config_path(entry) self.assertIsNotNone(path) self.assertTrue(path.is_file(), f"missing {path}") text = path.read_text(encoding="utf-8") self.assertIn(f"model_path: {entry.repo}", text) # The upstream pipeline budgets 0.98 of the card across its # colocated stages; 0.80 leaves transient-allocation headroom on # 24 GB cards (request-time CUDA OOM at the 0.85 default). self.assertRegex(text, r"gpu_memory_fraction:\s*0\.80") # The engine's 2048-frame default (~27 s at 75 fps) silently # truncates a full 250-word sub-chunk; per-request values are # clamped to this factory cap server-side. 3000 frames (~40 s) # is the most the pinned 4095-token admission window allows # after the prompt tokens. self.assertRegex(text, r"max_new_tokens:\s*3000") def test_entry_sends_the_raised_frame_cap_per_request(self): entry = entry_by_key("higgs_audio_v3_tts") self.assertEqual(entry.max_new_tokens, 3000) def test_entry_caps_sub_requests_for_the_admission_window(self): """The server pins prompt + generation at 4096 tokens for Higgs; 80 words (~30-40 s at 75 fps) narrates inside the 3000-frame cap, and the pre-flight popup offers the clamp for a run.""" entry = entry_by_key("higgs_audio_v3_tts") self.assertEqual(entry.chunk_words, 80) # A full CHUNK_SIZE sub-chunk does NOT fit one Higgs request. self.assertLess(entry.chunk_words, 250) self.assertIsNone(entry_by_key("zonos2").chunk_words) class EngineMemoryBudgetTests(unittest.TestCase): """Colocated AR pipelines with an unpinned engine budget OOM on 24 GB. Upstream auto-sizes the engine's sglang static pool to nearly all free VRAM when mem_fraction_static is unset (qwen3_tts, moss_tts, voxtral): on a 24 GB card the first /v1/audio/speech request aborts with "CUDA out of memory. Tried to allocate ~100 MiB" and every retry fails identically. These vendored configs pin the pool below auto-size, leaving several GB for the colocated vocoder, CUDA graphs, and other GPU processes — 0.70 for the small models. moss_tts pins 0.87 instead: its ~17.1 GB of v1.5 bf16 weights alone exceed the 0.70 budget on the ~20.7 GB the engine profiler sees on a 24 GB card (its colocated tokenizer stage is already resident), so a 0.70 pin would abort the boot with "Loaded weights leave no GPU memory for the KV cache" instead of fixing the request-time OOM. """ # (catalog key, engine stage name, pinned mem_fraction_static). # Voxtral's engine stage is tts_generation; every other engine-bearing # pipeline names it tts_engine. moss_tts_local is intentionally absent — # its upstream config class already budgets its colocated stages # explicitly (0.15 preprocessing / 0.67 AR / 0.18 vocoder), and pinning # would fight that logic (the codec reserve derives from the fractions). # s2-pro's OOM is model size, not budgeting (upstream issue #359). PINNED = ( ("qwen3_tts_0_6b_base", "tts_engine", 0.70), ("qwen3_tts_0_6b_customvoice", "tts_engine", 0.70), ("qwen3_tts_1_7b_base", "tts_engine", 0.70), ("qwen3_tts_1_7b_voicedesign", "tts_engine", 0.70), ("moss_tts", "tts_engine", 0.87), ("voxtral_tts", "tts_generation", 0.70), ) UNPINNED = ("moss_tts_local", "fish_s2_pro") def test_engine_budget_is_pinned_under_the_engine_stage(self): for key, stage, fraction in self.PINNED: with self.subTest(key=key, stage=stage, fraction=fraction): path = config_path(entry_by_key(key)) self.assertIsNotNone(path) self.assertTrue(path.is_file(), f"missing {path}") text = path.read_text(encoding="utf-8") self.assertIn( "stages:\n" f" {stage}:\n" " engine:\n" f" mem_fraction_static: {fraction:.2f}\n", text) def test_already_budgeted_and_oversized_models_stay_unpinned(self): for key in self.UNPINNED: with self.subTest(key=key): path = config_path(entry_by_key(key)) self.assertIsNotNone(path) self.assertTrue(path.is_file(), f"missing {path}") text = path.read_text(encoding="utf-8") self.assertNotIn("mem_fraction_static", text) class Fp8FallbackTests(unittest.TestCase): """FP8-only pipelines fall back to a vendored bf16 config on old GPUs.""" ZONOS2 = "zonos2" def _fallback(self, capability): return patch("backends.sglomni.gpu.compute_capability", return_value=capability) def test_fallback_config_declares_the_repo_and_disables_fp8(self): entry = entry_by_key("zonos2") path = catalog.fallback_config_path(entry) self.assertIsNotNone(path) self.assertTrue(path.is_file(), f"missing {path}") text = path.read_text(encoding="utf-8") self.assertIn(f"model_path: {entry.repo}", text) self.assertRegex(text, r"fp8:\s*false") # bf16 weights need a bigger static pool than the builder's 0.5 # default (24 GB card: >=0.64 for any KV cache at all). self.assertRegex(text, r"mem_fraction_static:\s*0\.70") def test_only_fp8_models_carry_a_fallback(self): for entry in ENTRIES: if entry.fp8_moe: self.assertIsNotNone(entry.fp8_min_compute_capability, entry.key) self.assertIsNotNone(entry.bf16_config, entry.key) else: self.assertIsNone(catalog.fallback_config_path(entry)) def test_fallback_needed_below_the_capability_floor(self): entry = entry_by_key("zonos2") with patch("backends.sglomni.gpu.compute_capability", return_value=(8, 6)): self.assertTrue(status.needs_fp8_fallback(entry)) self.assertEqual(status.launch_config_path(entry), catalog.fallback_config_path(entry)) note = status.gpu_fallback_note(entry) self.assertIn("bf16", note) self.assertIn("8.9", note) def test_no_fallback_at_or_above_the_capability(self): entry = entry_by_key("zonos2") for capability in ((8, 9), (9, 0), (10, 0)): with self.subTest(capability=capability): with patch("backends.sglomni.gpu.compute_capability", return_value=capability): self.assertFalse(status.needs_fp8_fallback(entry)) self.assertIsNone(status.gpu_fallback_note(entry)) self.assertEqual(status.launch_config_path(entry), config_path(entry)) def test_no_fallback_without_an_answerable_gpu(self): # A GPU this tool cannot read keeps upstream defaults instead of # second-guessing the host. entry = entry_by_key("zonos2") with patch("backends.sglomni.gpu.compute_capability", return_value=None): self.assertFalse(status.needs_fp8_fallback(entry)) self.assertIsNone(status.gpu_fallback_note(entry)) self.assertEqual(status.launch_config_path(entry), config_path(entry)) def test_non_fp8_models_never_fall_back(self): for entry in ENTRIES: if entry.key == "zonos2": continue with patch("backends.sglomni.gpu.compute_capability", return_value=(1, 0)): self.assertFalse(status.needs_fp8_fallback(entry)) def test_spec_launches_the_bf16_config_on_an_old_gpu(self): entry = entry_by_key("zonos2") with patch("backends.sglomni.gpu.compute_capability", return_value=(8, 6)): spec = status.build_spec(entry) self.assertIn("--config", spec.argv) self.assertEqual(Path(spec.argv[spec.argv.index("--config") + 1]), catalog.fallback_config_path(entry)) def test_spec_keeps_the_default_pipeline_on_modern_gpus(self): entry = entry_by_key("zonos2") with patch("backends.sglomni.gpu.compute_capability", return_value=(9, 0)): spec = status.build_spec(entry) self.assertNotIn("--config", spec.argv) def test_detect_tags_a_fallback_model(self): entry = entry_by_key("zonos2") with patch("backends.sglomni.status._is_installed", return_value=True), \ patch("backends.sglomni.status.installed_entries", return_value=[entry]), \ patch("backends.sglomni.gpu.compute_capability", return_value=(8, 6)), \ patch.object(servers, "manages", return_value=False), \ patch.object(status, "_detect_remote", return_value=([], {})): st = status.detect() self.assertIn("zonos2 (bf16 fallback)", next(line for line in st.details if line.startswith("models: "))) def test_install_prints_the_fallback_note(self): out = io.StringIO() with patch("backends.sglomni.models.prepare_env", return_value=0), \ patch("backends.common.pip_install", return_value=0), \ patch("backends.sglomni.models._hf_download_prefix", return_value=["hf"]), \ patch("backends.common.run_console_subprocess", return_value=0), \ patch("backends.sglomni.gpu.compute_capability", return_value=(8, 6)), \ redirect_stdout(out): rc = models.install_model("zonos2") self.assertEqual(rc, 0) self.assertIn("bf16", out.getvalue()) class DetectTests(unittest.TestCase): """detect() reports install state, models, and the running model.""" def _detect(self, *, script=False, module=False, entries=()): entry = entries[0] if entries else None spec = [status.build_spec(entry)] if entry else [] with patch("backends.sglomni.status._is_installed", return_value=script or module), \ patch("backends.sglomni.status.installed_entries", return_value=list(entries)), \ patch.object(servers, "manages", return_value=False), \ patch.object(status, "_detect_remote", return_value=([], {})): return status.detect(), spec def test_not_installed(self): st, _spec = self._detect() self.assertFalse(st.installed) self.assertFalse(st.configured) self.assertFalse(st.running) self.assertEqual(st.servers, []) self.assertEqual(st.partial, "") def test_installed_without_models_is_partial(self): st, _spec = self._detect(script=True) self.assertTrue(st.installed) self.assertFalse(st.configured) self.assertEqual(st.partial, "installed (no models)") def test_configured_reports_a_spec_and_ready(self): entry = entry_by_key("higgs_audio_v3_tts") st, _spec = self._detect(script=True, entries=[entry]) self.assertTrue(st.configured) self.assertTrue(st.ready) self.assertEqual(len(st.servers), 1) self.assertIn(entry.repo, st.servers[0].argv) class ProbeIdentityTests(unittest.TestCase): """A healthy sglang-omni /health identifies the sglomni backend.""" def _urlopen_returning(self, payloads): calls = {"index": 0} def fake_urlopen(url, timeout=3.0): if calls["index"] >= len(payloads): # Past the scripted payloads (e.g. the gradio fallback # probe): behave like a 404 — urlopen raises, _get_json # maps that to None. calls["index"] += 1 raise urllib.error.URLError("HTTP 404") payload = payloads[calls["index"]] calls["index"] += 1 response = MagicMock() response.__enter__.return_value = response response.read.return_value = json.dumps(payload).encode("utf-8") return response return fake_urlopen, calls def test_healthy_server_with_stages_identifies_sglomni(self): health = {"status": "healthy", "running": True, "stages": ["preprocessing", "tts_generation", "vocoder"]} fake_urlopen, _calls = self._urlopen_returning([health]) with patch.object(probe.common, "server_running", return_value=True), \ patch("backends.probe.urllib.request.urlopen", fake_urlopen): self.assertEqual(probe.identify_server("http://127.0.0.1:8100"), probe.IDENTITY_SGLOMNI) def test_unhealthy_server_is_not_sglomni(self): health = {"status": "unhealthy", "running": False, "stages": []} fake_urlopen, _calls = self._urlopen_returning([health]) with patch.object(probe.common, "server_running", return_value=True), \ patch("backends.probe.urllib.request.urlopen", fake_urlopen): self.assertIsNone(probe.identify_server("http://127.0.0.1:8100")) def test_served_model_read_from_v1_models(self): payload = {"object": "list", "data": [ {"id": "bosonai/higgs-audio-v3-tts-4b", "root": "bosonai/higgs-audio-v3-tts-4b"}]} fake_urlopen, _calls = self._urlopen_returning([payload]) with patch("backends.probe.urllib.request.urlopen", fake_urlopen): self.assertEqual( probe.sglomni_served_model("http://127.0.0.1:8100"), "bosonai/higgs-audio-v3-tts-4b") def test_served_model_none_on_garbage(self): fake_urlopen, _calls = self._urlopen_returning([{"data": []}]) with patch("backends.probe.urllib.request.urlopen", fake_urlopen): self.assertIsNone( probe.sglomni_served_model("http://127.0.0.1:8100")) def test_voice_names_read_from_uploaded_voices(self): payload = {"uploaded_voice_names": ["narrator", "second narrator"]} fake_urlopen, _calls = self._urlopen_returning([payload]) with patch("backends.probe.urllib.request.urlopen", fake_urlopen): self.assertEqual( probe.sglomni_voice_names("http://127.0.0.1:8100"), ["narrator", "second narrator"]) class ModelsScreenTests(unittest.TestCase): """models_screen: the audio.cpp-style checkbox tree driving steps. The tree is scripted (like the qwen models_screen tests): each fake render records its arguments and returns the next scripted answer; the task-view run executes its steps inline so delegation to install/uninstall_model is observable. """ FAMILY_INDEX = {option["key"]: index for index, family in enumerate(install_tree_families(list(ENTRIES))) for option in family["options"]} def _screen(self, answers, *, installed=(), package=True, confirm=True, extra=()): """Run models_screen with scripted tree answers; record calls. Returns ``(rc, trees, confirms, flashes, runs)``: trees holds one (title, kwargs) per render, confirms every uninstall question, flashes every (text, kind), and runs each (title, step titles) while executing its steps' work inline. """ import contextlib from backends.sglomni import wizard choices = list(answers) trees, confirms, flashes, runs = [], [], [], [] def fake_tree(stdscr, title, families, **kwargs): trees.append((title, kwargs)) return choices.pop(0) def fake_confirm(scr, question, **kwargs): confirms.append((question, kwargs)) return confirm def fake_flash(scr, text, kind="warn"): flashes.append((text, kind)) def fake_run(scr, title, steps, **kwargs): runs.append((title, [step.title for step in steps])) for step in steps: step.work(None, None) return 0 patches = [ patch.object(wizard, "_is_installed", return_value=package), patch.object(models, "installed_keys", return_value=list(installed)), patch.object(wizard.tui, "checkbox_tree", fake_tree), patch.object(wizard.tui, "confirm", fake_confirm), patch.object(wizard.tui, "flash", fake_flash), patch.object(wizard.taskview, "run_steps", fake_run), *extra, ] with contextlib.ExitStack() as stack: for ctx in patches: stack.enter_context(ctx) rc = wizard.models_screen(None) return rc, trees, confirms, flashes, runs def test_tree_is_the_audio_cpp_modify_flow(self): from backends.sglomni import wizard first = ENTRIES[0] rc, trees, _confirms, _flashes, _runs = self._screen( [wizard._GO_BACK], installed=(first.key,)) self.assertEqual(rc, 0) title, kwargs = trees[0] self.assertEqual(title, "Select SGLang-Omni Models") # The installed model starts checked (a modify list), Confirm is # pre-focused, and an empty selection is a valid answer. self.assertEqual(kwargs["checked"], {(self.FAMILY_INDEX[first.key], first.key)}) self.assertTrue(kwargs["start_on_buttons"]) self.assertTrue(kwargs["allow_empty"]) self.assertIs(kwargs["back_value"], wizard._GO_BACK) def test_checking_a_model_runs_one_install_step(self): from backends.sglomni import wizard entry = ENTRIES[0] requested = [] def capture(key, *, emit=None, cancel=None): requested.append(key) return 0 picked = [(self.FAMILY_INDEX[entry.key], entry.key)] rc, _trees, confirms, flashes, runs = self._screen( [picked, wizard._GO_BACK], extra=[patch.object(models, "install_model", side_effect=capture)]) self.assertEqual(rc, 0) self.assertEqual(requested, [entry.key]) self.assertEqual(runs, [("Configure SGLang-Omni", [f"Install {entry.label}"])]) self.assertEqual(confirms, []) self.assertEqual(flashes[-1], ("SGLang-Omni models updated: 1 installed.", "ok")) def test_unchecking_confirms_then_deletes_the_weights(self): from backends.sglomni import wizard entry = ENTRIES[0] rc, _trees, confirms, flashes, runs = self._screen( [[], wizard._GO_BACK], installed=(entry.key,)) self.assertEqual(rc, 0) # An empty selection is accepted (allow_empty) and uninstalls # everything installed: one confirm, one removal step. question, kwargs = confirms[0] self.assertEqual(question, "Remove cached weights for 1 model?") self.assertIn(entry.label, kwargs["body"]) self.assertEqual(runs, [("Configure SGLang-Omni", [f"Delete {entry.label} weights"])]) self.assertEqual(flashes[-1], ("SGLang-Omni models updated: 1 removed.", "ok")) def test_uninstall_step_stops_nothing_and_deletes_real_weights(self): # The removal step is uninstall_model itself: with a redirected # HF cache the seeded weight directory is really deleted. from backends.sglomni import wizard entry = ENTRIES[0] with tempfile.TemporaryDirectory() as td: directory = Path(td) / ("models--" + entry.repo.replace("/", "--")) (directory / "refs").mkdir(parents=True) (directory / "refs" / "main").write_text("hash\n") (directory / "snapshots" / "abc").mkdir(parents=True) (directory / "snapshots" / "abc" / "weights.safetensors") \ .write_bytes(b"x") rc, _trees, _confirms, _flashes, _runs = self._screen( [[], wizard._GO_BACK], installed=(entry.key,), extra=[ patch("backends.common.hf_cache_dir", return_value=Path(td)), patch.object(models, "_managed_running_repo", return_value=None), ]) self.assertEqual(rc, 0) self.assertFalse(directory.exists()) def test_declining_the_uninstall_confirm_runs_nothing(self): from backends.sglomni import wizard entry = ENTRIES[0] rc, trees, confirms, flashes, runs = self._screen( [[], wizard._GO_BACK], installed=(entry.key,), confirm=False) self.assertEqual(rc, 0) # The decline re-opens the tree (second render), no work happens. self.assertEqual(len(trees), 2) self.assertEqual(len(confirms), 1) self.assertEqual(runs, []) self.assertEqual(flashes, []) def test_install_without_package_flashes_guidance_instead(self): from backends.sglomni import wizard entry = ENTRIES[0] picked = [(self.FAMILY_INDEX[entry.key], entry.key)] rc, _trees, confirms, flashes, runs = self._screen( [picked, wizard._GO_BACK], package=False) self.assertEqual(rc, 0) self.assertEqual(runs, []) self.assertEqual(confirms, []) self.assertEqual(flashes, [( "Install the SGLang-Omni backend first " "(Configure Backends > Install Backend).", "warn")]) def test_unchanged_selection_re_opens_the_tree(self): from backends.sglomni import wizard entry = ENTRIES[0] picked = [(self.FAMILY_INDEX[entry.key], entry.key)] rc, trees, _confirms, flashes, runs = self._screen( [picked, wizard._GO_BACK], installed=(entry.key,)) self.assertEqual(rc, 0) self.assertEqual(len(trees), 2) self.assertEqual(runs, []) self.assertEqual(flashes, []) def test_mixed_selection_removes_before_downloading(self): # Unchecking the installed model and checking another in one # confirm: a single run whose removal step precedes the download. from backends.sglomni import wizard gone, added = ENTRIES[0], ENTRIES[1] picked = [(self.FAMILY_INDEX[added.key], added.key)] rc, _trees, confirms, flashes, runs = self._screen( [picked, wizard._GO_BACK], installed=(gone.key,)) self.assertEqual(rc, 0) self.assertEqual(len(confirms), 1) self.assertEqual(runs, [( "Configure SGLang-Omni", [f"Delete {gone.label} weights", f"Install {added.label}"])]) self.assertEqual( flashes[-1], ("SGLang-Omni models updated: 1 installed, 1 removed.", "ok")) class SetupWizardTests(unittest.TestCase): """_wizard: the setup tree reconciles models like the Configure screen.""" FAMILY_INDEX = {option["key"]: index for index, family in enumerate(install_tree_families(list(ENTRIES))) for option in family["options"]} def _wizard(self, answers, *, installed=(), package=True, confirm=True): """Run _wizard with scripted tree answers; return (settings, trees, confirms).""" from backends.sglomni import wizard trees, confirms = [], [] def fake_tree(stdscr, title, families, **kwargs): trees.append((title, kwargs)) return answers.pop(0) def fake_confirm(scr, question, **kwargs): confirms.append(question) return confirm args = wizard.build_parser().parse_args([]) with patch.object(wizard, "_preflight", return_value=[]), \ patch.object(wizard, "_gpu_warning", return_value=None), \ patch.object(wizard, "_is_installed", return_value=package), \ patch.object(models, "installed_keys", return_value=list(installed)), \ patch.object(wizard.tui, "checkbox_tree", fake_tree), \ patch.object(wizard.tui, "confirm", fake_confirm), \ patch.object(wizard.tui, "flash", lambda *a, **k: None): settings = wizard._wizard(None, args) return settings, trees, confirms def test_esc_aborts(self): from backends.sglomni import wizard settings, _trees, confirms = self._wizard([wizard._GO_BACK]) self.assertIsNone(settings) self.assertEqual(confirms, []) def test_modify_flow_installs_new_and_keeps_installed(self): from backends.sglomni import wizard first, second = ENTRIES[0], ENTRIES[1] # The installed model stays checked (kept as-is); the new one is # added — the diff installs the new one only. picked = [(self.FAMILY_INDEX[first.key], first.key), (self.FAMILY_INDEX[second.key], second.key)] settings, _trees, confirms = self._wizard( [picked], installed=(first.key,)) self.assertEqual(settings["keys"], [second.key]) self.assertEqual(settings["uninstall_keys"], []) self.assertEqual(confirms, []) def test_unchecking_requires_a_confirm_then_uninstalls(self): from backends.sglomni import wizard first = ENTRIES[0] settings, _trees, confirms = self._wizard([[]], installed=(first.key,)) self.assertEqual(settings["keys"], []) self.assertEqual(settings["uninstall_keys"], [first.key]) self.assertEqual(confirms, ["Remove cached weights for 1 model?"]) def test_declined_confirm_re_opens_the_tree(self): from backends.sglomni import wizard first = ENTRIES[0] settings, trees, confirms = self._wizard( [[], wizard._GO_BACK], installed=(first.key,), confirm=False) self.assertIsNone(settings) self.assertEqual(len(trees), 2) self.assertEqual(len(confirms), 1) def test_empty_tree_installs_the_package_only(self): from backends.sglomni import wizard settings, _trees, confirms = self._wizard([[]]) self.assertEqual(settings["keys"], []) self.assertEqual(settings["uninstall_keys"], []) self.assertEqual(confirms, []) def test_setup_tree_starts_minimized(self): """Like audio.cpp: no expand_all — every family starts collapsed.""" from backends.sglomni import wizard _settings, trees, _confirms = self._wizard([wizard._GO_BACK]) self.assertEqual(trees[0][0], "Select SGLang-Omni Models to Install") self.assertFalse(trees[0][1].get("expand_all", False)) def test_steps_remove_before_downloading(self): from backends.sglomni import wizard gone, added = ENTRIES[0], ENTRIES[1] steps = wizard._execute_steps({ "do_python": False, "do_install": False, "uninstall_keys": [gone.key], "keys": [added.key]}) self.assertEqual([step.title for step in steps], [f"Delete {gone.label} weights", f"Install {added.label}"]) def test_run_tui_drives_the_wizard_in_a_curses_session(self): # The standalone CLI wraps the wizard in its own curses session # (passing a real screen through) and runs the model work as a # console tail afterwards. from backends.sglomni import wizard settings = {"keys": [], "uninstall_keys": [], "do_python": False, "do_install": False} screens = [] with patch.object(wizard, "_wizard", side_effect=lambda scr, args: screens.append(scr) or settings), \ patch.object(wizard, "_execute", return_value=7) as execute, \ patch("curses.wrapper", side_effect=lambda fn: fn("SCREEN")): rc = wizard.run_tui(wizard.build_parser().parse_args([])) self.assertEqual(rc, 7) self.assertEqual(screens, ["SCREEN"]) execute.assert_called_once_with(settings) class SystemDepTests(unittest.TestCase): """Docs-required system binaries are checked at install time.""" def test_dac_models_require_ffmpeg(self): # The Fish Audio and ZONOS2 pipelines shell out to ffmpeg; the # install flow must warn like it does for the Qwen entries' sox. for key in ("fish_s2_pro", "zonos2"): with self.subTest(key=key): entry = entry_by_key(key) self.assertEqual(entry.system_dep, "ffmpeg") self.assertIn("ffmpeg", entry.system_hint) def test_missing_system_dep_warns(self): with patch("shutil.which", return_value=None): message = models.system_dep_missing(entry_by_key("fish_s2_pro")) self.assertIsNotNone(message) self.assertIn("ffmpeg", message) def test_present_system_dep_is_silent(self): with patch("shutil.which", return_value="/usr/bin/ffmpeg"): self.assertIsNone( models.system_dep_missing(entry_by_key("fish_s2_pro"))) def test_missing_system_dep_warns_during_install(self): out = io.StringIO() with patch("shutil.which", return_value=None), \ patch("backends.sglomni.models.prepare_env", return_value=0), \ patch("backends.common.pip_install", return_value=0), \ patch("backends.sglomni.models._hf_download_prefix", return_value=["hf"]), \ patch("backends.common.run_console_subprocess", return_value=0), \ redirect_stdout(out): rc = models.install_model("fish_s2_pro") self.assertEqual(rc, 0) self.assertIn("ffmpeg", out.getvalue()) class UninstallModelServerStopTests(unittest.TestCase): """uninstall_model stops a managed server hosting the model.""" def test_stops_a_server_hosting_the_model(self): entry = ENTRIES[0] with patch.object(models, "_managed_running_repo", return_value=entry.repo), \ patch.object(models, "delete_model_weights"), \ patch.object(servers, "stop") as stop: rc = models.uninstall_model(entry.key) self.assertEqual(rc, 0) stop.assert_called_once_with(constants.SERVER_NAME) def test_leaves_a_server_hosting_something_else_alone(self): entry = ENTRIES[0] with patch.object(models, "_managed_running_repo", return_value=ENTRIES[1].repo), \ patch.object(servers, "stop") as stop, \ patch.object(models, "delete_model_weights") as weights: rc = models.uninstall_model(entry.key) self.assertEqual(rc, 0) stop.assert_not_called() weights.assert_called_once_with([entry]) class WizardUninstallTests(unittest.TestCase): """wizard.uninstall: stop the server, delete weights, remove the venv. The venv removal IS the cleanup (the heavyweight CUDA stack is the install): there is deliberately no pip-uninstall phase ahead of it. """ def _run(self, *, pid=None, cancel_after=None, venv=True, pythons=True, rmtree_removes=True): """Run uninstall with a temp venv tree; return (rc, output, stop_calls, weights, pip_uninstall). CANCEL_AFTER N lets the first N cancel checks pass (None = none of them do). """ from backends.sglomni import wizard out = io.StringIO() with tempfile.TemporaryDirectory() as td: venv_dir = Path(td) / "sglomni" pythons_dir = Path(td) / "pythons" for make, directory in ((venv, venv_dir), (pythons, pythons_dir)): if make: directory.mkdir() state = {"stops": [], "cancel": 0} def fake_stop(name): state["stops"].append(name) def fake_cancel(_cancel): state["cancel"] += 1 return cancel_after is not None \ and state["cancel"] > cancel_after real_rmtree = shutil.rmtree def fake_rmtree(path, ignore_errors=False): if rmtree_removes: real_rmtree(path, ignore_errors=True) with patch.object(servers, "pid_for", return_value=pid), \ patch.object(servers, "stop", side_effect=fake_stop), \ patch.object(models, "delete_model_weights") as weights, \ patch.object(common, "pip_uninstall") as pip_uninstall, \ patch.object(wizard, "SGLOMNI_ENV", venv_dir), \ patch.object(envs, "PYTHON_INSTALL_DIR", pythons_dir), \ patch.object(common, "cancel_requested", side_effect=fake_cancel), \ patch("shutil.rmtree", side_effect=fake_rmtree), \ redirect_stdout(out): rc = wizard.uninstall() existed = venv_dir.exists() or pythons_dir.exists() return rc, out.getvalue(), state["stops"], weights, \ pip_uninstall, existed def test_removes_the_venv_tree_without_a_pip_phase(self): rc, out, _stops, weights, pip_uninstall, existed = self._run() self.assertEqual(rc, 0) weights.assert_called_once_with() pip_uninstall.assert_not_called() self.assertFalse(existed) def test_stops_only_a_running_managed_server(self): _rc, _out, stops, _w, _p, _e = self._run(pid=123) self.assertEqual(stops, ["sglomni"]) _rc, _out, stops, _w, _p, _e = self._run(pid=None) self.assertEqual(stops, []) def test_cancel_before_anything_removes_nothing(self): rc, _out, _stops, weights, _p, _e = self._run(cancel_after=0) self.assertEqual(rc, 130) weights.assert_not_called() def test_cancel_before_the_venv_removal_keeps_the_dirs(self): rc, _out, _stops, weights, _p, existed = self._run(cancel_after=1) self.assertEqual(rc, 130) weights.assert_called_once_with() self.assertTrue(existed) def test_a_stuck_venv_warns_but_the_uninstall_succeeds(self): rc, out, _stops, _w, _p, existed = self._run(rmtree_removes=False) self.assertEqual(rc, 0) self.assertTrue(existed) self.assertIn("Could not fully remove", out) def test_missing_venv_dirs_are_fine(self): rc, _out, _stops, _w, _p, _e = self._run(venv=False, pythons=False) self.assertEqual(rc, 0) class WizardUpdateTests(unittest.TestCase): """wizard.update: package upgrade, then a companion refresh. The refresh re-runs every installed model's extras (a satisfied pin is a pip no-op), healing version drift the import probe cannot see. """ def _run(self, *, pid=None, venv=True, pip_rc=0, companions_rc=0, installed=(0, 1), cancel_after=None): from backends.sglomni import wizard out = io.StringIO() entries = [ENTRIES[i] for i in installed] state = {"pip": [], "companions": [], "stops": [], "cancel": 0} def fake_stop(name): state["stops"].append(name) def fake_cancel(_cancel): state["cancel"] += 1 return cancel_after is not None \ and state["cancel"] > cancel_after def fake_pip(specs, **kwargs): state["pip"].append((list(specs), kwargs)) return pip_rc def fake_companions(entry, **kwargs): state["companions"].append((entry.key, kwargs.get("force"))) return companions_rc with patch.object(servers, "pid_for", return_value=pid), \ patch.object(servers, "stop", side_effect=fake_stop), \ patch.object(envs, "env_exists", return_value=venv), \ patch.object(common, "pip_install", side_effect=fake_pip), \ patch.object(models, "installed_entries", return_value=entries), \ patch.object(models, "install_companions", side_effect=fake_companions), \ patch.object(common, "cancel_requested", side_effect=fake_cancel), \ redirect_stdout(out): rc = wizard.update() return rc, out.getvalue(), state def test_no_venv_is_a_no_op(self): rc, _out, state = self._run(venv=False) self.assertEqual(rc, 0) self.assertEqual(state["pip"], []) self.assertEqual(state["companions"], []) def test_stops_the_server_then_upgrades_the_package(self): rc, _out, state = self._run(pid=123) self.assertEqual(rc, 0) self.assertEqual(state["stops"], ["sglomni"]) specs, kwargs = state["pip"][0] self.assertEqual(specs, ["sglang-omni"]) self.assertTrue(kwargs["upgrade"]) self.assertEqual(kwargs["extra_args"], ["--pre"]) def test_upgrade_failure_skips_the_companion_refresh(self): rc, _out, state = self._run(pip_rc=23) self.assertEqual(rc, 23) self.assertEqual(state["companions"], []) def test_refreshes_every_installed_models_companions(self): rc, _out, state = self._run(installed=(0, 4)) self.assertEqual(rc, 0) self.assertEqual(state["companions"], [(ENTRIES[0].key, True), (ENTRIES[4].key, True)]) def test_companion_failure_warns_but_the_update_succeeds(self): rc, out, _state = self._run(companions_rc=23) self.assertEqual(rc, 0) self.assertIn("companion packages", out) def test_cancel_before_pip_removes_nothing(self): rc, _out, state = self._run(cancel_after=0) self.assertEqual(rc, 130) self.assertEqual(state["pip"], []) class NonInteractiveCliTests(unittest.TestCase): """_collect_from_flags: the flag-driven (non-TUI) setup path. Positional keys are the space-separated twin of --models: both feed the same deduplicated install list. """ def _collect(self, argv, *, installed=True): from backends.sglomni import wizard parser = wizard.build_parser() args = parser.parse_args(argv) with patch.object(wizard, "_preflight", return_value=[]), \ patch.object(wizard, "_gpu_warning", return_value=None), \ patch.object(wizard, "_is_installed", return_value=installed): return wizard._collect_from_flags(args, parser) def test_positional_keys_are_installed(self): settings = self._collect(["higgs_audio_v3_tts"]) self.assertEqual(settings["keys"], ["higgs_audio_v3_tts"]) self.assertTrue(settings["do_python"]) self.assertFalse(settings["do_install"]) def test_positional_and_models_flags_merge_in_order(self): settings = self._collect(["moss_tts", "--models", "higgs_audio_v3_tts,moss_tts"]) self.assertEqual(settings["keys"], ["moss_tts", "higgs_audio_v3_tts"]) def test_all_installs_every_catalog_model(self): settings = self._collect(["--all"]) self.assertEqual(settings["keys"], [entry.key for entry in ENTRIES]) def test_unknown_key_stops_the_run(self): with self.assertRaises(SystemExit): self._collect(["nope"]) def test_no_keys_installs_the_package_only(self): settings = self._collect([], installed=False) self.assertEqual(settings["keys"], []) self.assertTrue(settings["do_install"]) def test_skip_flags_are_honored(self): settings = self._collect(["--skip-python", "--skip-install", "higgs_audio_v3_tts"]) self.assertFalse(settings["do_python"]) self.assertFalse(settings["do_install"]) if __name__ == "__main__": unittest.main()