Lap Tran
81ccf548e5
CI / build (push) Has been cancelled
Batch 1 (High severity): - Bug 1: --audit cat3 now checks .automaton/tasks/ paths - Bug 4: Verdict PASS/FAIL uses structured ## Status: line parsing - Bug 5: register-guards.sh checks .json/.jsonc, writes plugin key, strips comments - Bug 7: --can-edit/--scope-check path prefix uses os.sep boundary Batch 2 (Medium/Low severity): - Bug 2: migrate-project.sh find command parentheses for -prune binding - Bug 3: vram_detect model prefix matching with known-suffix whitelist - Bug 6: dashboard reads .state file before artifact heuristic fallback - Bug 8: removed wildcard CORS, added security headers (nosniff, DENY) - Bug 9: stale-task detection uses .state.lastedit instead of .state mtime - Bug 10: TEST_PLAN.md maps to test_design (was implement) 249 tests pass (up from 235). All 10 tasks driven through full workflow to completion.
306 lines
9.9 KiB
Python
306 lines
9.9 KiB
Python
"""Tests for scripts/vram_detect.py."""
|
|
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
import scripts.vram_detect as vram
|
|
|
|
|
|
def test_lookup_model_context() -> None:
|
|
assert vram._lookup_model_context("gpt-4o") == 128_000
|
|
assert vram._lookup_model_context("claude-3-5-sonnet") == 200_000
|
|
assert vram._lookup_model_context("unknown-model") == 0
|
|
|
|
|
|
def test_parse_token_value() -> None:
|
|
assert vram._parse_token_value("128k") == 128_000
|
|
assert vram._parse_token_value("128K") == 128_000
|
|
assert vram._parse_token_value("128000") == 128_000
|
|
assert vram._parse_token_value("nonsense") is None
|
|
|
|
|
|
def test_extract_value() -> None:
|
|
assert vram._extract_value("- **Model**: gpt-4o") == "gpt-4o"
|
|
assert vram._extract_value("model = gpt-4o # comment") == "gpt-4o"
|
|
assert vram._extract_value("- **Model**: auto") == "auto"
|
|
|
|
|
|
def test_parse_config_model(tmp_path: Path) -> None:
|
|
config = tmp_path / "config.md"
|
|
config.write_text(
|
|
"## VRAM Configuration\n"
|
|
"- **Auto-detect**: Yes\n"
|
|
"- **Target context**: 16k tokens\n"
|
|
"- **Headroom**: 25%\n"
|
|
"\n"
|
|
"## Model Configuration\n"
|
|
"- **Model**: gpt-4o\n"
|
|
"- **Override context window**: 128k\n"
|
|
)
|
|
model, override = vram._parse_config_model(config)
|
|
assert model == "gpt-4o"
|
|
assert override == 128_000
|
|
|
|
|
|
def test_parse_config_model_skips_code_blocks(tmp_path: Path) -> None:
|
|
"""Example code blocks should not be parsed as live config."""
|
|
config = tmp_path / "config.md"
|
|
config.write_text(
|
|
"## VRAM Configuration\n"
|
|
"- **Headroom**: 25%\n"
|
|
"```\n"
|
|
"- **Headroom**: 99%\n"
|
|
"```\n"
|
|
)
|
|
parsed = vram.parse_vram_config(config)
|
|
assert parsed["headroom_pct"] == 25
|
|
|
|
|
|
def test_parse_vram_config_manual(tmp_path: Path) -> None:
|
|
config = tmp_path / "config.md"
|
|
config.write_text(
|
|
"## VRAM Configuration\n"
|
|
"- **Auto-detect**: No\n"
|
|
"- **Target context**: 8k\n"
|
|
"- **Headroom**: 30%\n"
|
|
"- **Max peak context per sub-task**: 5.6k\n"
|
|
)
|
|
parsed = vram.parse_vram_config(config)
|
|
assert parsed["auto_detect"] is False
|
|
assert parsed["target_context_kb"] == 8000
|
|
assert parsed["headroom_pct"] == 30
|
|
assert parsed["max_peak_kb"] == 5600
|
|
|
|
|
|
def test_recommend_context_api_model() -> None:
|
|
config = {"auto_detect": True, "headroom_pct": 25, "target_context_kb": 0, "max_peak_kb": 0}
|
|
_, recommended_kb, max_peak_kb = vram.recommend_context(
|
|
gpu_vram_gb=0,
|
|
ram_gb=16,
|
|
model_context_kb=128_000,
|
|
overhead_tokens=4000,
|
|
config=config,
|
|
)
|
|
assert recommended_kb > 0
|
|
assert max_peak_kb > 0
|
|
assert recommended_kb <= 128_000 * 0.75 # after headroom
|
|
|
|
|
|
def test_recommend_context_manual_mode() -> None:
|
|
config = {"auto_detect": False, "headroom_pct": 30, "target_context_kb": 8000, "max_peak_kb": 0}
|
|
headroom, recommended_kb, max_peak_kb = vram.recommend_context(
|
|
gpu_vram_gb=0,
|
|
ram_gb=16,
|
|
model_context_kb=0,
|
|
overhead_tokens=0,
|
|
config=config,
|
|
)
|
|
assert headroom == 30
|
|
assert recommended_kb == 8000
|
|
assert max_peak_kb == 5600
|
|
|
|
|
|
def test_extract_model_from_file_respects_10kb_limit(tmp_path: Path) -> None:
|
|
"""Only the first 10KB of an API config file is scanned."""
|
|
env_file = tmp_path / ".env"
|
|
# Put the model name far beyond 10KB.
|
|
env_file.write_text("x" * 11_000 + "\nMODEL=far-away-model\n")
|
|
model = vram._extract_model_from_file(env_file)
|
|
assert model is None
|
|
|
|
|
|
# --- Cross-platform fixtures (module-level multiline string constants) ---
|
|
|
|
MEMINFO_LINUX_16GB = "MemTotal: 16384000 kB\nMemAvailable: 8192000 kB\n"
|
|
MEMINFO_LINUX_32GB = "MemTotal: 32768000 kB\nMemAvailable: 16384000 kB\n"
|
|
|
|
APPLE_M2_PROFILER = """\
|
|
Graphics/Displays:
|
|
Apple M2:
|
|
Chipset Model: Apple M2
|
|
Type: GPU
|
|
Bus: Built-In
|
|
Total Number of Cores: 10
|
|
Vendor: Apple (0x106b)
|
|
Metal: Supported, version 2
|
|
"""
|
|
|
|
WMIC_VIDEOCONTROLLER = """\
|
|
AdapterRAM=8589934592
|
|
Name=NVIDIA GeForce RTX 3060
|
|
"""
|
|
|
|
WMIC_COMPUTERSYSTEM = """\
|
|
TotalPhysicalMemory=34359738368
|
|
"""
|
|
|
|
OLLAMA_LIST = """\
|
|
NAME ID SIZE MODIFIED
|
|
llama-3.1-8b abc 4.7GB 2 days ago
|
|
"""
|
|
|
|
|
|
class _FakeResult:
|
|
"""Minimal stand-in for subprocess.CompletedProcess."""
|
|
|
|
def __init__(self, stdout: str = "", returncode: int = 0) -> None:
|
|
self.stdout = stdout
|
|
self.returncode = returncode
|
|
|
|
|
|
def _make_fake_run(responses: dict[str, str]):
|
|
"""Build a subprocess.run fake that maps cmd[0] (or joined PS string) to stdout."""
|
|
|
|
def fake_run(cmd, *args, **kwargs):
|
|
name = cmd[0]
|
|
if name == "powershell":
|
|
joined = cmd[-1] if len(cmd) > 4 else ""
|
|
for key, out in responses.items():
|
|
if key in joined:
|
|
return _FakeResult(out)
|
|
return _FakeResult("", 1)
|
|
return _FakeResult(responses.get(name, ""))
|
|
|
|
return fake_run
|
|
|
|
|
|
def _patch_meminfo(monkeypatch, text: str) -> None:
|
|
"""Patch Path.exists/read_text to serve `text` for /proc/meminfo only."""
|
|
orig_exists = Path.exists
|
|
orig_read_text = Path.read_text
|
|
|
|
def fake_exists(self):
|
|
if self.as_posix() == "/proc/meminfo":
|
|
return True
|
|
return orig_exists(self)
|
|
|
|
def fake_read_text(self, *args, **kwargs):
|
|
if self.as_posix() == "/proc/meminfo":
|
|
return text
|
|
return orig_read_text(self, *args, **kwargs)
|
|
|
|
monkeypatch.setattr(Path, "exists", fake_exists)
|
|
monkeypatch.setattr(Path, "read_text", fake_read_text)
|
|
|
|
|
|
def test_detect_ram_linux(monkeypatch) -> None:
|
|
monkeypatch.setattr(vram.platform, "system", lambda: "Linux")
|
|
_patch_meminfo(monkeypatch, MEMINFO_LINUX_16GB)
|
|
total, available = vram.detect_ram()
|
|
assert (total, available) == (16_384_000, 8_192_000)
|
|
|
|
|
|
def test_detect_ram_macos(monkeypatch) -> None:
|
|
monkeypatch.setattr(vram.platform, "system", lambda: "Darwin")
|
|
monkeypatch.setattr(vram.shutil, "which", lambda name: name)
|
|
monkeypatch.setattr(
|
|
vram.subprocess, "run", _make_fake_run({"sysctl": "34359738368"})
|
|
)
|
|
total, available = vram.detect_ram()
|
|
assert total == 33_554_432
|
|
assert available == total
|
|
|
|
|
|
def test_detect_ram_windows(monkeypatch) -> None:
|
|
monkeypatch.setattr(vram.platform, "system", lambda: "Windows")
|
|
monkeypatch.setattr(vram.shutil, "which", lambda name: name)
|
|
monkeypatch.setattr(
|
|
vram.subprocess, "run", _make_fake_run({"wmic": WMIC_COMPUTERSYSTEM})
|
|
)
|
|
total, available = vram.detect_ram()
|
|
assert total == 33_554_432
|
|
|
|
|
|
def test_detect_gpu_vram_nvidia_linux(monkeypatch) -> None:
|
|
monkeypatch.setattr(vram.platform, "system", lambda: "Linux")
|
|
monkeypatch.setattr(vram.shutil, "which", lambda name: name)
|
|
monkeypatch.setattr(
|
|
vram.subprocess, "run", _make_fake_run({"nvidia-smi": "24576\n"})
|
|
)
|
|
total, per_gpu, num = vram.detect_gpu_vram()
|
|
assert (total, per_gpu, num) == (25_165_824, 25_165_824, 1)
|
|
|
|
|
|
def test_detect_gpu_vram_apple_silicon(monkeypatch) -> None:
|
|
monkeypatch.setattr(vram.platform, "system", lambda: "Darwin")
|
|
monkeypatch.setattr(vram.shutil, "which", lambda name: name)
|
|
monkeypatch.setattr(
|
|
vram.subprocess,
|
|
"run",
|
|
_make_fake_run(
|
|
{"system_profiler": APPLE_M2_PROFILER, "sysctl": "17179869184"}
|
|
),
|
|
)
|
|
total, per_gpu, num = vram.detect_gpu_vram()
|
|
assert total > 0
|
|
assert num >= 1
|
|
|
|
|
|
def test_detect_gpu_vram_windows_wmic(monkeypatch) -> None:
|
|
monkeypatch.setattr(vram.platform, "system", lambda: "Windows")
|
|
monkeypatch.setattr(vram.shutil, "which", lambda name: name)
|
|
monkeypatch.setattr(
|
|
vram.subprocess, "run", _make_fake_run({"wmic": WMIC_VIDEOCONTROLLER})
|
|
)
|
|
total, per_gpu, num = vram.detect_gpu_vram()
|
|
assert total == 8_388_608
|
|
|
|
|
|
def test_lookup_model_context_unknown_returns_zero() -> None:
|
|
assert vram._lookup_model_context("completely-unknown-model") == 0
|
|
|
|
|
|
def test_lookup_model_context_prefix_match() -> None:
|
|
assert vram._lookup_model_context("deepseek-r1:7b") == 64_000
|
|
assert vram._lookup_model_context("llama-3.1-8b-instruct") == 128_000
|
|
|
|
|
|
def test_lookup_model_context_no_false_prefix_match() -> None:
|
|
"""phi-4-mini should NOT match phi-4 (different model, wrong context)."""
|
|
assert vram._lookup_model_context("phi-4-mini-instruct") == 0
|
|
assert vram._lookup_model_context("gpt-4o-foo-unknown") == 0
|
|
|
|
|
|
def test_detect_model_context_ollama_probe(monkeypatch, tmp_path: Path) -> None:
|
|
monkeypatch.setattr(vram.Path, "home", lambda: tmp_path)
|
|
monkeypatch.setattr(
|
|
vram.shutil, "which", lambda name: name if name == "ollama" else None
|
|
)
|
|
monkeypatch.setattr(
|
|
vram.subprocess, "run", _make_fake_run({"ollama": OLLAMA_LIST})
|
|
)
|
|
context = vram.detect_model_context(None, tmp_path)
|
|
assert context == 128_000
|
|
|
|
|
|
def test_run_command_windows_powershell_wrapper(monkeypatch) -> None:
|
|
monkeypatch.setattr(vram.platform, "system", lambda: "Windows")
|
|
monkeypatch.setattr(vram.shutil, "which", lambda name: name)
|
|
captured: dict = {}
|
|
|
|
def fake_run(cmd, *args, **kwargs):
|
|
captured["cmd"] = cmd
|
|
return _FakeResult("ok", 0)
|
|
|
|
monkeypatch.setattr(vram.subprocess, "run", fake_run)
|
|
result = vram.run_command(
|
|
["Get-CimInstance", "Win32_VideoController", "-Property", "AdapterRAM"]
|
|
)
|
|
assert result == "ok"
|
|
cmd = captured["cmd"]
|
|
assert cmd[0] == "powershell"
|
|
assert "-NoProfile" in cmd
|
|
assert "-NoLogo" in cmd
|
|
assert "-Command" in cmd
|
|
assert "Get-CimInstance" in cmd[-1]
|
|
assert "Win32_VideoController" in cmd[-1]
|
|
|
|
|
|
def test_detect_ram_linux_regression(monkeypatch) -> None:
|
|
"""LOCKED regression guard — exact match on (total_kb, available_kb)."""
|
|
monkeypatch.setattr(vram.platform, "system", lambda: "Linux")
|
|
_patch_meminfo(monkeypatch, MEMINFO_LINUX_32GB)
|
|
total, available = vram._detect_ram_linux()
|
|
assert (total, available) == (32_768_000, 16_384_000)
|