import importlib.util import io import json import os 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.assertNotIn("--stage", ben2_download) self.assertNotIn("--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_allocates_contiguous_progress_slices(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) starts = [item[1] for item in stage_ranges] ends = [item[2] for item in stage_ranges] self.assertTrue( all(starts[i] < starts[i + 1] for i in range(len(starts) - 1)) ) self.assertEqual(ends[-1], 1.0) self.assertEqual(starts[0], bootstrap.ENVIRONMENT_SLICE[0]) runtime_ranges = stage_ranges[:2] self.assertEqual( [item[0] for item in runtime_ranges], [1, 2], ) self.assertTrue( all(item[0] is None for item in stage_ranges[4:]) ) 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_bundled_uv_missing(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.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, "bundled_uv_archive", return_value=None, ): with self.assertRaisesRegex( RuntimeError, "Bundled uv is missing", ): bootstrap.install(args) def test_install_uses_bundled_vendor_and_uv_without_network(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) venv_dir = root / ".venv" venv_dir.mkdir() bundled = root / "bundled" bundled.mkdir() (bundled / "depth-anything-3.tar.gz").write_bytes(b"da3") (bundled / "ben2.tar.gz").write_bytes(b"ben2") (bundled / "utils3d.tar.gz").write_bytes(b"utils3d") (bundled / "moge.tar.gz").write_bytes(b"moge") target, archive_name = bootstrap.uv_download_target() archive = bundled / 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 [] with patch.object( bootstrap, "bundled_vendor_directory", return_value=bundled, ), patch.object( bootstrap, "resolve_indexes", return_value=None, ), 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 self.assertFalse( any( "github.com" in str(value) for command in all_commands for value in command ) ) 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)) 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_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_bundled_archive(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) uv_path = root / "uv" bundled = root / "bundled" bundled.mkdir(parents=True) target, archive_name = bootstrap.uv_download_target() archive = bundled / 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)) with patch.object( bootstrap, "bundled_vendor_directory", return_value=bundled, ): bootstrap.install_uv(uv_path, str(root / "status.json")) self.assertEqual(uv_path.read_bytes(), b"uv-binary") def test_install_uv_raises_when_bundled_missing(self): bootstrap = load_install_env() with tempfile.TemporaryDirectory() as directory: root = Path(directory) uv_path = root / "uv" bundled = root / "bundled" bundled.mkdir() with patch.object( bootstrap, "bundled_vendor_directory", return_value=bundled, ): with self.assertRaisesRegex( RuntimeError, "Bundled uv is missing", ): bootstrap.install_uv(uv_path, str(root / "status.json")) if __name__ == "__main__": unittest.main()