Azornes-Comfyui-Model-Resolver/tests/test_python_download_pause.py

import hashlib
import os
import requests
import tempfile
import threading
import time
from types import SimpleNamespace
from unittest.mock import MagicMock, patch

from core.download import api as download_api
from core.download import orchestrator
from core.download import state as download_state


def test_python_download_control_lifecycle():
    controls = {}
    lock = threading.Lock()
    dl_id = "test-dl-1"

    assert not download_state.is_python_download_active(dl_id, controls, lock)

    event = download_state.register_python_download_control(dl_id, controls, lock)
    assert download_state.is_python_download_active(dl_id, controls, lock)
    assert event.is_set()

    # Pause
    paused = download_state.pause_python_download(dl_id, controls, lock)
    assert paused is True
    assert not event.is_set()

    # Resume
    resumed = download_state.resume_python_download(dl_id, controls, lock)
    assert resumed is True
    assert event.is_set()

    # Unregister
    download_state.unregister_python_download_control(dl_id, controls, lock)
    assert not download_state.is_python_download_active(dl_id, controls, lock)
    assert dl_id not in controls


def test_pause_and_resume_download_facade_integration():
    download_id = "py-dl-42"
    controls = {}
    controls_lock = threading.Lock()
    progress = {
        download_id: {
            "status": "downloading",
            "speed": 500000,
            "downloaded": 1024,
            "total_size": 4096,
            "download_backend": "python",
        }
    }
    persisted_updates = []

    def fake_update_state(dl_id, **kwargs):
        persisted_updates.append((dl_id, kwargs))

    facade = SimpleNamespace(
        cancelled_downloads=set(),
        download_progress=progress,
        download_lock=threading.Lock(),
        python_download_controls=controls,
        python_download_controls_lock=controls_lock,
        is_python_download_active=lambda did: did in controls,
        pause_python_download=lambda did: download_state.pause_python_download(
            did, controls, controls_lock
        ),
        resume_python_download=lambda did: download_state.resume_python_download(
            did, controls, controls_lock
        ),
        _set_download_progress_status=lambda did, st, **kw: (
            progress[did].update(status=st, **kw)
        ),
        update_download_progress_state=fake_update_state,
        _queue_aria2_desired_state=MagicMock(),
    )

    # Register active python download
    event = download_state.register_python_download_control(
        download_id, controls, controls_lock
    )

    # Test pause
    pause_res = download_api.pause_download(download_id, dependencies=facade)
    assert pause_res == {"success": True, "message": "Download paused"}
    assert not event.is_set()
    assert progress[download_id]["status"] == "paused"
    assert progress[download_id]["speed"] == 0
    assert ("py-dl-42", {"status": "paused"}) in persisted_updates

    # Pause again while already paused
    pause_again = download_api.pause_download(download_id, dependencies=facade)
    assert pause_again == {"success": True, "message": "Download paused"}

    # Test resume
    resume_res = download_api.resume_download(download_id, dependencies=facade)
    assert resume_res == {"success": True, "message": "Download resumed"}
    assert event.is_set()
    assert progress[download_id]["status"] == "downloading"
    assert ("py-dl-42", {"status": "downloading"}) in persisted_updates


def test_pause_and_resume_reject_when_cancelled():
    download_id = "cancelling-dl"
    facade = SimpleNamespace(
        cancelled_downloads={download_id},
        is_python_download_active=lambda did: True,
    )
    p_res = download_api.pause_download(download_id, dependencies=facade)
    assert p_res == {"success": False, "error": "Download is being cancelled"}

    r_res = download_api.resume_download(download_id, dependencies=facade)
    assert r_res == {"success": False, "error": "Download is being cancelled"}


def test_in_flight_pause_and_resume_streaming():
    """Verify that a Python download pauses chunk writing until resumed, then completes."""
    with tempfile.TemporaryDirectory() as tmpdir:
        dest_path = os.path.join(tmpdir, "test_model.safetensors")
        download_id = "stream-pause-test"

        chunks = [b"chunk1_", b"chunk2_", b"chunk3_"]
        chunk1_yielded = threading.Event()
        allow_chunk2 = threading.Event()

        class MockResponse:
            def __init__(self):
                self.headers = {"content-length": str(sum(len(c) for c in chunks))}
                self.status_code = 200

            def raise_for_status(self):
                pass

            def iter_content(self, chunk_size=None):
                yield chunks[0]
                chunk1_yielded.set()
                allow_chunk2.wait(timeout=5)
                yield chunks[1]
                yield chunks[2]

            def close(self):
                pass

        mock_resp = MockResponse()
        controls = {}
        controls_lock = threading.Lock()
        progress = {}
        cancelled_set = set()

        facade = SimpleNamespace(
            hashlib=download_api.context.hashlib,
            os=os,
            requests=download_api.context.requests,
            time=time,
            urlparse=download_api.context.urlparse,
            deque=download_api.context.deque,
            download_progress=progress,
            download_lock=threading.Lock(),
            cancelled_downloads=cancelled_set,
            SPEED_HISTORY_SIZE=10,
            CHUNK_SIZE=1024,
            CLI_LOG_INTERVAL=60,
            log=download_api.log,
            _extract_expected_sha256=lambda m: None,
            _download_backend_from_settings=lambda: "python",
            _download_huggingface_xet=lambda *a, **k: None,
            get_filename_from_path=os.path.basename,
            host_matches_domain=lambda *a: False,
            _strip_sensitive_url_params=lambda u: u,
            build_download_headers=lambda u, h: {},
            format_bytes=lambda b: f"{b}B",
            _delete_python_partial_download_file=lambda p: (
                os.path.exists(p) and os.unlink(p)
            ),
            _sanitize_download_error=str,
            extract_response_file_size=lambda r: sum(len(c) for c in chunks),
            write_model_resolver_metadata=lambda *a, **k: None,
            invalidate_model_caches=lambda: None,
            register_python_download_control=lambda did: download_state.register_python_download_control(
                did, controls, controls_lock
            ),
            unregister_python_download_control=lambda did: download_state.unregister_python_download_control(
                did, controls, controls_lock
            ),
            is_python_download_active=lambda did: did in controls,
            pause_python_download=lambda did: download_state.pause_python_download(
                did, controls, controls_lock
            ),
            resume_python_download=lambda did: download_state.resume_python_download(
                did, controls, controls_lock
            ),
            update_download_progress_state=MagicMock(),
        )

        def mock_request(*args, **kwargs):
            return mock_resp, "https://example.com/test_model.safetensors", {}

        with patch.object(orchestrator, "request_download_response", side_effect=mock_request):
            download_thread = threading.Thread(
                target=orchestrator.download_file,
                args=("https://example.com/test_model.safetensors", dest_path, download_id),
                kwargs={"dependencies": facade},
            )
            download_thread.start()

            # Wait for first chunk to be yielded
            assert chunk1_yielded.wait(timeout=2)

            # Pause download
            download_state.pause_python_download(download_id, controls, controls_lock)

            # Now allow chunk 2 to be yielded into the loop
            allow_chunk2.set()

            # Wait to confirm loop entered paused state on chunk 2
            deadline = time.time() + 2
            while time.time() < deadline:
                with facade.download_lock:
                    if progress.get(download_id, {}).get("status") == "paused":
                        break
                time.sleep(0.05)
            assert progress[download_id]["status"] == "paused"
            assert progress[download_id]["speed"] == 0

            # Resume download
            download_state.resume_python_download(download_id, controls, controls_lock)

            download_thread.join(timeout=5)
            assert not download_thread.is_alive()

            # Download completed
            assert progress[download_id]["status"] == "completed"
            assert os.path.exists(dest_path)
            with open(dest_path, "rb") as f:
                assert f.read() == b"chunk1_chunk2_chunk3_"


def test_in_flight_pause_and_cancel():
    """Verify that cancelling while paused unblocks the thread and cleans up."""
    with tempfile.TemporaryDirectory() as tmpdir:
        dest_path = os.path.join(tmpdir, "cancel_model.safetensors")
        download_id = "cancel-pause-test"

        first_chunk_event = threading.Event()

        class InfiniteResponse:
            def __init__(self):
                self.headers = {"content-length": "1000000"}
                self.status_code = 200

            def raise_for_status(self):
                pass

            def iter_content(self, chunk_size=None):
                first_chunk_event.set()
                while True:
                    yield b"infinite_data_"
                    time.sleep(0.05)

            def close(self):
                pass

        mock_resp = InfiniteResponse()
        controls = {}
        controls_lock = threading.Lock()
        progress = {}
        cancelled_set = set()

        facade = SimpleNamespace(
            hashlib=download_api.context.hashlib,
            os=os,
            requests=download_api.context.requests,
            time=time,
            urlparse=download_api.context.urlparse,
            deque=download_api.context.deque,
            download_progress=progress,
            download_lock=threading.Lock(),
            cancelled_downloads=cancelled_set,
            SPEED_HISTORY_SIZE=10,
            CHUNK_SIZE=1024,
            CLI_LOG_INTERVAL=60,
            log=download_api.log,
            _extract_expected_sha256=lambda m: None,
            _download_backend_from_settings=lambda: "python",
            _download_huggingface_xet=lambda *a, **k: None,
            get_filename_from_path=os.path.basename,
            host_matches_domain=lambda *a: False,
            _strip_sensitive_url_params=lambda u: u,
            build_download_headers=lambda u, h: {},
            format_bytes=lambda b: f"{b}B",
            _delete_python_partial_download_file=lambda p: (
                os.path.exists(p) and os.unlink(p)
            ),
            _sanitize_download_error=str,
            extract_response_file_size=lambda r: 1000000,
            write_model_resolver_metadata=lambda *a, **k: None,
            invalidate_model_caches=lambda: None,
            register_python_download_control=lambda did: download_state.register_python_download_control(
                did, controls, controls_lock
            ),
            unregister_python_download_control=lambda did: download_state.unregister_python_download_control(
                did, controls, controls_lock
            ),
            is_python_download_active=lambda did: did in controls,
            pause_python_download=lambda did: download_state.pause_python_download(
                did, controls, controls_lock
            ),
            resume_python_download=lambda did: download_state.resume_python_download(
                did, controls, controls_lock
            ),
            update_download_progress_state=MagicMock(),
        )

        def mock_request(*args, **kwargs):
            return mock_resp, "https://example.com/cancel_model.safetensors", {}

        with patch.object(orchestrator, "request_download_response", side_effect=mock_request):
            download_thread = threading.Thread(
                target=orchestrator.download_file,
                args=("https://example.com/cancel_model.safetensors", dest_path, download_id),
                kwargs={"dependencies": facade},
            )
            download_thread.start()

            assert first_chunk_event.wait(timeout=2)

            # Pause download
            download_state.pause_python_download(download_id, controls, controls_lock)

            # Confirm paused
            deadline = time.time() + 2
            while time.time() < deadline:
                with facade.download_lock:
                    if progress.get(download_id, {}).get("status") == "paused":
                        break
                time.sleep(0.05)
            assert progress[download_id]["status"] == "paused"

            # Now cancel download while paused
            cancelled_set.add(download_id)

            # Waking up event if done by state
            download_state.cancel_download(
                download_id,
                SimpleNamespace(
                    download_progress=progress,
                    download_lock=facade.download_lock,
                    cancelled_downloads=cancelled_set,
                    aria2_lock=threading.Lock(),
                    aria2_transfers={},
                    aria2_desired_states={},
                    xet_transfers={},
                    xet_transfers_lock=threading.Lock(),
                    python_download_controls=controls,
                    python_download_controls_lock=controls_lock,
                    set_download_progress_status=MagicMock(),
                    force_remove_aria2_transfer=MagicMock(),
                    thread_factory=threading.Thread,
                    log=facade.log,
                ),
            )

            download_thread.join(timeout=3)
            assert not download_thread.is_alive()
            assert progress[download_id]["status"] == "cancelled"
            assert not os.path.exists(f"{dest_path}.part")
            assert not os.path.exists(dest_path)


def test_python_download_416_complete_file_recovered_and_published():
    """Verify HTTP 416 on resume treats an already complete .part file as finished and publishes it."""
    with tempfile.TemporaryDirectory() as tmpdir:
        dest_path = os.path.join(tmpdir, "model_416.safetensors")
        part_path = f"{dest_path}.part"
        download_id = "dl-416-complete"

        data = b"X" * 1024
        with open(part_path, "wb") as f:
            f.write(data)

        expected_hash = hashlib.sha256(data).hexdigest()
        controls = {}
        controls_lock = threading.Lock()
        progress = {}
        cancelled_set = set()

        facade = SimpleNamespace(
            hashlib=hashlib,
            os=os,
            requests=requests,
            time=time,
            urlparse=download_api.context.urlparse,
            deque=download_api.context.deque,
            download_progress=progress,
            download_lock=threading.Lock(),
            cancelled_downloads=cancelled_set,
            SPEED_HISTORY_SIZE=10,
            CHUNK_SIZE=1024,
            CLI_LOG_INTERVAL=60,
            log=download_api.log,
            load_settings=lambda: {},
            PYTHON_DOWNLOAD_CONNECT_TIMEOUT=10,
            PYTHON_DOWNLOAD_READ_TIMEOUT=15,
            PYTHON_DOWNLOAD_MAX_RETRIES=2,
            PYTHON_DOWNLOAD_RETRY_DELAY=0.1,
            _extract_expected_sha256=lambda m: (m or {}).get("sha256"),
            _download_backend_from_settings=lambda: "python",
            _download_huggingface_xet=lambda *a, **k: None,
            get_filename_from_path=os.path.basename,
            host_matches_domain=lambda *a: False,
            _strip_sensitive_url_params=lambda u: u,
            build_download_headers=lambda u, h: dict(h or {}),
            format_bytes=lambda b: f"{b}B",
            _delete_python_partial_download_file=lambda p: os.path.exists(p) and os.unlink(p),
            _sanitize_download_error=str,
            extract_response_file_size=lambda r: (
                1024 if getattr(r, "status_code", None) == 416 else 1024
            ),
            write_model_resolver_metadata=lambda *a, **k: os.path.join(tmpdir, "meta.json"),
            invalidate_model_caches=lambda: None,
            register_python_download_control=lambda did: download_state.register_python_download_control(
                did, controls, controls_lock
            ),
            unregister_python_download_control=lambda did: download_state.unregister_python_download_control(
                did, controls, controls_lock
            ),
        )

        mock_416_resp = MagicMock()
        mock_416_resp.status_code = 416
        mock_416_resp.headers = {"Content-Range": "bytes */1024"}
        http_err = requests.exceptions.HTTPError(response=mock_416_resp)

        def mock_request(*args, **kwargs):
            raise http_err

        with patch.object(orchestrator, "request_download_response", side_effect=mock_request):
            result = orchestrator.download_file(
                "https://example.com/model.safetensors",
                dest_path,
                download_id,
                metadata={"sha256": expected_hash},
                dependencies=facade,
            )

            assert result["success"] is True
            assert result["size"] == 1024
            assert result["sha256"] == expected_hash
            assert result["sha256_verified"] is True
            assert not os.path.exists(part_path)
            assert os.path.exists(dest_path)
            with open(dest_path, "rb") as f:
                assert f.read() == data


def test_python_download_416_size_mismatch_resets_and_redownloads():
    """Verify HTTP 416 on resume resets .part file and redownloads when size does not match."""
    with tempfile.TemporaryDirectory() as tmpdir:
        dest_path = os.path.join(tmpdir, "model_416_mismatch.safetensors")
        part_path = f"{dest_path}.part"
        download_id = "dl-416-mismatch"

        with open(part_path, "wb") as f:
            f.write(b"Z" * 500)

        new_data = b"CORRECT_FULL_CONTENT"
        expected_hash = hashlib.sha256(new_data).hexdigest()

        controls = {}
        controls_lock = threading.Lock()
        progress = {}
        cancelled_set = set()

        class MockFreshResponse:
            def __init__(self):
                self.status_code = 200
                self.headers = {"content-length": str(len(new_data))}

            def raise_for_status(self):
                pass

            def iter_content(self, chunk_size=None):
                yield new_data

            def close(self):
                pass

        fresh_resp = MockFreshResponse()

        facade = SimpleNamespace(
            hashlib=hashlib,
            os=os,
            requests=requests,
            time=time,
            urlparse=download_api.context.urlparse,
            deque=download_api.context.deque,
            download_progress=progress,
            download_lock=threading.Lock(),
            cancelled_downloads=cancelled_set,
            SPEED_HISTORY_SIZE=10,
            CHUNK_SIZE=1024,
            CLI_LOG_INTERVAL=60,
            log=download_api.log,
            load_settings=lambda: {},
            PYTHON_DOWNLOAD_CONNECT_TIMEOUT=10,
            PYTHON_DOWNLOAD_READ_TIMEOUT=15,
            PYTHON_DOWNLOAD_MAX_RETRIES=2,
            PYTHON_DOWNLOAD_RETRY_DELAY=0.1,
            _extract_expected_sha256=lambda m: (m or {}).get("sha256"),
            _download_backend_from_settings=lambda: "python",
            _download_huggingface_xet=lambda *a, **k: None,
            get_filename_from_path=os.path.basename,
            host_matches_domain=lambda *a: False,
            _strip_sensitive_url_params=lambda u: u,
            build_download_headers=lambda u, h: dict(h or {}),
            format_bytes=lambda b: f"{b}B",
            _delete_python_partial_download_file=lambda p: os.path.exists(p) and os.unlink(p),
            _sanitize_download_error=str,
            extract_response_file_size=lambda r: (
                1000 if getattr(r, "status_code", None) == 416 else len(new_data)
            ),
            write_model_resolver_metadata=lambda *a, **k: os.path.join(tmpdir, "meta.json"),
            invalidate_model_caches=lambda: None,
            register_python_download_control=lambda did: download_state.register_python_download_control(
                did, controls, controls_lock
            ),
            unregister_python_download_control=lambda did: download_state.unregister_python_download_control(
                did, controls, controls_lock
            ),
        )

        mock_416_resp = MagicMock()
        mock_416_resp.status_code = 416
        mock_416_resp.headers = {"Content-Range": "bytes */1000"}
        http_err = requests.exceptions.HTTPError(response=mock_416_resp)

        call_count = 0

        def mock_request(fcd, did, url, headers=None, **kwargs):
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                assert "Range" in (headers or {})
                raise http_err
            assert "Range" not in (headers or {})
            return fresh_resp, url, {}

        with patch.object(orchestrator, "request_download_response", side_effect=mock_request):
            result = orchestrator.download_file(
                "https://example.com/model.safetensors",
                dest_path,
                download_id,
                metadata={"sha256": expected_hash},
                dependencies=facade,
            )

            assert result["success"] is True
            assert result["size"] == len(new_data)
            assert result["sha256"] == expected_hash
            assert not os.path.exists(part_path)
            assert os.path.exists(dest_path)
            with open(dest_path, "rb") as f:
                assert f.read() == new_data


def test_python_download_stall_auto_reconnects():
    """Verify that a socket ReadTimeout stall triggers automatic reconnect and resumes from offset."""
    with tempfile.TemporaryDirectory() as tmpdir:
        dest_path = os.path.join(tmpdir, "model_stall.safetensors")
        part_path = f"{dest_path}.part"
        download_id = "dl-stall-test"

        chunk1 = b"PART_ONE_"
        chunk2 = b"PART_TWO_DONE"
        total_data = chunk1 + chunk2
        expected_hash = hashlib.sha256(total_data).hexdigest()

        controls = {}
        controls_lock = threading.Lock()
        progress = {}
        cancelled_set = set()

        class MockStallingResponse:
            def __init__(self):
                self.status_code = 200
                self.headers = {"content-length": str(len(total_data))}

            def raise_for_status(self):
                pass

            def iter_content(self, chunk_size=None):
                yield chunk1
                raise requests.exceptions.ReadTimeout("Read timed out without data")

            def close(self):
                pass

        class MockResumeResponse:
            def __init__(self):
                self.status_code = 206
                self.headers = {
                    "content-length": str(len(chunk2)),
                    "content-range": f"bytes {len(chunk1)}-{len(total_data)-1}/{len(total_data)}",
                }

            def raise_for_status(self):
                pass

            def iter_content(self, chunk_size=None):
                yield chunk2

            def close(self):
                pass

        resp1 = MockStallingResponse()
        resp2 = MockResumeResponse()

        facade = SimpleNamespace(
            hashlib=hashlib,
            os=os,
            requests=requests,
            time=time,
            urlparse=download_api.context.urlparse,
            deque=download_api.context.deque,
            download_progress=progress,
            download_lock=threading.Lock(),
            cancelled_downloads=cancelled_set,
            SPEED_HISTORY_SIZE=10,
            CHUNK_SIZE=1024,
            CLI_LOG_INTERVAL=60,
            log=download_api.log,
            load_settings=lambda: {},
            PYTHON_DOWNLOAD_CONNECT_TIMEOUT=10,
            PYTHON_DOWNLOAD_READ_TIMEOUT=15,
            PYTHON_DOWNLOAD_MAX_RETRIES=2,
            PYTHON_DOWNLOAD_RETRY_DELAY=0.01,
            _extract_expected_sha256=lambda m: (m or {}).get("sha256"),
            _download_backend_from_settings=lambda: "python",
            _download_huggingface_xet=lambda *a, **k: None,
            get_filename_from_path=os.path.basename,
            host_matches_domain=lambda *a: False,
            _strip_sensitive_url_params=lambda u: u,
            build_download_headers=lambda u, h: dict(h or {}),
            format_bytes=lambda b: f"{b}B",
            _delete_python_partial_download_file=lambda p: os.path.exists(p) and os.unlink(p),
            _sanitize_download_error=str,
            extract_response_file_size=lambda r: (
                len(total_data) if getattr(r, "status_code", None) == 200 else len(chunk2)
            ),
            write_model_resolver_metadata=lambda *a, **k: os.path.join(tmpdir, "meta.json"),
            invalidate_model_caches=lambda: None,
            register_python_download_control=lambda did: download_state.register_python_download_control(
                did, controls, controls_lock
            ),
            unregister_python_download_control=lambda did: download_state.unregister_python_download_control(
                did, controls, controls_lock
            ),
        )

        attempts = 0

        def mock_request(fcd, did, url, headers=None, **kwargs):
            nonlocal attempts
            attempts += 1
            if attempts == 1:
                return resp1, url, {}
            assert headers.get("Range") == f"bytes={len(chunk1)}-"
            return resp2, url, {}

        with patch.object(orchestrator, "request_download_response", side_effect=mock_request):
            result = orchestrator.download_file(
                "https://example.com/model_stall.safetensors",
                dest_path,
                download_id,
                metadata={"sha256": expected_hash},
                dependencies=facade,
            )

            assert result["success"] is True
            assert result["size"] == len(total_data)
            assert result["sha256"] == expected_hash
            assert not os.path.exists(part_path)
            assert os.path.exists(dest_path)
            with open(dest_path, "rb") as f:
                assert f.read() == total_data