AnyImage/tests/test_download_progress.py

import importlib.util
import json
import os
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch


PROJECT_ROOT = Path(__file__).resolve().parents[1]


def load_worker():
    specification = importlib.util.spec_from_file_location(
        "snapform_download_worker",
        PROJECT_ROOT / "src" / "snapform" / "worker.py",
    )
    module = importlib.util.module_from_spec(specification)
    specification.loader.exec_module(module)
    return module


class FakeTqdm:
    def __init__(self, *args, **kwargs):
        self.n = kwargs.get("initial", 0)
        self.total = kwargs.get("total", 0)

    def update(self, amount=1):
        self.n += amount
        return True

    def refresh(self, *args, **kwargs):
        return True


class DownloadProgressTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.worker = load_worker()

    def test_reports_aggregated_download_bytes(self):
        with tempfile.TemporaryDirectory() as directory:
            status_path = Path(directory) / "status.json"
            progress_class = self.worker.create_download_progress_class(
                status_path,
                "owner/model",
                tqdm_type=FakeTqdm,
                clock=lambda: 1.0,
            )
            progress = progress_class(
                desc="Downloading bytes",
                total=100,
                initial=0,
                unit="B",
            )

            progress.update(25)

            status = json.loads(status_path.read_text(encoding="utf-8"))
            self.assertAlmostEqual(status["progress"], 0.275)
            self.assertEqual(status["message"], "downloading owner/model")

    def test_reports_download_speed(self):
        class FakeTqdmWithRate(FakeTqdm):
            @property
            def format_dict(self):
                return {"rate": 14_000_000}

        with tempfile.TemporaryDirectory() as directory:
            status_path = Path(directory) / "status.json"
            progress_class = self.worker.create_download_progress_class(
                status_path,
                "owner/model",
                tqdm_type=FakeTqdmWithRate,
                clock=lambda: 1.0,
                label="da3 small",
            )
            progress = progress_class(
                desc="Downloading bytes",
                total=100,
                initial=0,
                unit="B",
            )

            progress.update(50)

            status = json.loads(status_path.read_text(encoding="utf-8"))
            self.assertEqual(
                status["message"],
                "downloading da3 small 13.4 MB/s",
            )

    def test_staged_download_prefixes_speed_with_downloading(self):
        class FakeTqdmWithRate(FakeTqdm):
            @property
            def format_dict(self):
                return {"rate": 14_000_000}

        with tempfile.TemporaryDirectory() as directory:
            status_path = Path(directory) / "status.json"
            progress_class = self.worker.create_download_progress_class(
                status_path,
                "owner/model",
                tqdm_type=FakeTqdmWithRate,
                clock=lambda: 1.0,
                label="da3 small",
                stage=6,
                stages=8,
                stage_label="da3 small",
            )
            progress = progress_class(
                desc="Downloading bytes",
                total=100,
                initial=0,
                unit="B",
            )

            progress.update(50)

            status = json.loads(status_path.read_text(encoding="utf-8"))
            self.assertEqual(status["message"], "downloading 13.4 MB/s")
            self.assertEqual(status["stage"], 6)
            self.assertEqual(status["stages"], 8)
            self.assertEqual(status["stage_label"], "da3 small")

    def test_ignores_unrelated_progress_bars(self):
        with tempfile.TemporaryDirectory() as directory:
            status_path = Path(directory) / "status.json"
            progress_class = self.worker.create_download_progress_class(
                status_path,
                "owner/model",
                tqdm_type=FakeTqdm,
                clock=lambda: 1.0,
            )
            progress = progress_class(
                desc="Reconstructing",
                total=100,
                initial=0,
                unit="B",
            )

            progress.update(25)

            self.assertFalse(status_path.exists())

    def test_download_endpoints_keeps_direct_endpoint(self):
        with patch.dict(os.environ, {"HF_ENDPOINT": ""}):
            endpoints = self.worker.download_endpoints()

        self.assertEqual(
            endpoints,
            ["https://huggingface.co"],
        )

    def test_download_endpoints_respects_environment_override(self):
        with patch.dict(os.environ, {"HF_ENDPOINT": "https://hf.example"}):
            endpoints = self.worker.download_endpoints()

        self.assertEqual(
            endpoints,
            ["https://hf.example"],
        )

    def test_download_endpoints_strips_trailing_slash(self):
        with patch.dict(os.environ, {"HF_ENDPOINT": "https://hf.example/"}):
            endpoints = self.worker.download_endpoints()

        self.assertEqual(endpoints, ["https://hf.example"])


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