import importlib.util import json import os import subprocess import tempfile import time import unittest import zipfile from pathlib import Path from types import SimpleNamespace from unittest.mock import patch PROJECT_ROOT = Path(__file__).resolve().parents[1] def load_bootstrap_runtime(): specification = importlib.util.spec_from_file_location( "snapform_bootstrap_runtime", PROJECT_ROOT / "src" / "snapform" / "bootstrap_runtime.py", ) module = importlib.util.module_from_spec(specification) specification.loader.exec_module(module) return module class BootstrapRuntimeTest(unittest.TestCase): def test_installs_da3_from_archives_without_gsplat(self): bootstrap = load_bootstrap_runtime() with tempfile.TemporaryDirectory() as directory: root = Path(directory) runtime_dir = root / "runtime" runtime_dir.mkdir() uv_path = root / "uv" uv_path.write_bytes(b"uv") args = SimpleNamespace( uv=str(uv_path), runtime_dir=str(runtime_dir), models_dir=str(root / "models"), backend="cu128", da3_commit="test-commit", ben2_commit="test-ben2-commit", status=str(root / "status.json"), ) with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object( bootstrap, "run", ) as mocked_run, patch.object( bootstrap, "run_install", ) as mocked_run_install: 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 ] da3_command = next( command for command in commands if any("depth-anything-3 @" in str(value) for value in command) ) self.assertIn( f"addict=={bootstrap.ADDICT_VERSION}", da3_command, ) self.assertTrue( any( "https://github.com/ByteDance-Seed/Depth-Anything-3/" "archive/test-commit.tar.gz" in str(value) for value in da3_command ) ) self.assertFalse( any("git+" in str(value) for value in da3_command) ) self.assertFalse( any( "gsplat" in str(value) for command in commands for value in command ) ) self.assertFalse( any( "--no-build-isolation" in str(value) for command in commands for value in command ) ) ben2_command = next( command for command in commands if any("ben2 @" in str(value) for value in command) ) self.assertTrue( any( "https://github.com/PramaLLC/BEN2/archive/" "test-ben2-commit.tar.gz" in str(value) for value in ben2_command ) ) self.assertFalse( any("git+" in str(value) for value in ben2_command) ) host_command = next( command for command in commands if f"fastapi=={bootstrap.FASTAPI_VERSION}" in command ) self.assertIn(f"uvicorn=={bootstrap.UVICORN_VERSION}", host_command) self.assertIn("setuptools", host_command) self.assertNotIn("ninja", host_command) check_command = next( command for command in commands if any("import torch;" in str(value) for value in command) ) self.assertTrue( any( "from ben2 import BEN_Base" in str(value) for value in check_command ) ) self.assertTrue( any("import fastapi, uvicorn" in str(value) for value in check_command) ) self.assertFalse( any("import gsplat" in str(value) for value in check_command) ) default_downloads = [ command for command in commands if any( repository in command for repository in ( "depth-anything/DA3-SMALL", "depth-anything/DA3-BASE", "PramaLLC/BEN2", ) ) ] self.assertEqual(len(default_downloads), 3) ben2_download = next( command for command in default_downloads if "PramaLLC/BEN2" in command ) patterns = [ ben2_download[index + 1] for index, value in enumerate(ben2_download) if value == "--allow-pattern" ] self.assertEqual(patterns, ["config.json", "model.safetensors"]) self.assertIn("--stage", ben2_download) self.assertIn("--stage-label", ben2_download) self.assertEqual(ben2_download[-1], "ben2 base") def test_darwin_install_skips_cuda_index_and_xformers(self): bootstrap = load_bootstrap_runtime() with tempfile.TemporaryDirectory() as directory: root = Path(directory) runtime_dir = root / "runtime" runtime_dir.mkdir() uv_path = root / "uv" uv_path.write_bytes(b"uv") args = SimpleNamespace( uv=str(uv_path), runtime_dir=str(runtime_dir), models_dir=str(root / "models"), backend="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, "preflight", return_value=({}, False)), patch.object( bootstrap, "run", ), patch.object( bootstrap, "run_install", ) as mocked_run_install: bootstrap.install(args) commands = [ call.args[4] for call in mocked_run_install.call_args_list ] torch_command = next( command for command in commands if any( f"torch=={bootstrap.TORCH_VERSION}" in str(value) for value in command ) ) self.assertNotIn("--index-url", torch_command) self.assertFalse( any( "xformers" in str(value) for command in commands for value in command ) ) def test_intel_mac_is_rejected_with_clear_error(self): bootstrap = load_bootstrap_runtime() args = SimpleNamespace( uv="uv", runtime_dir="runtime", models_dir="models", backend="cu128", da3_commit="test-commit", ben2_commit="test-ben2-commit", status="status.json", ) with patch.object( bootstrap.sys, "platform", "darwin", ), patch.object( bootstrap.platform, "machine", return_value="x86_64", ), patch.object(bootstrap.shutil, "rmtree") as mocked_rmtree: with self.assertRaisesRegex(RuntimeError, "Apple Silicon"): bootstrap.install(args) mocked_rmtree.assert_not_called() def test_run_stage_writes_stage_progress(self): bootstrap = load_bootstrap_runtime() 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_bootstrap_runtime() 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_bootstrap_runtime() 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_bootstrap_runtime() with tempfile.TemporaryDirectory() as directory: root = Path(directory) runtime_dir = root / "runtime" runtime_dir.mkdir() uv_path = root / "uv" uv_path.write_bytes(b"uv") args = SimpleNamespace( uv=str(uv_path), runtime_dir=str(runtime_dir), models_dir=str(root / "models"), backend="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, "preflight", return_value=({}, False)), patch.object( bootstrap, "run_stage", side_effect=fake_run_stage, ), patch.object(bootstrap, "run"): 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.RUNTIME_SLICE_START) 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_bootstrap_runtime() 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, [ "3 packages", "3/3 packages", "3/3 packages", ], ) self.assertEqual(mocked_write.call_args_list[-1].args[1], 0.02) 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_bootstrap_runtime() 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("3 packages", messages) self.assertIn("3/3 packages", messages) def test_run_install_maps_progress_into_stage_range(self): bootstrap = load_bootstrap_runtime() 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_formats_elapsed_time(self): bootstrap = load_bootstrap_runtime() self.assertEqual(bootstrap.format_elapsed(12), "12s") self.assertEqual(bootstrap.format_elapsed(83), "1m 23s") def test_preflight_rejects_low_disk_space(self): bootstrap = load_bootstrap_runtime() args = SimpleNamespace( backend="cu128", runtime_dir=str(Path.cwd()), ) with patch.object( bootstrap.shutil, "disk_usage", return_value=SimpleNamespace(free=1024**3), ), patch.object(bootstrap, "reachable", return_value=True): with self.assertRaisesRegex(RuntimeError, "disk space"): bootstrap.preflight(args) def test_preflight_rejects_unreachable_network(self): bootstrap = load_bootstrap_runtime() args = SimpleNamespace( backend="cu128", runtime_dir=str(Path.cwd()), ) with patch.object( bootstrap.shutil, "disk_usage", return_value=SimpleNamespace(free=100 * 1024**3), ), patch.object(bootstrap, "reachable", return_value=False): with self.assertRaisesRegex(RuntimeError, "Cannot reach"): bootstrap.preflight(args) def test_preflight_passes_when_checks_ok(self): bootstrap = load_bootstrap_runtime() args = SimpleNamespace( backend="cu128", runtime_dir=str(Path.cwd()), ) with patch.object( bootstrap.shutil, "disk_usage", return_value=SimpleNamespace(free=100 * 1024**3), ), patch.object(bootstrap, "reachable", return_value=True): resolved, used_mirror = bootstrap.preflight(args) self.assertFalse(used_mirror) self.assertEqual( resolved["pypi"], "https://pypi.org/simple/", ) def test_preflight_falls_back_to_mirrors(self): bootstrap = load_bootstrap_runtime() args = SimpleNamespace( backend="cu128", runtime_dir=str(Path.cwd()), ) def fake_reachable(url, timeout=bootstrap.PREFLIGHT_TIMEOUT): if "github.com" in url: return True if "tuna" in url or "hf-mirror" in url or "aliyun" in url: return True return False with patch.object( bootstrap.shutil, "disk_usage", return_value=SimpleNamespace(free=100 * 1024**3), ), patch.object( bootstrap, "reachable", side_effect=fake_reachable, ): resolved, used_mirror = bootstrap.preflight(args) self.assertTrue(used_mirror) self.assertEqual( resolved["pypi"], "https://pypi.tuna.tsinghua.edu.cn/simple", ) self.assertEqual(resolved["hf"], "https://hf-mirror.com") self.assertEqual( resolved["torch"], "https://mirrors.aliyun.com/pytorch-wheels/cu128", ) self.assertEqual( bootstrap._resolved_torch_index, "https://mirrors.aliyun.com/pytorch-wheels/cu128", ) def test_preflight_uses_environment_overrides(self): bootstrap = load_bootstrap_runtime() args = SimpleNamespace( backend="cu128", runtime_dir=str(Path.cwd()), ) with patch.dict( os.environ, { "PIP_INDEX_URL": "https://pypi.example/simple", "HF_ENDPOINT": "https://hf.example", "SNAPFORM_TORCH_INDEX": "https://torch.example/whl/cu128", }, ), patch.object( bootstrap.shutil, "disk_usage", return_value=SimpleNamespace(free=100 * 1024**3), ), patch.object(bootstrap, "reachable", return_value=True): resolved, used_mirror = bootstrap.preflight(args) self.assertFalse(used_mirror) self.assertEqual(resolved["pypi"], "https://pypi.example/simple") self.assertEqual(resolved["hf"], "https://hf.example") self.assertEqual( resolved["torch"], "https://torch.example/whl/cu128", ) def test_child_environment_injects_resolved_mirrors(self): bootstrap = load_bootstrap_runtime() previous_pypi = bootstrap._resolved_pypi_index previous_hf = bootstrap._resolved_hf_endpoint bootstrap._resolved_pypi_index = ( "https://pypi.tuna.tsinghua.edu.cn/simple" ) bootstrap._resolved_hf_endpoint = "https://hf-mirror.com" try: env = bootstrap.child_environment() finally: bootstrap._resolved_pypi_index = previous_pypi bootstrap._resolved_hf_endpoint = previous_hf self.assertEqual( env["PIP_INDEX_URL"], "https://pypi.tuna.tsinghua.edu.cn/simple", ) self.assertEqual(env["HF_ENDPOINT"], "https://hf-mirror.com") def test_child_environment_injects_registry_proxy(self): bootstrap = load_bootstrap_runtime() with patch.object( bootstrap.urllib.request, "getproxies", return_value={ "http": "http://proxy:8080", "https": "http://proxy:8080", }, ): env = bootstrap.child_environment() self.assertEqual(env["HTTP_PROXY"], "http://proxy:8080") self.assertEqual(env["HTTPS_PROXY"], "http://proxy:8080") def test_child_environment_keeps_explicit_proxy(self): bootstrap = load_bootstrap_runtime() with patch.object( bootstrap.urllib.request, "getproxies", return_value={"https": "http://proxy:8080"}, ): env = bootstrap.child_environment( {"HTTPS_PROXY": "http://explicit:3128"} ) self.assertEqual(env["HTTPS_PROXY"], "http://explicit:3128") def test_ready_json_records_installed_packages(self): bootstrap = load_bootstrap_runtime() with tempfile.TemporaryDirectory() as directory: root = Path(directory) runtime_dir = root / "runtime" runtime_dir.mkdir() uv_path = root / "uv" uv_path.write_bytes(b"uv") args = SimpleNamespace( uv=str(uv_path), runtime_dir=str(runtime_dir), models_dir=str(root / "models"), backend="cu128", da3_commit="test-commit", ben2_commit="test-ben2-commit", status=str(root / "status.json"), ) with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object( bootstrap, "run", ), patch.object( bootstrap, "run_stage", return_value=["torch==2.9.1+cu128", "numpy==1.26.4"], ): bootstrap.install(args) ready = json.loads( (runtime_dir / "snapform_ready.json").read_text( encoding="utf-8" ) ) self.assertEqual( ready["packages"], ["numpy==1.26.4", "torch==2.9.1+cu128"], ) def test_broken_runtime_is_cleaned_before_reinstall(self): bootstrap = load_bootstrap_runtime() with tempfile.TemporaryDirectory() as directory: root = Path(directory) runtime_dir = root / "runtime" runtime_dir.mkdir() (runtime_dir / "partial-python.exe").write_text("junk") args = SimpleNamespace( uv="uv", runtime_dir=str(runtime_dir), models_dir=str(root / "models"), backend="cu128", da3_commit="test-commit", ben2_commit="test-ben2-commit", status=str(root / "status.json"), ) with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object( bootstrap, "run", ), patch.object( bootstrap, "run_install", ): bootstrap.install(args) self.assertFalse((runtime_dir / "partial-python.exe").exists()) self.assertTrue((runtime_dir / "snapform_ready.json").is_file()) def test_ready_runtime_is_not_cleaned(self): bootstrap = load_bootstrap_runtime() with tempfile.TemporaryDirectory() as directory: root = Path(directory) runtime_dir = root / "runtime" runtime_dir.mkdir(parents=True) (runtime_dir / "Scripts").mkdir() (runtime_dir / "Scripts" / "python.exe").write_text("python") (runtime_dir / "snapform_ready.json").write_text( json.dumps({"runtime_schema": bootstrap.RUNTIME_SCHEMA}), encoding="utf-8", ) (runtime_dir / "marker.txt").write_text("keep") args = SimpleNamespace( uv="uv", runtime_dir=str(runtime_dir), models_dir=str(root / "models"), backend="cu128", da3_commit="test-commit", ben2_commit="test-ben2-commit", status=str(root / "status.json"), ) with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object( bootstrap, "run", ), patch.object( bootstrap, "run_install", ): bootstrap.install(args) self.assertTrue((runtime_dir / "marker.txt").is_file()) def test_model_download_failure_keeps_runtime_ready(self): bootstrap = load_bootstrap_runtime() def fake_run(command, environment=None): if any("DA3-SMALL" in str(value) for value in command): raise subprocess.CalledProcessError(1, command) with tempfile.TemporaryDirectory() as directory: root = Path(directory) runtime_dir = root / "runtime" runtime_dir.mkdir() args = SimpleNamespace( uv="uv", runtime_dir=str(runtime_dir), models_dir=str(root / "models"), backend="cu128", da3_commit="test-commit", ben2_commit="test-ben2-commit", status=str(root / "status.json"), ) with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object( bootstrap, "run", side_effect=fake_run, ), patch.object(bootstrap, "run_install"): bootstrap.install(args) self.assertTrue((runtime_dir / "snapform_ready.json").is_file()) status = json.loads( (root / "status.json").read_text(encoding="utf-8") ) self.assertEqual(status["progress"], 1.0) self.assertIn("Runtime is ready", status["message"]) def test_install_downloads_missing_uv(self): bootstrap = load_bootstrap_runtime() with tempfile.TemporaryDirectory() as directory: root = Path(directory) runtime_dir = root / "runtime" runtime_dir.mkdir() uv_path = root / "uv" / "uv.exe" args = SimpleNamespace( uv=str(uv_path), runtime_dir=str(runtime_dir), models_dir=str(root / "models"), backend="cu128", da3_commit="test-commit", ben2_commit="test-ben2-commit", status=str(root / "status.json"), ) with patch.object(bootstrap, "preflight", return_value=({}, False)), 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") as mocked_download: bootstrap.install(args) mocked_download.assert_called_once() call_args = mocked_download.call_args self.assertEqual(call_args.args[0], uv_path) self.assertEqual(call_args.args[1], str(root / "status.json")) self.assertEqual(call_args.kwargs["stage"], 1) self.assertEqual(call_args.kwargs["stage_label"], "uv") def test_install_reuses_existing_uv(self): bootstrap = load_bootstrap_runtime() with tempfile.TemporaryDirectory() as directory: root = Path(directory) runtime_dir = root / "runtime" runtime_dir.mkdir() uv_path = root / "uv.exe" uv_path.write_bytes(b"uv") args = SimpleNamespace( uv=str(uv_path), runtime_dir=str(runtime_dir), models_dir=str(root / "models"), backend="cu128", da3_commit="test-commit", ben2_commit="test-ben2-commit", status=str(root / "status.json"), ) with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object( bootstrap, "run", ), patch.object( bootstrap, "run_install", ), patch.object(bootstrap, "download_uv") as mocked_download: bootstrap.install(args) mocked_download.assert_not_called() def test_extract_uv_unpacks_standalone_binary(self): bootstrap = load_bootstrap_runtime() 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") if __name__ == "__main__": unittest.main()