AnyImage/tests/test_bootstrap_runtime.py

import importlib.util
import json
import tempfile
import unittest
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_and_non_compiled_gsplat(self):
        bootstrap = load_bootstrap_runtime()
        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, "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,
        )
        gsplat_call = next(
            call
            for call in mocked_run_install.call_args_list
            if any("gsplat @" in str(value) for value in call.args[4])
        )
        gsplat_command = gsplat_call.args[4]
        self.assertIn("--no-build-isolation", gsplat_command)
        self.assertTrue(
            any(bootstrap.GSPLAT_COMMIT in str(value) for value in gsplat_command)
        )
        self.assertEqual(gsplat_call.args[5], {"BUILD_NO_CUDA": "1"})

        ben2_command = next(
            command
            for command in commands
            if any("ben2 @" in str(value) for value in command)
        )
        self.assertTrue(
            any("test-ben2-commit" 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)
        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)
        )
        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], "Download BEN2 Base")

    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,
                    12,
                    "Install 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"], 12)
            self.assertEqual(status["stage_label"], "Install PyTorch")

    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",
                "Downloading torch (10MiB)",
                " Downloaded torch",
                "Prepared 3 packages in 2.0s",
                "Installed 3 packages in 0.5s",
            ]
        )
        with tempfile.TemporaryDirectory() as directory:
            status_path = Path(directory) / "status.json"
            with patch.object(
                bootstrap.subprocess,
                "Popen",
                return_value=fake_process,
            ):
                bootstrap.run_install(
                    str(status_path),
                    2,
                    12,
                    "Install PyTorch",
                    ["uv", "pip", "install", "torch"],
                )

            status = json.loads(status_path.read_text(encoding="utf-8"))
            self.assertEqual(status["progress"], 0.95)
            self.assertEqual(status["message"], "Installing packages")
            self.assertEqual(status["stage"], 2)
            self.assertEqual(status["stages"], 12)
            self.assertEqual(status["stage_label"], "Install PyTorch")


if __name__ == "__main__":
    unittest.main()