import tempfile
from pathlib import Path
from unittest import TestCase
from app.model_manager import LocalModelManager, validate_model_file
class ModelManagerTests(TestCase):
def test_validate_model_checks_magic_and_size(self) -> None:
with tempfile.TemporaryDirectory() as tmpdir:
model_path = Path(tmpdir) / "bad.gguf"
model_path.write_bytes(b"GGXX" + b"\x00" * 1024)
with self.assertRaises(ValueError):
validate_model_file(model_path, min_size_mb=0.0001, expected_sha256=None)
def test_validate_model_checks_sha256(self) -> None:
with tempfile.TemporaryDirectory() as tmpdir:
model_path = Path(tmpdir) / "test.gguf"
# Create valid GGUF with known content
model_path.write_bytes(b"GGUF" + b"\x00" * (2 * 1024 * 1024))
# Should pass with correct SHA
import hashlib
sha = hashlib.sha256()
with model_path.open("rb") as f:
sha.update(f.read())
correct_sha = sha.hexdigest()
validate_model_file(model_path, min_size_mb=1.0, expected_sha256=correct_sha)
# Should fail with wrong SHA
with self.assertRaisesRegex(ValueError, "SHA256 mismatch"):
validate_model_file(model_path, min_size_mb=1.0, expected_sha256="badbadbad")
def test_self_test_normalization_handles_quotes_and_punctuation(self) -> None:
manager = LocalModelManager(model_path=Path(__file__))
# Test various formats that should all normalize to OK
test_cases = [
' "OK.?!\n" ',
"'OK!'",
' OK. ',
'"OK"',
'OK.!?.',
]
for test_input in test_cases:
normalized = manager._normalize_self_test_output(test_input)
self.assertEqual(normalized, "OK", f"Failed for input: {test_input!r}")
