1 files changed, 77 insertions, 1 deletions
diff --git a/app/tests/test_backends_sglomni.py b/app/tests/test_backends_sglomni.py
index a11aa92..1fd5f5a 100644
--- a/app/tests/test_backends_sglomni.py
+++ b/app/tests/test_backends_sglomni.py
@@ -11,7 +11,7 @@ from pathlib import Path
 from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
-from backends import envs, probe, servers
+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, \
@@ -184,6 +184,82 @@ class ModelInstallStateTests(unittest.TestCase):
         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")
+
+    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_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."""