diff options
Diffstat (limited to 'app/tests/test_backends_sglomni.py')
| -rw-r--r-- | app/tests/test_backends_sglomni.py | 828 |
1 files changed, 828 insertions, 0 deletions
diff --git a/app/tests/test_backends_sglomni.py b/app/tests/test_backends_sglomni.py new file mode 100644 index 0000000..a11aa92 --- /dev/null +++ b/app/tests/test_backends_sglomni.py @@ -0,0 +1,828 @@ +"""Tests for the SGLang-Omni backend package (backends/sglomni).""" + +import io +import json +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 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.object(models, "_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 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_omits_config_when_none_needed(self): + entry = entry_by_key("higgs_audio_v3_tts") + spec = status.build_spec(entry) + self.assertNotIn("--config", spec.argv) + + +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 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.object(models, "_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_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) + + +if __name__ == "__main__": + unittest.main() |