Files
AnyImage/tests/test_install_env.py
nan c5b77c5c85 Add ONNX model inference implementations for MoGe-2, BEN2, and DA3
- Introduced `onnx_moge2.py` for MoGe-2 ONNX inference with NumPy post-processing.
- Added `onnx_ben2.py` for BEN2 ONNX inference with Pillow and NumPy preprocessing.
- Created `onnx_da3.py` for Depth Anything 3 inference with ONNX Runtime and NumPy.
- Implemented shared ONNX Runtime session helpers in `onnx_runtime.py`.
- Developed unit tests for server functionalities in `test_server.py`, including model caching and job management.
- Enhanced job server to handle frame extraction without requiring models.
2026-08-17 10:37:55 +08:00

122 lines
4.3 KiB
Python

import importlib.util
import json
import sys
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
PROJECT_ROOT = Path(__file__).resolve().parents[1]
def load_environment():
path = PROJECT_ROOT / "src" / "snapform" / "blendjob" / "environment.py"
specification = importlib.util.spec_from_file_location(
"blendjob_environment_test",
path,
)
module = importlib.util.module_from_spec(specification)
specification.loader.exec_module(module)
return module
class EnvironmentTest(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.environment = load_environment()
def config(self):
return {
"python": "3.12",
"packages": ["numpy==2.4.2"],
"platform_packages": {
"windows": ["onnxruntime-directml==1.24.4"],
"default": ["onnxruntime==1.24.4"],
},
}
def test_normalizes_public_dictionary(self):
self.assertEqual(
self.environment.normalized_environment(self.config()),
self.config(),
)
def test_requires_python_and_list_packages(self):
with self.assertRaises(ValueError):
self.environment.normalized_environment({})
with self.assertRaises(TypeError):
self.environment.normalized_environment(
{"python": "3.12", "packages": "numpy"}
)
def test_environment_digest_tracks_configuration(self):
first = self.environment.environment_digest(self.config())
reordered = {
"platform_packages": self.config()["platform_packages"],
"packages": ["numpy==2.4.2"],
"python": "3.12",
}
self.assertEqual(first, self.environment.environment_digest(reordered))
reordered["packages"].append("pillow==12.1.1")
self.assertNotEqual(first, self.environment.environment_digest(reordered))
def test_selects_platform_packages_and_server_packages(self):
with patch.object(self.environment.sys, "platform", "win32"):
packages = self.environment.resolved_packages(self.config())
self.assertIn("fastapi==0.139.2", packages)
self.assertIn("uvicorn==0.51.0", packages)
self.assertIn("numpy==2.4.2", packages)
self.assertIn("onnxruntime-directml==1.24.4", packages)
self.assertNotIn("onnxruntime==1.24.4", packages)
def test_uses_default_platform_packages(self):
with patch.object(self.environment.sys, "platform", "linux"):
packages = self.environment.resolved_packages(self.config())
self.assertIn("onnxruntime==1.24.4", packages)
def test_install_writes_manifest_from_environment_hash(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
uv = root / "tools" / "uv"
uv.parent.mkdir()
uv.touch()
status = root / "environment-status.json"
with patch.object(
self.environment,
"install_uv",
return_value=uv,
), patch.object(
self.environment,
"run_command",
), patch.object(
self.environment,
"install_packages",
) as install_packages:
self.environment.install(root, self.config(), status, stages=3)
manifest = json.loads((root / "manifest.json").read_text("utf-8"))
final_status = json.loads(status.read_text("utf-8"))
self.assertEqual(
manifest["environment_hash"],
self.environment.environment_digest(self.config()),
)
self.assertIn("numpy==2.4.2", manifest["packages"])
self.assertEqual(final_status["stage"], 1)
self.assertEqual(final_status["stages"], 3)
self.assertEqual(final_status["progress"], 1.0)
self.assertTrue(install_packages.called)
def test_environment_python_is_inside_runtime_storage(self):
path = self.environment.environment_python("storage")
expected = (
Path("storage/.venv/Scripts/python.exe")
if sys.platform == "win32"
else Path("storage/.venv/bin/python")
)
self.assertEqual(path, expected)
if __name__ == "__main__":
unittest.main()