path: root/app/tests/test_backends_sglomni.py

"""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()