9db42d6cc1
With --no-create-home, os.path.expanduser("~") resolves to "/" causing
kraken get to write to /.local/share/htrmopo. Replace with
os.environ.get("HTRMOPO_DIR", "/app/models/.htrmopo") so the path is
explicit and override-friendly without a home directory.
Adds two tests verifying env-var resolution and ~-free default.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
95 lines
3.8 KiB
Python
95 lines
3.8 KiB
Python
"""Unit tests for ensure_blla_model.main()."""
|
|
|
|
import importlib
|
|
import os
|
|
from unittest.mock import MagicMock, call, patch
|
|
|
|
import ensure_blla_model
|
|
|
|
|
|
# ─── HTRMOPO_DIR env var resolution ──────────────────────────────────────────
|
|
|
|
|
|
def test_htrmopo_dir_reads_from_env_var():
|
|
"""HTRMOPO_DIR uses the HTRMOPO_DIR env var when set, not ~ expansion."""
|
|
with patch.dict(os.environ, {"HTRMOPO_DIR": "/custom/htrmopo"}):
|
|
importlib.reload(ensure_blla_model)
|
|
result = ensure_blla_model.HTRMOPO_DIR
|
|
importlib.reload(ensure_blla_model)
|
|
assert result == "/custom/htrmopo"
|
|
|
|
|
|
def test_htrmopo_dir_default_is_fixed_path():
|
|
"""Default HTRMOPO_DIR is a fixed path not derived from ~ (no-create-home safe)."""
|
|
clean_env = {k: v for k, v in os.environ.items() if k != "HTRMOPO_DIR"}
|
|
with patch.dict(os.environ, clean_env, clear=True):
|
|
importlib.reload(ensure_blla_model)
|
|
result = ensure_blla_model.HTRMOPO_DIR
|
|
importlib.reload(ensure_blla_model)
|
|
assert "~" not in result
|
|
assert not result.startswith("/.")
|
|
|
|
|
|
# ─── Model already loadable ───────────────────────────────────────────────────
|
|
|
|
|
|
def test_main_returns_early_when_model_is_loadable():
|
|
"""When the model exists and loads cleanly, no download or rename occurs."""
|
|
with (
|
|
patch("os.path.exists", return_value=True),
|
|
patch.object(ensure_blla_model, "_model_is_loadable", return_value=True),
|
|
patch.object(ensure_blla_model, "_download_blla") as mock_download,
|
|
patch("os.rename") as mock_rename,
|
|
):
|
|
ensure_blla_model.main()
|
|
|
|
mock_download.assert_not_called()
|
|
mock_rename.assert_not_called()
|
|
|
|
|
|
# ─── Model exists but is incompatible ─────────────────────────────────────────
|
|
|
|
|
|
def test_main_replaces_incompatible_model():
|
|
"""An incompatible model is renamed and replaced with a fresh download."""
|
|
fake_path = "/app/models/blla.mlmodel"
|
|
downloaded_path = "/tmp/downloaded.mlmodel"
|
|
|
|
with (
|
|
patch.object(ensure_blla_model, "BLLA_MODEL_PATH", fake_path),
|
|
patch("os.path.exists", return_value=True),
|
|
patch.object(ensure_blla_model, "_model_is_loadable", return_value=False),
|
|
patch.object(ensure_blla_model, "_download_blla", return_value=downloaded_path),
|
|
patch("os.rename") as mock_rename,
|
|
patch("shutil.copy2") as mock_copy,
|
|
patch("os.makedirs"),
|
|
):
|
|
ensure_blla_model.main()
|
|
|
|
mock_rename.assert_called_once_with(fake_path, fake_path + ".incompatible")
|
|
mock_copy.assert_called_once_with(downloaded_path, fake_path)
|
|
|
|
|
|
# ─── Model missing ────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_main_downloads_when_model_missing():
|
|
"""When the model file doesn't exist at all, it is downloaded without rename."""
|
|
fake_path = "/app/models/blla.mlmodel"
|
|
downloaded_path = "/tmp/downloaded.mlmodel"
|
|
|
|
with (
|
|
patch.object(ensure_blla_model, "BLLA_MODEL_PATH", fake_path),
|
|
patch("os.path.exists", return_value=False),
|
|
patch.object(ensure_blla_model, "_model_is_loadable") as mock_loadable,
|
|
patch.object(ensure_blla_model, "_download_blla", return_value=downloaded_path),
|
|
patch("os.rename") as mock_rename,
|
|
patch("shutil.copy2") as mock_copy,
|
|
patch("os.makedirs"),
|
|
):
|
|
ensure_blla_model.main()
|
|
|
|
mock_loadable.assert_not_called()
|
|
mock_rename.assert_not_called()
|
|
mock_copy.assert_called_once_with(downloaded_path, fake_path)
|