188 lines
7.0 KiB
Python
188 lines
7.0 KiB
Python
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()
|