testsAndMisc/python_pkg/wsg_grabber/tests/test_net_resume.py

"""Tests for resuming and discarding partial downloads."""

from __future__ import annotations

from typing import TYPE_CHECKING
from unittest.mock import MagicMock

import pytest

from python_pkg.wsg_grabber import net
from python_pkg.wsg_grabber.constants import API_HEADERS
from python_pkg.wsg_grabber.models import Outcome

if TYPE_CHECKING:
    from pathlib import Path

_BODY = b"fake webm bytes, but real bytes" * 40


def _response(
    status: int = 200,
    *,
    body: bytes = b"",
    headers: dict[str, str] | None = None,
    text: str = "",
) -> MagicMock:
    """Build a stand-in for a requests Response.

    Args:
        status: HTTP status code.
        body: Payload streamed by ``iter_content``.
        headers: Response headers.
        text: Body for JSON decoding.

    Returns:
        MagicMock: Response double.
    """
    response = MagicMock()
    response.status_code = status
    response.headers = headers or {}
    response.text = text
    response.content = text.encode("utf-8")
    response.iter_content.return_value = [body] if body else []
    return response


def _transfer(tmp_path: Path, *, md5: str | None = None, size: int = 0) -> net.Transfer:
    """Build a Transfer pointing at a sandboxed part file.

    Args:
        tmp_path: Per-test temporary directory.
        md5: Expected digest; defaults to the real digest of ``_BODY``.
        size: Expected file size.

    Returns:
        net.Transfer: Transfer description.
    """
    return net.Transfer(
        url="https://i.4cdn.org/wsg/1.webm",
        part_path=tmp_path / "1.webm.part",
        expected_md5=md5 if md5 is not None else net.digest_of(_BODY),
        expected_size=size,
        should_stop=lambda: False,
    )


def test_an_oversized_leftover_is_discarded_before_requesting(
    tmp_path: Path,
) -> None:
    part = tmp_path / "1.webm.part"
    part.write_bytes(b"x" * 9999)
    session = MagicMock()
    session.get.return_value = _response(200, body=_BODY)
    transfer = net.Transfer(
        url="https://i.4cdn.org/wsg/1.webm",
        part_path=part,
        expected_md5=net.digest_of(_BODY),
        expected_size=len(_BODY),
        should_stop=lambda: False,
    )
    result = net.download(session, transfer)
    assert "Range" not in session.get.call_args.kwargs["headers"]
    assert result.outcome is Outcome.COMPLETED


def test_a_rejected_range_restarts_from_zero(tmp_path: Path) -> None:
    part = tmp_path / "1.webm.part"
    part.write_bytes(_BODY[:100])
    session = MagicMock()
    session.get.return_value = _response(416)
    transfer = net.Transfer(
        url="https://i.4cdn.org/wsg/1.webm",
        part_path=part,
        expected_md5=net.digest_of(_BODY),
        expected_size=len(_BODY),
        should_stop=lambda: False,
    )
    result = net.download(session, transfer)
    assert result.outcome is Outcome.TRANSIENT
    assert not part.exists()


def test_shutdown_aborts_mid_stream_and_keeps_the_partial(tmp_path: Path) -> None:
    """Stopping must be immediate and must leave a resumable .part behind."""
    session = MagicMock()
    response = _response(200)
    response.iter_content.return_value = [b"a" * 10, b"b" * 10, b"c" * 10]
    session.get.return_value = response
    transfer = net.Transfer(
        url="https://i.4cdn.org/wsg/1.webm",
        part_path=tmp_path / "1.webm.part",
        expected_md5="whatever" + "=" * 16,
        expected_size=30,
        should_stop=lambda: True,
    )
    result = net.download(session, transfer)
    assert result.outcome is Outcome.ABORTED


def test_empty_chunks_are_skipped(tmp_path: Path) -> None:
    """requests yields keep-alive chunks; they must not corrupt the digest."""
    session = MagicMock()
    response = _response(200)
    response.iter_content.return_value = [b"", _BODY, b""]
    session.get.return_value = response
    result = net.download(session, _transfer(tmp_path))
    assert result.outcome is Outcome.COMPLETED


def test_size_of_missing_file_is_zero(tmp_path: Path) -> None:
    assert net.size_of(tmp_path / "nope") == 0


def test_a_compressed_body_is_refused(tmp_path: Path) -> None:
    """We ask for identity; requests would silently expand a gzip bomb."""
    session = MagicMock()
    session.get.return_value = _response(
        200,
        body=_BODY,
        headers={"Content-Encoding": "gzip"},
    )
    result = net.download(session, _transfer(tmp_path))
    assert result.outcome is Outcome.TRANSIENT
    assert result.message is not None
    assert "Content-Encoding" in result.message


def test_a_body_larger_than_advertised_is_abandoned(tmp_path: Path) -> None:
    """A server that lies about its size must not fill the disk."""
    session = MagicMock()
    response = _response(200)
    response.iter_content.return_value = [b"x" * (1024 * 1024)] * 4
    session.get.return_value = response
    transfer = net.Transfer(
        url="https://i.4cdn.org/wsg/1.webm",
        part_path=tmp_path / "1.webm.part",
        expected_md5=net.digest_of(_BODY),
        expected_size=16,
        should_stop=lambda: False,
    )
    result = net.download(session, transfer)
    assert result.outcome is Outcome.TRANSIENT
    assert not transfer.part_path.exists()


def test_an_unknown_size_is_not_treated_as_over_budget(tmp_path: Path) -> None:
    """fsize 0 means the API did not say; that must not block the download."""
    session = MagicMock()
    session.get.return_value = _response(200, body=_BODY)
    result = net.download(session, _transfer(tmp_path, size=0))
    assert result.outcome is Outcome.COMPLETED


def test_an_oversized_api_response_is_refused() -> None:
    session = MagicMock()
    response = _response(200)
    response.content = b"[" + b"1," * 20_000_000 + b"1]"
    session.get.return_value = response
    with pytest.raises(ValueError, match="exceeded"):
        net.get_json(session, "https://a.4cdn.org/wsg/threads.json")


def test_the_api_session_refuses_compression() -> None:
    assert API_HEADERS["Accept-Encoding"] == "identity"