"""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