Update addon tests, bootstrap runtime, download progress, and worker functionality

- Modified status messages and stage labels in `test_addon.py` for clarity and consistency.
- Added tests to ensure proper host termination and command execution in `test_addon.py`.
- Enhanced `test_bootstrap_runtime.py` with additional tests for installation behavior, including handling of CUDA and MPS devices, and improved error handling for disk space and network checks.
- Updated download progress reporting in `test_download_progress.py` to include download speed and improved message formatting.
- Added new tests in `test_worker.py` for device selection, NPZ file validation, and model downloading behavior with Hugging Face endpoints.
- Bumped version in `uv.lock` to 0.9.30.

tests/test_worker.py

@@ -1,6 +1,7 @@
 import importlib.util
 import inspect
 import json
+import os
 import sys
 import tempfile
 import unittest
@@ -9,6 +10,8 @@ from pathlib import Path
 from types import ModuleType, SimpleNamespace
 from unittest.mock import patch
 
+import numpy as np
+
 
 PROJECT_ROOT = Path(__file__).resolve().parents[1]
 
@@ -61,6 +64,57 @@ class WorkerTest(unittest.TestCase):
             inspect.getsource(self.worker.remove_background),
         )
 
+    def test_select_device_supports_mps_and_auto_fallback(self):
+        mps_torch = SimpleNamespace(
+            cuda=SimpleNamespace(is_available=lambda: False),
+            backends=SimpleNamespace(
+                mps=SimpleNamespace(is_available=lambda: True),
+            ),
+        )
+        mps_torch.device = lambda device_type: SimpleNamespace(
+            type=device_type
+        )
+
+        self.assertEqual(
+            self.worker.select_device("mps", mps_torch).type,
+            "mps",
+        )
+        self.assertEqual(
+            self.worker.select_device("auto", mps_torch).type,
+            "mps",
+        )
+        self.assertEqual(
+            self.worker.select_device("cpu", mps_torch).type,
+            "cpu",
+        )
+
+        cuda_torch = SimpleNamespace(
+            cuda=SimpleNamespace(is_available=lambda: True),
+        )
+        cuda_torch.device = lambda device_type: SimpleNamespace(
+            type=device_type
+        )
+        self.assertEqual(
+            self.worker.select_device("auto", cuda_torch).type,
+            "cuda",
+        )
+
+        cpu_only_torch = SimpleNamespace(
+            cuda=SimpleNamespace(is_available=lambda: False),
+            backends=SimpleNamespace(
+                mps=SimpleNamespace(is_available=lambda: False),
+            ),
+        )
+        cpu_only_torch.device = lambda device_type: SimpleNamespace(
+            type=device_type
+        )
+        self.assertEqual(
+            self.worker.select_device("auto", cpu_only_torch).type,
+            "cpu",
+        )
+        with self.assertRaisesRegex(RuntimeError, "MPS was requested"):
+            self.worker.select_device("mps", cpu_only_torch)
+
     def test_optional_background_removal_runs_before_da3_model_load(self):
         source = inspect.getsource(self.worker.reconstruct)
 
@@ -121,6 +175,38 @@ class WorkerTest(unittest.TestCase):
         self.assertEqual(images, [image_path])
         self.assertIsNone(temporary_directory)
 
+    def test_npz_is_complete_requires_readable_archive(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            valid_path = root / "valid.npz"
+            np.savez_compressed(
+                valid_path,
+                depth=np.zeros((4, 4)),
+            )
+            self.assertTrue(self.worker.npz_is_complete(valid_path))
+
+            empty_path = root / "empty.npz"
+            empty_path.write_bytes(b"PK\x05\x06" + b"\x00" * 18)
+            self.assertFalse(self.worker.npz_is_complete(empty_path))
+
+    def test_wait_for_exported_paths_waits_for_complete_npz(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            npz_path = root / "exports" / "mini_npz" / "results.npz"
+            npz_path.parent.mkdir(parents=True)
+            npz_path.touch()
+            with patch.object(
+                self.worker,
+                "npz_is_complete",
+                side_effect=[False, True],
+            ):
+                result = self.worker.wait_for_exported_paths(
+                    root,
+                    "mini_npz",
+                    timeout=2.0,
+                )
+            self.assertEqual(result, [npz_path])
+
     def test_rejects_unsupported_input_file(self):
         with tempfile.TemporaryDirectory() as directory:
             input_path = Path(directory) / "input.txt"
@@ -149,6 +235,10 @@ class WorkerTest(unittest.TestCase):
             self.worker,
             "create_download_progress_class",
             return_value=object,
+        ), patch.object(
+            self.worker,
+            "resolve_hf_endpoint",
+            return_value=(self.worker.HF_ENDPOINT, False),
         ):
             args.model_dir = str(Path(directory) / "model")
             args.status = str(Path(directory) / "status.json")
@@ -159,6 +249,44 @@ class WorkerTest(unittest.TestCase):
             ["config.json", "model.safetensors"],
         )
 
+    def test_download_model_switches_to_mirror_before_hf_import(self):
+        calls = []
+        fake_hugging_face = ModuleType("huggingface_hub")
+        fake_hugging_face.snapshot_download = lambda **options: calls.append(
+            options
+        )
+        args = SimpleNamespace(
+            repository="PramaLLC/BEN2",
+            model_dir="model",
+            status="status.json",
+            allow_pattern=[],
+        )
+
+        with tempfile.TemporaryDirectory() as directory, patch.dict(
+            sys.modules,
+            {"huggingface_hub": fake_hugging_face},
+        ), patch.object(
+            self.worker,
+            "create_download_progress_class",
+            return_value=object,
+        ), patch.object(
+            self.worker,
+            "resolve_hf_endpoint",
+            return_value=("https://hf-mirror.com", True),
+        ), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
+            self.worker,
+            "write_status",
+        ) as mocked_write:
+            args.model_dir = str(Path(directory) / "model")
+            args.status = str(Path(directory) / "status.json")
+            self.worker.download_model(args)
+            self.assertEqual(
+                os.environ["HF_ENDPOINT"],
+                "https://hf-mirror.com",
+            )
+        messages = [call.args[2] for call in mocked_write.call_args_list]
+        self.assertTrue(any("(mirror)" in message for message in messages))
+
     def test_ben2_background_removal_writes_fast_transparent_png(self):
         inference_options = []
         save_options = []
@@ -279,7 +407,10 @@ class WorkerTest(unittest.TestCase):
             else:
                 mini_directory = output_directory / "exports" / "mini_npz"
                 mini_directory.mkdir(parents=True)
-                (mini_directory / "results.npz").touch()
+                np.savez_compressed(
+                    mini_directory / "results.npz",
+                    depth=np.zeros((14, 14)),
+                )
 
         fake_torch = ModuleType("torch")
         fake_torch.device = lambda device_type: SimpleNamespace(