| """First-run setup decisions (no packages are installed by these tests)."""
|
|
|
| from __future__ import annotations
|
|
|
| import json
|
| import os
|
| from pathlib import Path
|
| import tempfile
|
| import unittest
|
|
|
| from unittest import mock
|
|
|
| from scyllasband import download
|
| from scyllasband.bootstrap import _extra, missing_packages, plan_setup, wanted_backend, wants_onnx
|
|
|
| MISSING = {"ai_edge_litert": "ai-edge-litert"}
|
|
|
|
|
| def plan(**overrides):
|
| state = dict(checkout=True, venv_ready=False, in_venv=False, in_repo_venv=False, python_ok=True, previous=None)
|
| state.update(overrides)
|
| return plan_setup(MISSING, **state)
|
|
|
|
|
| class PlanSetupTest(unittest.TestCase):
|
| def test_nothing_missing(self):
|
| self.assertEqual(plan_setup({}, checkout=True, venv_ready=False, in_venv=False, in_repo_venv=False, python_ok=True), "ok")
|
|
|
| def test_fresh_checkout_creates_the_venv(self):
|
| self.assertEqual(plan(), "create-venv")
|
|
|
| def test_existing_venv_is_used_without_asking(self):
|
| self.assertEqual(plan(venv_ready=True), "use-venv")
|
| self.assertEqual(plan(venv_ready=True, in_venv=True), "use-venv")
|
|
|
| def test_inside_a_virtualenv_installs_there(self):
|
| self.assertEqual(plan(venv_ready=True, in_venv=True, in_repo_venv=True), "install-here")
|
| self.assertEqual(plan(in_venv=True), "install-here")
|
|
|
| def test_old_python_is_refused_before_creating_anything(self):
|
| self.assertEqual(plan(python_ok=False), "fail-python")
|
|
|
| def test_installed_copy_only_gets_advice(self):
|
| self.assertEqual(plan(checkout=False), "fail-installed")
|
|
|
| def test_never_loops(self):
|
| self.assertEqual(plan(in_venv=True, in_repo_venv=True, venv_ready=True, previous="installed"), "fail-loop")
|
| self.assertEqual(plan(venv_ready=True, previous="use-venv"), "fail-loop")
|
| self.assertEqual(plan(venv_ready=True, in_venv=True, in_repo_venv=True, previous="use-venv"), "install-here")
|
|
|
|
|
| class WantsOnnxTest(unittest.TestCase):
|
| def test_backend_flag(self):
|
| self.assertTrue(wants_onnx(["speak", "--backend", "onnx", "hi"]))
|
| self.assertTrue(wants_onnx(["speak", "--backend=onnx"]))
|
| self.assertFalse(wants_onnx(["speak", "--backend", "litert"]))
|
| self.assertFalse(wants_onnx(["download", "--flavor", "onnx"]))
|
|
|
| def test_bundle_manifest(self):
|
| with tempfile.TemporaryDirectory() as tmp:
|
| Path(tmp, "manifest.json").write_text(json.dumps(dict(preferred_backends=["onnx"])))
|
| self.assertTrue(wants_onnx(["speak", "--bundle", tmp]))
|
| self.assertFalse(wants_onnx(["speak", "--bundle", tmp, "--backend", "litert"]))
|
|
|
|
|
| class CudaExtraTest(unittest.TestCase):
|
| def test_gpu_on_onnx_installs_onnx_runtime_for_cuda(self):
|
| self.assertEqual(_extra(["speak", "--backend", "onnx", "--compute-units", "gpu", "hi"]), "cuda")
|
| self.assertEqual(_extra(["speak", "--backend", "onnx", "--compute-units=cpu,vector_estimator=cuda"]), "cuda")
|
| self.assertEqual(_extra(["speak", "--backend", "onnx", "--compute-units", "cpu"]), "onnx")
|
| self.assertEqual(_extra(["speak", "--backend", "onnx"]), "onnx")
|
| self.assertIsNone(_extra(["speak", "--backend", "litert", "--compute-units", "gpu"]))
|
| with mock.patch("importlib.util.find_spec", return_value=None):
|
| missing = missing_packages(["speak", "--backend", "onnx", "--compute-units", "gpu"])
|
| self.assertEqual(missing["onnxruntime"], "onnxruntime-gpu")
|
|
|
|
|
| class WantedBackendTest(unittest.TestCase):
|
| def test_apple_bundles_name_their_runtime(self):
|
| with tempfile.TemporaryDirectory() as tmp:
|
| Path(tmp, "manifest.json").write_text(json.dumps(dict(preferred_backends=["coreml"])))
|
| self.assertEqual(wanted_backend(["speak", "--bundle", tmp]), "coreml")
|
| self.assertEqual(wanted_backend(["speak", "--backend=coreai"]), "coreai")
|
| self.assertIsNone(wanted_backend(["download"]))
|
|
|
| def test_model_commands_use_the_installed_bundle_this_machine_picks(self):
|
| with tempfile.TemporaryDirectory() as tmp:
|
| for flavor in ("litert", "coreml"):
|
| Path(tmp, flavor).mkdir()
|
| Path(tmp, flavor, "manifest.json").write_text("{}")
|
| with mock.patch.object(download, "macos_major", return_value=26):
|
| self.assertEqual(wanted_backend(["speak", "--models-dir", tmp, "hi"]), "coreml")
|
| with mock.patch.object(download, "macos_major", return_value=None):
|
| self.assertEqual(wanted_backend(["speak", "--models-dir", tmp, "hi"]), "litert")
|
|
|
|
|
| class ChooseFlavorTest(unittest.TestCase):
|
| def choose(self, major, sdk=None, answers=None):
|
| replies = iter(answers or [])
|
| with mock.patch.object(download, "macos_major", return_value=major), \
|
| mock.patch.object(download, "xcode_sdk_major", return_value=sdk), open(os.devnull, "w") as quiet:
|
| return download.choose_flavors(interactive=answers is not None, ask=lambda _: next(replies), out=quiet)
|
|
|
| def test_macs_run_core_ml_and_other_machines_litert(self):
|
| for major in (15, 26, 27):
|
| with mock.patch.object(download, "macos_major", return_value=major):
|
| self.assertEqual(download.recommended_flavor(), "coreml")
|
| with mock.patch.object(download, "macos_major", return_value=None):
|
| self.assertEqual(download.recommended_flavor(), "litert")
|
| self.assertEqual(self.choose(None), ["litert"])
|
|
|
| def test_core_ai_is_offered_only_to_setups_that_build_for_os_27(self):
|
| self.assertEqual(self.choose(26, sdk=26, answers=[]), ["coreml"])
|
| self.assertEqual(self.choose(26, sdk=27, answers=["y"]), ["coreml", "coreai"])
|
| self.assertEqual(self.choose(27, sdk=None, answers=[""]), ["coreml"])
|
| self.assertEqual(self.choose(27, answers=["maybe", "yes"]), ["coreml", "coreai"])
|
| self.assertEqual(self.choose(27), ["coreml"])
|
|
|
| def test_a_mac_with_both_bundles_uses_core_ml(self):
|
| with tempfile.TemporaryDirectory() as tmp:
|
| for flavor in ("coreai", "coreml"):
|
| Path(tmp, flavor).mkdir()
|
| Path(tmp, flavor, "manifest.json").write_text("{}")
|
| with mock.patch.object(download, "macos_major", return_value=27):
|
| self.assertEqual(download.default_installed_flavor(tmp), "coreml")
|
|
|
|
|
| if __name__ == "__main__":
|
| unittest.main()
|
|
|
