| from __future__ import annotations |
|
|
| import tempfile |
| import unittest |
| from pathlib import Path |
|
|
| from models.model_catalog import load_model_catalog |
| from training.export import ( |
| ToolStatus, |
| build_export_plan, |
| detect_llama_cpp_tools, |
| list_exported_files, |
| ) |
| from ui.export_tab import export_download_paths |
|
|
|
|
| class ExportPlanTest(unittest.TestCase): |
| def test_detects_llama_cpp_tools_without_requiring_install(self) -> None: |
| statuses = detect_llama_cpp_tools(which_func=lambda name: f"C:/tools/{name}.exe") |
|
|
| self.assertEqual( |
| [status.name for status in statuses], |
| ["llama-server", "llama-cli", "llama-quantize"], |
| ) |
| self.assertTrue(all(status.available for status in statuses)) |
|
|
| def test_builds_non_executing_official_gguf_plan(self) -> None: |
| catalog = load_model_catalog("config/models.yaml") |
| tools = [ToolStatus("llama-quantize", False, "")] |
|
|
| plan = build_export_plan(catalog["minicpm_v46"], "Q4_K_M", "exports", tools) |
| data = plan.as_dict() |
|
|
| self.assertEqual(data["official_gguf_repo"], "openbmb/MiniCPM-V-4.6-gguf") |
| self.assertIn("huggingface-cli", data["download_command"]) |
| self.assertIn("llama-quantize", data["quantize_command"]) |
| self.assertFalse(data["executes_commands"]) |
| self.assertFalse(data["startup_downloads"]) |
|
|
| def test_rejects_unknown_quantization(self) -> None: |
| catalog = load_model_catalog("config/models.yaml") |
|
|
| with self.assertRaises(ValueError): |
| build_export_plan(catalog["minicpm5_1b"], "Q2_UNKNOWN") |
|
|
| def test_lists_exported_files(self) -> None: |
| with tempfile.TemporaryDirectory() as tmp: |
| output = Path(tmp) |
| file_path = output / "model.gguf" |
| file_path.write_bytes(b"gguf") |
|
|
| rows = list_exported_files(output) |
|
|
| self.assertEqual(rows, [[str(file_path), "4"]]) |
|
|
| def test_export_download_paths_prioritizes_existing_planned_files(self) -> None: |
| catalog = load_model_catalog("config/models.yaml") |
| tools = [ToolStatus("llama-quantize", False, "")] |
|
|
| with tempfile.TemporaryDirectory() as tmp: |
| root = Path(tmp) |
| model_dir = root / "minicpm_v46" |
| model_dir.mkdir() |
| extra_file = root / "notes.txt" |
| extra_file.write_text("manual note", encoding="utf-8") |
|
|
| plan = build_export_plan(catalog["minicpm_v46"], "Q4_K_M", root, tools) |
| planned_file = model_dir / plan.official_gguf_file |
| planned_file.write_bytes(b"planned gguf") |
|
|
| rows = list_exported_files(root) |
| paths = export_download_paths(plan.as_dict(), rows) |
|
|
| self.assertEqual(paths[0], str(planned_file)) |
| self.assertIn(str(extra_file), paths) |
|
|
| def test_export_download_paths_omits_missing_planned_files(self) -> None: |
| catalog = load_model_catalog("config/models.yaml") |
| tools = [ToolStatus("llama-quantize", False, "")] |
|
|
| with tempfile.TemporaryDirectory() as tmp: |
| root = Path(tmp) |
| extra_file = root / "already-exported.gguf" |
| extra_file.write_bytes(b"existing") |
|
|
| plan = build_export_plan(catalog["minicpm_v46"], "Q4_K_M", root, tools) |
| paths = export_download_paths(plan.as_dict(), list_exported_files(root)) |
|
|
| self.assertEqual(paths, [str(extra_file)]) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
