Initial public release

tests/tools/models/test_cli.py

@@ -0,0 +1,187 @@
+from tests.support.paths import PROJECT_ROOT
+import importlib.util
+import sys
+import unittest
+from unittest.mock import patch
+from dataclasses import replace
+from pathlib import Path
+
+import pytest
+
+
+@pytest.mark.parametrize("command", ["prepare", "sync"])
+def test_model_commands_prepare_all_before_upload(command):
+    module = load_sync_models()
+    events = []
+    with (
+        patch("sys.argv", ["model", command]),
+        patch.object(module, "prepare_model", side_effect=lambda model: events.append(model)),
+        patch.object(module, "upload_models", side_effect=lambda: events.append("upload")),
+    ):
+        module.main()
+    assert events == [*module.MODEL_FILES, *(["upload"] if command == "sync" else [])]
+
+
+def test_prepare_failure_prevents_upload():
+    module = load_sync_models()
+    with (
+        patch("sys.argv", ["model", "sync"]),
+        patch.object(module, "prepare_model", side_effect=RuntimeError("invalid model")),
+        patch.object(module, "upload_models") as upload,
+    ):
+        with pytest.raises(RuntimeError, match="invalid model"):
+            module.main()
+    upload.assert_not_called()
+
+
+def test_valid_model_cache_skips_all_preparation():
+    module = load_sync_models()
+    with (
+        patch.object(module, "verify_file", return_value=True),
+        patch.object(module, "build_model") as build,
+        patch.object(module, "download_file") as download,
+        patch.object(module.subprocess, "run") as run,
+    ):
+        for model in module.MODEL_FILES:
+            module.prepare_model(model)
+    build.assert_not_called()
+    download.assert_not_called()
+    run.assert_not_called()
+
+
+@pytest.mark.parametrize("valid", [True, False])
+def test_isolated_export_validates_all_outputs_before_replacing(tmp_path, valid):
+    module = load_sync_models()
+    digest = module.hashlib.sha256(b"new").hexdigest()
+    models = tuple(module.ModelFile("HAT_GAN_X4_SHARPER", "hat-gan-x4-sharper", name, 3, digest) for name in ("a.onnx", "b.onnx"))
+
+    def export(command, **kwargs):
+        destination = Path(command[-1])
+        (destination / "a.onnx").write_bytes(b"new")
+        (destination / "b.onnx").write_bytes(b"new" if valid else b"bad")
+
+    with patch.object(module, "MODELS_DIR", tmp_path), patch.object(module, "MODEL_FILES", models), patch.object(module.subprocess, "run", side_effect=export):
+        models[1].destination.parent.mkdir()
+        models[1].destination.write_bytes(b"old")
+        if valid:
+            module.prepare_model(models[0])
+            assert all(model.destination.read_bytes() == b"new" for model in models)
+        else:
+            with pytest.raises(RuntimeError, match="SHA-256"):
+                module.prepare_model(models[0])
+            assert not models[0].destination.exists()
+            assert models[1].destination.read_bytes() == b"old"
+        assert not list(tmp_path.glob(".hat-gan-x4-sharper-*"))
+
+
+def test_local_export_failure_cleans_temporary_output(tmp_path):
+    module = load_sync_models()
+    model = replace(next(model for model in module.MODEL_FILES if model.build), folder="test")
+
+    def export(exporter, source, destination):
+        destination.write_bytes(b"partial")
+        raise RuntimeError("export failed")
+
+    with (
+        patch.object(module, "MODELS_DIR", tmp_path),
+        patch.object(module, "PROJECT_ROOT", tmp_path),
+        patch.object(module, "prepare_source", return_value=tmp_path / "weights.pth"),
+        patch("tools.models.exports.export_model", side_effect=export),
+    ):
+        model.destination.parent.mkdir()
+        model.destination.write_bytes(b"original")
+        with pytest.raises(RuntimeError, match="export failed"):
+            module.build_model(model)
+        assert model.destination.read_bytes() == b"original"
+        assert not model.destination.with_name(model.filename + ".part").exists()
+
+
+def load_sync_models():
+    path = PROJECT_ROOT / "tools" / "models" / "cli.py"
+    spec = importlib.util.spec_from_file_location("sync_models", path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+class SyncModelsTests(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.sync = load_sync_models()
+
+    def test_r2_is_only_an_upload_destination(self):
+        for model in self.sync.MODEL_FILES:
+            self.assertNotIn("models.omoolab.xyz", model.source_url)
+            if model.build:
+                self.assertFalse(model.source_url)
+
+    def test_local_exports_use_official_pytorch_weights(self):
+        builds = self.sync.MODEL_BUILDS
+        self.assertEqual(
+            set(builds),
+            {"REALESRGAN_GENERAL_WDN_X4V3"},
+        )
+        for build in builds.values():
+            self.assertTrue(build.source_url.startswith("https://github.com/"))
+            self.assertTrue(build.source_filename.endswith(".pth"))
+            self.assertEqual(len(build.source_sha256), 64)
+
+    def test_prepare_model_builds_missing_local_export(self):
+        model = next(model for model in self.sync.MODEL_FILES if model.build)
+        with (
+            patch.object(self.sync, "verify_file", return_value=False),
+            patch.object(self.sync, "build_model") as build,
+            patch.object(self.sync, "download_file") as download,
+        ):
+            self.sync.prepare_model(model)
+        build.assert_called_once_with(model)
+        download.assert_not_called()
+
+    def test_prepare_model_downloads_upstream_onnx(self):
+        model = next(
+            model
+            for model in self.sync.MODEL_FILES
+            if model.source_url and not model.build
+        )
+        with (
+            patch.object(self.sync, "verify_file", return_value=False),
+            patch.object(self.sync, "download_file") as download,
+        ):
+            self.sync.prepare_model(model)
+        download.assert_called_once_with(
+            model.source_url,
+            model.destination,
+            model.size,
+            model.sha256,
+        )
+
+    def test_upload_copies_declared_files_without_remote_deletion(self):
+        model = self.sync.ModelFile(
+            key="TEST",
+            folder="folder",
+            filename="model.onnx",
+            size=1,
+            sha256="hash",
+        )
+        with (
+            patch.object(self.sync.shutil, "which", return_value="rclone"),
+            patch.object(self.sync.subprocess, "run") as run,
+        ):
+            self.sync.upload_models((model,))
+        command = run.call_args.args[0]
+        self.assertEqual(command[1], "copyto")
+        self.assertNotIn("sync", command)
+        self.assertEqual(command[-2], model.remote)
+
+    def test_verify_file_checks_size_and_sha256(self):
+        path = PROJECT_ROOT / "tools" / "models" / "cli.py"
+        size = path.stat().st_size
+        digest = self.sync.file_sha256(path)
+        self.assertTrue(self.sync.verify_file(path, size, digest))
+        with self.assertRaisesRegex(RuntimeError, "unexpected size"):
+            self.sync.verify_file(path, size - 1, digest)
+
+
+if __name__ == "__main__":
+    unittest.main()