100 lines
3.3 KiB
Python
100 lines
3.3 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[2]
|
|
SCRIPT = ROOT / "scripts" / "configure_yolo_model.py"
|
|
|
|
|
|
def _run_configure(*args: str) -> subprocess.CompletedProcess[str]:
|
|
return subprocess.run(
|
|
[sys.executable, str(SCRIPT), *args, "--json"],
|
|
cwd=ROOT,
|
|
check=False,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
|
|
|
|
def test_configure_yolo_model_reports_no_local_model(tmp_path: Path) -> None:
|
|
result = _run_configure("--models-dir", str(tmp_path), "--env-file", str(tmp_path / ".env"))
|
|
payload = json.loads(result.stdout)
|
|
|
|
assert result.returncode == 2
|
|
assert payload["status"] == "no_model_found"
|
|
assert payload["will_download_models"] is False
|
|
assert payload["env_updates"] == {}
|
|
assert not (tmp_path / ".env").exists()
|
|
|
|
|
|
def test_configure_yolo_model_refuses_ambiguous_model_selection(tmp_path: Path) -> None:
|
|
(tmp_path / "a.pt").write_bytes(b"model-a")
|
|
(tmp_path / "b.onnx").write_bytes(b"model-b")
|
|
|
|
result = _run_configure("--models-dir", str(tmp_path), "--env-file", str(tmp_path / ".env"))
|
|
payload = json.loads(result.stdout)
|
|
|
|
assert result.returncode == 3
|
|
assert payload["status"] == "multiple_models_found"
|
|
assert len(payload["candidates"]) == 2
|
|
assert payload["env_updates"] == {}
|
|
assert not (tmp_path / ".env").exists()
|
|
|
|
|
|
def test_configure_yolo_model_dry_run_selects_single_model(tmp_path: Path) -> None:
|
|
model_path = tmp_path / "nested" / "detector.pt"
|
|
model_path.parent.mkdir()
|
|
model_path.write_bytes(b"model")
|
|
|
|
result = _run_configure(
|
|
"--models-dir",
|
|
str(tmp_path),
|
|
"--container-model-dir",
|
|
"/app/models",
|
|
"--env-file",
|
|
str(tmp_path / ".env"),
|
|
)
|
|
payload = json.loads(result.stdout)
|
|
|
|
assert result.returncode == 0
|
|
assert payload["status"] == "ready_to_apply"
|
|
assert payload["selected_host_model_path"] == str(model_path)
|
|
assert payload["selected_container_model_path"] == "/app/models/nested/detector.pt"
|
|
assert payload["env_updates"]["GEOINTEL_INSTALL_AI"] == "true"
|
|
assert payload["env_updates"]["YOLO_ENABLED"] == "true"
|
|
assert payload["env_updates"]["YOLO_MODELS_DIR"] == "/app/models"
|
|
assert payload["env_updates"]["YOLO_MODEL_PATH"] == "/app/models/nested/detector.pt"
|
|
assert payload["will_download_models"] is False
|
|
assert not (tmp_path / ".env").exists()
|
|
|
|
|
|
def test_configure_yolo_model_apply_updates_existing_env_file(tmp_path: Path) -> None:
|
|
model_path = tmp_path / "detector.engine"
|
|
model_path.write_bytes(b"model")
|
|
env_file = tmp_path / ".env"
|
|
env_file.write_text("GEOINTEL_FRONTEND_PORT=1202\nYOLO_ENABLED=false\n", encoding="utf-8")
|
|
|
|
result = _run_configure(
|
|
"--models-dir",
|
|
str(tmp_path),
|
|
"--container-model-dir",
|
|
"/app/models",
|
|
"--env-file",
|
|
str(env_file),
|
|
"--apply",
|
|
)
|
|
payload = json.loads(result.stdout)
|
|
|
|
assert result.returncode == 0
|
|
assert payload["status"] == "applied"
|
|
contents = env_file.read_text(encoding="utf-8")
|
|
assert "GEOINTEL_FRONTEND_PORT=1202" in contents
|
|
assert "GEOINTEL_INSTALL_AI=true" in contents
|
|
assert "YOLO_ENABLED=true" in contents
|
|
assert "YOLO_MODELS_DIR=/app/models" in contents
|
|
assert "YOLO_MODEL_PATH=/app/models/detector.engine" in contents
|