Files
AnyImage/tests/test_model_download_jobs.py
nan 4c902ed7db feat(node_assets): add image plane geometry nodes and validation scripts
- Implemented `image_plane.py` to build the Image Plane Geometry Nodes group with customizable parameters for geometry, subdivide, and thickness.
- Created `image_plane_validation.py` to validate the functionality of the image plane, ensuring correct geometry generation and material assignment.
- Developed `material_layers.py` to construct shader node groups for image materials, including depth and normal handling.
- Authored `plane_layouts.py` to define node arrangements for image geometry assets.
- Introduced `shadeless.py` to create a shadeless shader node group that displays color without scene lighting while preserving shadows and reflections.
2026-09-08 23:24:27 +08:00

155 lines
5.2 KiB
Python

import os
import sys
import tempfile
import unittest
from pathlib import Path
from types import ModuleType, SimpleNamespace
from unittest.mock import patch
from server_test_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"])