AnyImage/tests/tools/models/test_cli.py

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_pisa_export_validates_all_outputs_before_replacing(tmp_path, valid):
    module = load_sync_models()
    digest = module.hashlib.sha256(b"new").hexdigest()
    models = tuple(module.ModelFile("PISA_SR", "pisa-sr", 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(".pisa-sr-*"))


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", "SAFMN_REAL_X4"},
        )
        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()