import importlib.util import io import hashlib import json import os import shutil import subprocess import tarfile import tempfile import time import unittest import zipfile from pathlib import Path from types import SimpleNamespace from unittest.mock import Mock, patch PROJECT_ROOT = Path(__file__).resolve().parents[1] def load_install_env(): specification = importlib.util.spec_from_file_location( "snapform_install_env", PROJECT_ROOT / "src" / "snapform" / "install_env.py", ) module = importlib.util.module_from_spec(specification) specification.loader.exec_module(module) return module class InstallEnvTest(unittest.TestCase): def create_bundled_vendor(self, root): vendor = root / "vendor" vendor.mkdir() (vendor / "depth-anything-3.tar.gz").write_bytes(b"da3") (vendor / "ben2.tar.gz").write_bytes(b"ben2") (vendor / "utils3d.tar.gz").write_bytes(b"utils3d") (vendor / "moge.tar.gz").write_bytes(b"moge") return vendor def test_installs_onnx_runtime_without_training_frameworks(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) venv_dir = root / ".venv" venv_dir.mkdir() self.create_bundled_vendor(root) uv_path = root / "uv" uv_path.write_bytes(b"uv") args = SimpleNamespace( uv=str(uv_path), venv_dir=str(venv_dir), models_dir=str(root / "models"), torch_variant="cu128", da3_commit="test-commit", ben2_commit="test-ben2-commit", status=str(root / "status.json"), ) with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object( bootstrap, "run", ) as mocked_run, patch.object( bootstrap, "run_install", ) as mocked_run_install, patch.object( bootstrap, "run_verify", return_value={}, ): bootstrap.install(args) commands = [call.args[0] for call in mocked_run.call_args_list] commands += [ call.args[4] for call in mocked_run_install.call_args_list ] service_command = next( command for command in commands if f"fastapi=={bootstrap.FASTAPI_VERSION}" in command ) self.assertIn(f"uvicorn=={bootstrap.UVICORN_VERSION}", service_command) self.assertIn( f"onnxruntime-directml=={bootstrap.ONNXRUNTIME_VERSION}", service_command, ) for forbidden in ("torch", "torchvision", "xformers", "ben2 @", "moge @"): self.assertFalse( any(forbidden in str(value) for value in service_command), forbidden, ) check_command = bootstrap.verify_command( bootstrap.venv_python(venv_dir) ) self.assertFalse(any("torch" in str(value) for value in check_command)) self.assertTrue( any("import fastapi, uvicorn" in str(value) for value in check_command) ) self.assertTrue( any("onnxruntime" in str(value) for value in check_command) ) self.assertFalse( any("import gsplat" in str(value) for value in check_command) ) self.assertTrue( any( "SNAPFORM_DEVICES=" in str(value) for value in check_command ) ) default_downloads = [ command for command in commands if any( repository in command for repository in ( "Ruicheng/moge-2-vits-normal-onnx", "onnx-community/BEN2-ONNX", ) ) ] self.assertEqual(len(default_downloads), 2) self.assertTrue(any("Ruicheng/moge-2-vits-normal-onnx" in map(str, command) for command in default_downloads)) ben2_download = next( command for command in default_downloads if "onnx-community/BEN2-ONNX" in command ) patterns = [ ben2_download[index + 1] for index, value in enumerate(ben2_download) if value == "--allow-pattern" ] self.assertEqual(patterns, ["onnx/model_fp16.onnx"]) self.assertIn("--stage", ben2_download) self.assertIn("--stage-label", ben2_download) self.assertEqual( ben2_download[ben2_download.index("--label") + 1], "BEN2 Base", ) def test_darwin_install_uses_default_onnx_runtime(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) venv_dir = root / ".venv" venv_dir.mkdir() self.create_bundled_vendor(root) uv_path = root / "uv" uv_path.write_bytes(b"uv") args = SimpleNamespace( uv=str(uv_path), venv_dir=str(venv_dir), models_dir=str(root / "models"), torch_variant="cu128", da3_commit="test-commit", ben2_commit="test-ben2-commit", status=str(root / "status.json"), ) with patch.object( bootstrap.sys, "platform", "darwin", ), patch.object( bootstrap.platform, "machine", return_value="arm64", ), patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object( bootstrap, "run", ), patch.object( bootstrap, "run_install", ) as mocked_run_install, patch.object( bootstrap, "run_verify", return_value={}, ): bootstrap.install(args) commands = [ call.args[4] for call in mocked_run_install.call_args_list ] runtime_command = next( command for command in commands if any( f"onnxruntime=={bootstrap.ONNXRUNTIME_VERSION}" in str(value) for value in command ) ) self.assertNotIn("--index-url", runtime_command) self.assertFalse( any( "xformers" in str(value) for command in commands for value in command ) ) def test_intel_mac_uses_portable_onnx_runtime(self): bootstrap = load_install_env() with patch.object( bootstrap.sys, "platform", "darwin", ): stages = bootstrap.install_stages( SimpleNamespace( models_dir="models", status="status.json", ), ".venv", "python", "uv", ) runtime = stages[1][1] self.assertIn( f"onnxruntime=={bootstrap.ONNXRUNTIME_VERSION}", runtime, ) def test_run_stage_writes_stage_progress(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: status_path = Path(directory) / "status.json" with patch.object(bootstrap, "run") as mocked_run: bootstrap.run_stage( str(status_path), 3, 8, "pytorch", ["uv", "pip", "install", "torch"], None, "plain", ) self.assertEqual( mocked_run.call_args.args[0], ["uv", "pip", "install", "torch"], ) status = json.loads(status_path.read_text(encoding="utf-8")) self.assertEqual(status["progress"], 1.0) self.assertEqual(status["stage"], 3) self.assertEqual(status["stages"], 8) self.assertEqual(status["stage_label"], "pytorch") def test_run_stage_writes_progress_range(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: status_path = Path(directory) / "status.json" with patch.object(bootstrap, "run"), patch.object( bootstrap, "write_status", ) as mocked_write: bootstrap.run_stage( str(status_path), 3, 8, "pytorch", ["uv", "pip", "install", "torch"], None, "plain", start=0.2, end=0.4, ) progresses = [ call.args[1] for call in mocked_write.call_args_list ] self.assertEqual(progresses[0], 0.2) self.assertEqual(progresses[-1], 0.4) def test_run_stage_passes_progress_range_to_download_command(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: status_path = Path(directory) / "status.json" command = [ "python", "worker.py", "download", "--repository", "depth-anything/DA3-SMALL", ] with patch.object(bootstrap, "run") as mocked_run: bootstrap.run_stage( str(status_path), 6, 8, "da3 small", command, None, "download", start=0.9, end=0.9333, ) called_command = mocked_run.call_args.args[0] self.assertIn("--progress-start", called_command) self.assertIn("--progress-end", called_command) self.assertEqual( called_command[called_command.index("--progress-start") + 1], "0.9", ) self.assertEqual( called_command[called_command.index("--progress-end") + 1], "0.9333", ) self.assertEqual(called_command[-1], "da3 small") def test_install_resets_progress_for_each_stage(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) venv_dir = root / ".venv" venv_dir.mkdir() self.create_bundled_vendor(root) uv_path = root / "uv" uv_path.write_bytes(b"uv") args = SimpleNamespace( uv=str(uv_path), venv_dir=str(venv_dir), models_dir=str(root / "models"), torch_variant="cu128", da3_commit="test-commit", ben2_commit="test-ben2-commit", status=str(root / "status.json"), ) stage_ranges = [] def fake_run_stage( status_path, stage, stages, label, command, environment, kind, start=0.0, end=1.0, ): stage_ranges.append((stage, start, end)) bootstrap.write_status( status_path, end, "", stage=stage, stages=stages, stage_label=label, ) return [] with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object( bootstrap, "run_stage", side_effect=fake_run_stage, ), patch.object(bootstrap, "run"), patch.object( bootstrap, "run_verify", return_value={}, ): bootstrap.install(args) self.assertEqual(stage_ranges[:2], [(1, 0.0, 0.05), (1, 0.05, 1.0)]) self.assertEqual(stage_ranges[2:], [(2, 0.0, 1.0), (3, 0.0, 1.0)]) status = json.loads( (root / "status.json").read_text(encoding="utf-8") ) self.assertEqual(status["progress"], 1.0) def test_run_install_tracks_package_download_events(self): bootstrap = load_install_env() class FakeProcess: def __init__(self, lines): self.stdout = lines def wait(self): return 0 fake_process = FakeProcess( [ "Resolved 3 packages in 1.0s", "DEBUG Adding transitive dependency for torch==2.9.1: filelock*", "Prepared 3 packages in 2.0s", "Installed 3 packages in 0.5s", "+ torch==2.9.1", "+ numpy==1.26.4", ] ) with patch.object( bootstrap.subprocess, "Popen", return_value=fake_process, ), patch.object(bootstrap, "write_status") as mocked_write: installed = bootstrap.run_install( "status.json", 2, 8, "pytorch deps", ["uv", "pip", "install", "torch"], ) messages = [call.args[2] for call in mocked_write.call_args_list] self.assertEqual( messages, [ "Resolving 3 packages", "Downloading packages 3/3", "Installing packages 3/3", ], ) self.assertEqual(mocked_write.call_args_list[0].args[1], 0.02) self.assertLess(mocked_write.call_args_list[-1].args[1], 0.03) first = mocked_write.call_args_list[0] self.assertEqual(first.kwargs["stage"], 2) self.assertEqual(first.kwargs["stages"], 8) self.assertEqual(first.kwargs["stage_label"], "pytorch deps") self.assertEqual(installed, ["torch==2.9.1", "numpy==1.26.4"]) def test_run_install_smooths_progress_between_events(self): bootstrap = load_install_env() class SlowLines: def __init__(self, lines, delay): self.lines = list(lines) self.delay = delay def __iter__(self): for line in self.lines: time.sleep(self.delay) yield line class FakeProcess: def __init__(self, lines): self.stdout = lines def wait(self): return 0 fake_process = FakeProcess( SlowLines( [ "Resolved 3 packages in 1.0s", "Prepared 3 packages in 2.0s", "Installed 3 packages in 0.5s", ], 0.3, ) ) with patch.object( bootstrap.subprocess, "Popen", return_value=fake_process, ), patch.object(bootstrap, "write_status") as mocked_write: bootstrap.run_install( "status.json", 2, 8, "pytorch deps", ["uv", "pip", "install", "torch"], ) progresses = [call.args[1] for call in mocked_write.call_args_list] self.assertGreater(progresses[-1], 0.02) self.assertTrue( all( progresses[i] <= progresses[i + 1] + 1e-9 for i in range(len(progresses) - 1) ) ) messages = [call.args[2] for call in mocked_write.call_args_list] self.assertIn("Resolving 3 packages", messages) self.assertIn("Downloading packages 3/3", messages) def test_run_install_maps_progress_into_stage_range(self): bootstrap = load_install_env() class FakeProcess: def __init__(self, lines): self.stdout = lines def wait(self): return 0 fake_process = FakeProcess( [ "Resolved 3 packages in 1.0s", "Prepared 3 packages in 2.0s", "Installed 3 packages in 0.5s", ] ) with patch.object( bootstrap.subprocess, "Popen", return_value=fake_process, ), patch.object(bootstrap, "write_status") as mocked_write: bootstrap.run_install( "status.json", 2, 8, "pytorch deps", ["uv", "pip", "install", "torch"], start=0.2, end=0.4, ) first = mocked_write.call_args_list[0] self.assertAlmostEqual(first.args[1], 0.204) self.assertEqual(first.kwargs["stage"], 2) self.assertEqual(first.kwargs["stages"], 8) self.assertEqual(first.kwargs["stage_label"], "pytorch deps") def test_run_install_heartbeat_reports_package_count(self): bootstrap = load_install_env() class SlowLines: def __iter__(self): yield "Resolved 3 packages in 1.0s\n" time.sleep(0.6) class FakeProcess: def __init__(self, lines): self.stdout = lines def wait(self): return 0 clock_state = {"now": 0.0} def fake_clock(): clock_state["now"] += 6.0 return clock_state["now"] fake_process = FakeProcess(SlowLines()) with patch.object( bootstrap.subprocess, "Popen", return_value=fake_process, ), patch.object(bootstrap, "write_status") as mocked_write: bootstrap._run_install_once( "status.json", 2, 2, "ONNX Runtime", ["uv", "pip", "install", "onnxruntime"], None, 0.10, 0.85, fake_clock, ) heartbeats = [ call for call in mocked_write.call_args_list if call.args[2].startswith("Downloading") ] self.assertTrue(heartbeats) self.assertTrue( any("3 packages" in call.args[2] for call in heartbeats) ) for call in heartbeats: self.assertNotIn("elapsed", call.args[2]) self.assertNotIn("eta_seconds", call.kwargs) def test_run_install_retries_across_indexes(self): bootstrap = load_install_env() previous = bootstrap._resolved_pypi_index bootstrap._resolved_pypi_index = None try: with patch.object( bootstrap, "_run_install_once", side_effect=[ subprocess.CalledProcessError(1, ["uv", "pip"]), ["numpy==1.26.4"], ], ) as mocked_install: installed = bootstrap.run_install( "status.json", 3, 8, "ai deps", ["uv", "pip", "install", "x"], ) finally: bootstrap._resolved_pypi_index = previous self.assertEqual(installed, ["numpy==1.26.4"]) self.assertEqual(mocked_install.call_count, 2) # 未指定源时默认国内镜像优先，第一个镜像失败后切到第二个镜像。 self.assertEqual( mocked_install.call_args_list[0].args[5], {"UV_INDEX_URL": bootstrap.PYPI_MIRRORS[0]}, ) self.assertEqual( mocked_install.call_args_list[1].args[5], {"UV_INDEX_URL": bootstrap.PYPI_MIRRORS[1]}, ) def test_run_install_uses_only_configured_index(self): bootstrap = load_install_env() previous = bootstrap._resolved_pypi_index bootstrap._resolved_pypi_index = ( "https://pypi.tuna.tsinghua.edu.cn/simple" ) try: with patch.object( bootstrap, "_run_install_once", return_value=["numpy==1.26.4"], ) as mocked_install: installed = bootstrap.run_install( "status.json", 3, 8, "ai deps", ["uv", "pip", "install", "x"], ) finally: bootstrap._resolved_pypi_index = previous self.assertEqual(installed, ["numpy==1.26.4"]) self.assertEqual(mocked_install.call_count, 1) self.assertEqual( mocked_install.call_args.args[5], {"UV_INDEX_URL": "https://pypi.tuna.tsinghua.edu.cn/simple"}, ) def test_resolve_indexes_uses_configured_pypi_index(self): bootstrap = load_install_env() bootstrap.resolve_indexes( SimpleNamespace(pypi_index="https://pypi.tuna.tsinghua.edu.cn/simple") ) self.assertEqual( bootstrap._resolved_pypi_index, "https://pypi.tuna.tsinghua.edu.cn/simple", ) def test_resolve_indexes_uses_environment_override(self): bootstrap = load_install_env() with patch.dict( os.environ, {"UV_INDEX_URL": "https://pypi.example/simple"}, ): bootstrap.resolve_indexes(SimpleNamespace(pypi_index=None)) self.assertEqual( bootstrap._resolved_pypi_index, "https://pypi.example/simple", ) def test_resolve_indexes_falls_back_to_pip_index_url(self): bootstrap = load_install_env() with patch.dict( os.environ, {"PIP_INDEX_URL": "https://pypi.example/simple"}, ): bootstrap.resolve_indexes(SimpleNamespace(pypi_index=None)) self.assertEqual( bootstrap._resolved_pypi_index, "https://pypi.example/simple", ) def test_child_environment_injects_resolved_pypi_index(self): bootstrap = load_install_env() previous_pypi = bootstrap._resolved_pypi_index bootstrap._resolved_pypi_index = ( "https://pypi.tuna.tsinghua.edu.cn/simple" ) try: env = bootstrap.child_environment() finally: bootstrap._resolved_pypi_index = previous_pypi self.assertEqual( env["UV_INDEX_URL"], "https://pypi.tuna.tsinghua.edu.cn/simple", ) self.assertEqual( env["PYTHONIOENCODING"], "utf-8:backslashreplace", ) self.assertEqual(env["PYTHONUTF8"], "1") def test_standard_streams_are_reconfigured_for_unicode_logs(self): bootstrap = load_install_env() stdout = Mock() stderr = Mock() with patch.object(bootstrap.sys, "stdout", stdout), patch.object( bootstrap.sys, "stderr", stderr, ): bootstrap.configure_standard_streams() stdout.reconfigure.assert_called_once_with( encoding="utf-8", errors="backslashreplace", ) stderr.reconfigure.assert_called_once_with( encoding="utf-8", errors="backslashreplace", ) def test_ready_json_records_installed_packages(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) venv_dir = root / ".venv" venv_dir.mkdir() self.create_bundled_vendor(root) uv_path = root / "uv" uv_path.write_bytes(b"uv") args = SimpleNamespace( uv=str(uv_path), venv_dir=str(venv_dir), models_dir=str(root / "models"), torch_variant="cu128", da3_commit="test-commit", ben2_commit="test-ben2-commit", status=str(root / "status.json"), ) with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object( bootstrap, "run", ), patch.object( bootstrap, "run_stage", return_value=["torch==2.9.1+cu128", "numpy==1.26.4"], ), patch.object( bootstrap, "run_verify", return_value={"cuda": True, "mps": False}, ): bootstrap.install(args) ready = json.loads( (root / "manifest.json").read_text( encoding="utf-8" ) ) self.assertEqual( ready["packages"], ["numpy==1.26.4", "torch==2.9.1+cu128"], ) self.assertEqual(ready["devices"], ["directml", "cpu"]) self.assertEqual( ready["manifest_revision"], bootstrap.MANIFEST_REVISION, ) def test_broken_env_is_cleaned_before_reinstall(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) venv_dir = root / ".venv" venv_dir.mkdir() self.create_bundled_vendor(root) (venv_dir / "partial-python.exe").write_text("junk") args = SimpleNamespace( uv="uv", venv_dir=str(venv_dir), models_dir=str(root / "models"), torch_variant="cu128", da3_commit="test-commit", ben2_commit="test-ben2-commit", status=str(root / "status.json"), ) with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object( bootstrap, "run", ), patch.object( bootstrap, "run_install", ), patch.object( bootstrap, "run_verify", return_value={}, ): bootstrap.install(args) self.assertFalse((venv_dir / "partial-python.exe").exists()) self.assertTrue( (root / "manifest.json").is_file() ) def test_installed_env_is_not_cleaned(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) venv_dir = root / ".venv" venv_dir.mkdir(parents=True) self.create_bundled_vendor(root) (venv_dir / "Scripts").mkdir() (venv_dir / "Scripts" / "python.exe").write_text("python") (root / "manifest.json").write_text( json.dumps( {"manifest_revision": bootstrap.MANIFEST_REVISION} ), encoding="utf-8", ) (venv_dir / "marker.txt").write_text("keep") args = SimpleNamespace( uv="uv", venv_dir=str(venv_dir), models_dir=str(root / "models"), torch_variant="cu128", da3_commit="test-commit", ben2_commit="test-ben2-commit", status=str(root / "status.json"), ) with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object( bootstrap, "run", ), patch.object( bootstrap, "run_install", ), patch.object( bootstrap, "run_verify", return_value={}, ): bootstrap.install(args) self.assertTrue((venv_dir / "marker.txt").is_file()) def test_model_download_failure_keeps_env_installed(self): bootstrap = load_install_env() commands = [] def fake_run(command, environment=None): commands.append(command) if any("moge-2-vits-normal" in str(value) for value in command): raise subprocess.CalledProcessError(1, command) with tempfile.TemporaryDirectory() as directory: root = Path(directory) venv_dir = root / ".venv" venv_dir.mkdir() self.create_bundled_vendor(root) args = SimpleNamespace( uv="uv", venv_dir=str(venv_dir), models_dir=str(root / "models"), torch_variant="cu128", da3_commit="test-commit", ben2_commit="test-ben2-commit", status=str(root / "status.json"), ) with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object( bootstrap, "run", side_effect=fake_run, ), patch.object(bootstrap, "run_install"), patch.object( bootstrap, "run_verify", return_value={}, ): bootstrap.install(args) self.assertTrue( (root / "manifest.json").is_file() ) status = json.loads( (root / "status.json").read_text(encoding="utf-8") ) self.assertEqual(status["progress"], 1.0) self.assertIn("Environment is installed", status["message"]) self.assertTrue( any( "onnx-community/BEN2-ONNX" in map(str, command) for command in commands ) ) def test_install_fails_when_all_uv_sources_are_unavailable(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) venv_dir = root / ".venv" venv_dir.mkdir() uv_path = root / "uv" / "uv.exe" args = SimpleNamespace( uv=str(uv_path), venv_dir=str(venv_dir), models_dir=str(root / "models"), torch_variant="cu128", da3_commit="test-commit", ben2_commit="test-ben2-commit", status=str(root / "status.json"), ) with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object( bootstrap, "run", ), patch.object( bootstrap, "run_install", ), patch.object( bootstrap.shutil, "which", return_value=None, ), patch.object( bootstrap.shutil, "rmtree", ), patch.object( bootstrap, "download_uv_archive", side_effect=RuntimeError("all sources unavailable"), ): with self.assertRaisesRegex( RuntimeError, "all sources unavailable", ): bootstrap.install(args) def test_install_downloads_and_extracts_uv(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) venv_dir = root / ".venv" venv_dir.mkdir() target, archive_name = bootstrap.uv_download_target() archive = root / f"uv-{target}.{archive_name}" if archive_name.endswith("zip"): with zipfile.ZipFile(archive, "w") as bundle: bundle.writestr("uv/uv.exe", b"uv-binary") else: with tarfile.open(archive, "w:gz") as tar: content = b"uv-binary" info = tarfile.TarInfo("uv") info.size = len(content) tar.addfile(info, io.BytesIO(content)) uv_path = root / "uv" / ("uv.exe" if os.name == "nt" else "uv") args = SimpleNamespace( uv=str(uv_path), venv_dir=str(venv_dir), models_dir=str(root / "models"), torch_variant="cu128", da3_commit="test-commit", ben2_commit="test-ben2-commit", status=str(root / "status.json"), ) run_commands = [] install_commands = [] def fake_run(command, environment=None): run_commands.append(command) def fake_run_install( status_path, stage, stages, label, command, environment=None, start=0.0, end=1.0, clock=None, ): install_commands.append(command) return [] def fake_download(destination, *_args): shutil.copy2(archive, destination) return destination with patch.object( bootstrap, "resolve_indexes", return_value=None, ), patch.object( bootstrap, "download_uv_archive", side_effect=fake_download, ), patch.object( bootstrap, "run", side_effect=fake_run, ), patch.object( bootstrap, "run_install", side_effect=fake_run_install, ), patch.object( bootstrap, "run_verify", return_value={}, ): bootstrap.install(args) self.assertEqual(uv_path.read_bytes(), b"uv-binary") all_commands = run_commands + install_commands runtime_command = next( command for command in install_commands if any("onnxruntime" in str(value) for value in command) ) self.assertFalse(any(" @ file:" in str(value) for value in runtime_command)) self.assertFalse(any("torch" in str(value) for value in runtime_command)) status = json.loads(Path(args.status).read_text(encoding="utf-8")) self.assertEqual(status["stage"], 4) self.assertEqual(status["stages"], 4) def test_install_reuses_existing_uv(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) venv_dir = root / ".venv" venv_dir.mkdir() self.create_bundled_vendor(root) uv_path = root / "uv.exe" uv_path.write_bytes(b"uv") args = SimpleNamespace( uv=str(uv_path), venv_dir=str(venv_dir), models_dir=str(root / "models"), torch_variant="cu128", da3_commit="test-commit", ben2_commit="test-ben2-commit", status=str(root / "status.json"), ) with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object( bootstrap, "run", ), patch.object( bootstrap, "run_install", ), patch.object( bootstrap, "run_verify", return_value={}, ): bootstrap.install(args) def test_extract_uv_unpacks_standalone_binary(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) archive_path = root / "uv.zip" with zipfile.ZipFile(archive_path, "w") as bundle: bundle.writestr("uv-x86_64-pc-windows-msvc/uv.exe", "binary") uv_path = root / "uv.exe" bootstrap.extract_uv(archive_path, uv_path) self.assertEqual(uv_path.read_bytes(), b"binary") def test_uv_download_falls_back_to_omoolab_and_verifies_checksum(self): bootstrap = load_install_env() content = io.BytesIO() with zipfile.ZipFile(content, "w") as archive: archive.writestr("uv/uv.exe", b"uv-binary") payload = content.getvalue() class Response(io.BytesIO): headers = {"Content-Length": str(len(payload))} with tempfile.TemporaryDirectory() as directory: root = Path(directory) destination = root / "uv.tmp" status = root / "status.json" urls = [] clock_state = {"now": 0.0} def clock(): clock_state["now"] += 0.25 return clock_state["now"] def open_url(request): urls.append(request.full_url) if len(urls) == 1: raise OSError("official unavailable") return Response(payload) with patch.object( bootstrap, "uv_download_target", return_value=("test-target", "zip"), ), patch.dict( bootstrap.UV_ARTIFACTS, { "test-target": ( "uv-test.zip", hashlib.sha256(payload).hexdigest(), ) }, ), patch.dict( bootstrap._uv_download_openers, { bootstrap.UV_OFFICIAL_BASE_URL: SimpleNamespace( open=open_url ), bootstrap.UV_MIRROR_BASE_URL: SimpleNamespace( open=open_url ), }, ), patch.object( bootstrap, "write_status", ) as write_status: result = bootstrap.download_uv_archive( destination, status, 0.0, 1.0, {}, clock=clock, ) self.assertEqual(result.read_bytes(), payload) self.assertTrue(urls[0].startswith(bootstrap.UV_OFFICIAL_BASE_URL)) self.assertTrue(urls[1].startswith(bootstrap.UV_MIRROR_BASE_URL)) self.assertEqual(bootstrap.UV_OFFICIAL_CONNECT_TIMEOUT, 3.0) self.assertEqual(bootstrap.UV_MIRROR_CONNECT_TIMEOUT, 15.0) messages = [call.args[2] for call in write_status.call_args_list] self.assertTrue(any("uv (mirror)" in value for value in messages)) self.assertTrue(any("B/s" in value for value in messages)) self.assertFalse(any("left" in value for value in messages)) def test_install_stages_use_only_runtime_packages(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) venv_dir = root / ".venv" venv_dir.mkdir() uv_path = root / "uv" uv_path.write_bytes(b"uv") args = SimpleNamespace( uv=str(uv_path), venv_dir=str(venv_dir), models_dir=str(root / "models"), torch_variant="cu128", da3_commit="test-commit", ben2_commit="test-ben2-commit", status=str(root / "status.json"), ) stages = bootstrap.install_stages( args, venv_dir, root / "python", uv_path, ) runtime_command = next( command for _label, command, _environment, kind in stages if kind == "install" and any("onnxruntime" in str(value) for value in command) ) self.assertFalse(any(" @ file:" in str(value) for value in runtime_command)) self.assertFalse(any("torch" in str(value) for value in runtime_command)) def test_install_stages_download_only_default_models(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) venv_dir = root / ".venv" venv_dir.mkdir() vendor = root / "vendor" vendor.mkdir() (vendor / "depth-anything-3.tar.gz").write_bytes(b"da3") (vendor / "ben2.tar.gz").write_bytes(b"ben2") (vendor / "utils3d.tar.gz").write_bytes(b"utils3d") (vendor / "moge.tar.gz").write_bytes(b"moge") uv_path = root / "uv" uv_path.write_bytes(b"uv") args = SimpleNamespace( uv=str(uv_path), venv_dir=str(venv_dir), models_dir=str(root / "models"), torch_variant="cu128", da3_commit="test-commit", ben2_commit="test-ben2-commit", status=str(root / "status.json"), ) stages = bootstrap.install_stages( args, venv_dir, root / "python", uv_path, ) downloads = { label: command for label, command, _environment, kind in stages if kind == "download" } self.assertIn( "moge-2-vits-normal-onnx", " ".join(map(str, downloads["MoGe-2 ViT-S Normal"])), ) self.assertIn("BEN2-ONNX", " ".join(map(str, downloads["BEN2 Base"]))) self.assertEqual( set(downloads), {"MoGe-2 ViT-S Normal", "BEN2 Base"}, ) def test_install_does_not_verify_env(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) venv_dir = root / ".venv" venv_dir.mkdir() self.create_bundled_vendor(root) uv_path = root / "uv" uv_path.write_bytes(b"uv") args = SimpleNamespace( uv=str(uv_path), venv_dir=str(venv_dir), models_dir=str(root / "models"), torch_variant="cu128", da3_commit="test-commit", ben2_commit="test-ben2-commit", status=str(root / "status.json"), ) with patch.object( bootstrap, "resolve_indexes", return_value=None, ), patch.object(bootstrap, "run"), patch.object( bootstrap, "run_install", ), patch.object( bootstrap, "run_verify", ) as mocked_verify: bootstrap.install(args) mocked_verify.assert_not_called() def test_run_verify_returns_detected_devices(self): bootstrap = load_install_env() class FakeProcess: def __init__(self, lines): self.stdout = lines def wait(self): return 0 fake_process = FakeProcess( ['SNAPFORM_DEVICES={"cuda": false, "mps": true}'] ) with patch.object( bootstrap.subprocess, "Popen", return_value=fake_process, ): report = bootstrap.run_verify( "python", "status.json", clock=lambda: 0.0, ) self.assertFalse(report["cuda"]) self.assertTrue(report["mps"]) def test_run_verify_writes_heartbeat_progress(self): bootstrap = load_install_env() class FakeProcess: def __init__(self, lines): self.stdout = lines def wait(self): return 0 clock_state = {"now": 0.0} def fake_clock(): clock_state["now"] += 6.0 return clock_state["now"] def delayed_empty_lines(): time.sleep(0.6) if False: yield fake_process = FakeProcess(delayed_empty_lines()) with patch.object( bootstrap.subprocess, "Popen", return_value=fake_process, ), patch.object(bootstrap, "write_status") as mocked_write: report = bootstrap.run_verify( "python", "status.json", clock=fake_clock, ) messages = [call.args[2] for call in mocked_write.call_args_list] self.assertTrue( any( "Checking Environment" in message for message in messages ) ) self.assertEqual(report, {}) def test_install_uv_extracts_downloaded_archive(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) uv_path = root / "uv" target, archive_name = bootstrap.uv_download_target() archive = root / f"uv-{target}.{archive_name}" if archive_name.endswith("zip"): with zipfile.ZipFile(archive, "w") as bundle: bundle.writestr("uv/uv.exe", b"uv-binary") else: with tarfile.open(archive, "w:gz") as tar: content = b"uv-binary" info = tarfile.TarInfo("uv") info.size = len(content) tar.addfile(info, io.BytesIO(content)) def fake_download(destination, *_args): shutil.copy2(archive, destination) return destination with patch.object( bootstrap, "download_uv_archive", side_effect=fake_download, ): bootstrap.install_uv(uv_path, str(root / "status.json")) self.assertEqual(uv_path.read_bytes(), b"uv-binary") def test_install_uv_raises_when_download_fails(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) uv_path = root / "uv" with patch.object( bootstrap, "download_uv_archive", side_effect=RuntimeError("all uv sources failed"), ): with self.assertRaisesRegex( RuntimeError, "all uv sources failed", ): bootstrap.install_uv(uv_path, str(root / "status.json")) if __name__ == "__main__": unittest.main()