Initial public release

tests/anyimage/server/test_model_download_jobs.py

@@ -0,0 +1,154 @@
+import os
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+from types import ModuleType, SimpleNamespace
+from unittest.mock import patch
+from tests.anyimage.server.support import load_server_module
+
+
+class ModelDownloadJobsTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.download = load_server_module("model_download")
+
+
+    def test_download_model_uses_actual_request_without_preflight(self):
+        calls = []
+        fake_hugging_face = ModuleType("huggingface_hub")
+        fake_hugging_face.snapshot_download = lambda **options: calls.append(
+            options
+        )
+        args = SimpleNamespace(
+            repository="Ruicheng/moge-2-vits-normal-onnx",
+            model_dir="model",
+            progress=lambda *_args, **_kwargs: None,
+            allow_pattern=[],
+        )
+
+        with tempfile.TemporaryDirectory() as directory, patch.dict(
+            sys.modules,
+            {"huggingface_hub": fake_hugging_face},
+        ), patch.object(
+            self.download,
+            "create_download_progress_class",
+            return_value=object,
+        ), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
+            self.download,
+            "_verify_model_complete",
+            return_value=None,
+        ):
+            args.model_dir = str(Path(directory) / "models" / "MoGe2")
+            self.download.download_model(args)
+
+        self.assertEqual(len(calls), 1)
+
+
+    def test_download_model_refuses_success_when_files_missing(self):
+        fake_hugging_face = ModuleType("huggingface_hub")
+        fake_hugging_face.snapshot_download = lambda **options: None
+        args = SimpleNamespace(
+            repository="PramaLLC/BEN2",
+            model_dir="model",
+            progress=lambda *_args, **_kwargs: None,
+            allow_pattern=[],
+        )
+
+        with tempfile.TemporaryDirectory() as directory, patch.dict(
+            sys.modules,
+            {"huggingface_hub": fake_hugging_face},
+        ), patch.object(
+            self.download,
+            "create_download_progress_class",
+            return_value=object,
+        ), patch.dict(os.environ, {"HF_ENDPOINT": ""}):
+            args.model_dir = str(Path(directory) / "model")
+            with self.assertRaisesRegex(
+                RuntimeError,
+                "All download sources failed",
+            ):
+                self.download.download_model(args)
+
+
+    def test_download_model_reports_all_failed_sources(self):
+        args = SimpleNamespace(
+            repository="Ruicheng/moge-2-vits-normal-onnx",
+            model_dir="model",
+            progress=lambda *_args, **_kwargs: None,
+            allow_pattern=[],
+            r2_directory="moge-2-vits-normal-onnx",
+            r2_file=["config.json|100|abc"],
+        )
+
+        with tempfile.TemporaryDirectory() as directory, patch.object(
+            self.download,
+            "download_from_hf",
+            side_effect=OSError("direct down"),
+        ), patch.object(
+            self.download,
+            "download_from_r2",
+            side_effect=OSError("mirror down"),
+        ):
+            args.model_dir = str(Path(directory) / "model")
+            with self.assertRaisesRegex(RuntimeError, "models.omoolab.xyz"):
+                self.download.download_model(args)
+
+
+    def test_download_model_prefers_direct_hf_then_r2(self):
+        args = SimpleNamespace(
+            repository="Ruicheng/moge-2-vits-normal-onnx",
+            model_dir="model",
+            progress=lambda *_args, **_kwargs: None,
+            allow_pattern=["config.json"],
+            r2_directory="moge-2-vits-normal-onnx",
+            r2_file=["config.json|100|abc"],
+        )
+        order = []
+
+        def fake_download_from_hf(*_args, **_kwargs):
+            order.append("hf")
+            raise OSError("hf down")
+
+        def fake_download_from_r2(*_args, **_kwargs):
+            order.append("r2")
+
+        with tempfile.TemporaryDirectory() as directory, patch.object(
+            self.download,
+            "download_from_hf",
+            side_effect=fake_download_from_hf,
+        ), patch.object(
+            self.download,
+            "download_from_r2",
+            side_effect=fake_download_from_r2,
+        ), patch.object(
+            self.download,
+            "_verify_model_complete",
+            return_value=None,
+        ):
+            args.model_dir = str(Path(directory) / "models" / "MoGe2")
+            self.download.download_model(args)
+
+        self.assertEqual(order, ["hf", "r2"])
+
+
+    def test_download_from_r2_downloads_declared_files(self):
+        args = SimpleNamespace(
+            repository="Ruicheng/moge-2-vits-normal-onnx",
+            model_dir="model",
+            progress=lambda *_args, **_kwargs: None,
+            allow_pattern=["config.json"],
+            r2_directory="moge-2-vits-normal-onnx",
+            r2_file=["config.json|100|abc"],
+        )
+        downloaded = []
+
+        with tempfile.TemporaryDirectory() as directory, patch.object(
+            self.download,
+            "_download_file_from_url",
+            side_effect=lambda *args, **kwargs: downloaded.append(args[1].name),
+        ):
+            args.model_dir = str(Path(directory) / "model")
+            self.download.download_from_r2(args, args.model_dir)
+
+        self.assertEqual(downloaded, ["config.json"])