unshackle/tests/downloaders/test_content_encoding.py

import gzip
import importlib
import socket
import threading

import pytest
import requests as rq

from unshackle.core.downloaders.requests import download, is_content_encoded


@pytest.mark.parametrize(
    ("value", "expected"),
    [
        (None, False),
        ("", False),
        ("   ", False),
        ("identity", False),
        ("IDENTITY", False),
        ("identity, identity", False),
        ("gzip", True),
        ("GZIP", True),
        (" gzip ", True),
        ("x-gzip", True),
        ("br", True),
        ("zstd", True),
        ("deflate", True),
        ("compress", True),
        ("dcb", True),
        ("gzip, br", True),
        ("identity, gzip", True),
        ("gzip, identity", True),
        ("gzip,,", True),
        ("some-future-coding", True),
    ],
)
def test_is_content_encoded(value, expected):
    assert is_content_encoded(value) is expected


PLAIN = b"WEBVTT\n\n" + b"".join(f"{i:05d} cue text\n".encode() for i in range(2000))


class _EncodingServer:
    """Serves one gzip body, honouring Range against the encoded representation."""

    def __init__(self):
        self.encoded = gzip.compress(PLAIN)
        self.sock = socket.socket()
        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.sock.bind(("127.0.0.1", 0))
        self.sock.listen(8)
        self.url = f"http://127.0.0.1:{self.sock.getsockname()[1]}/sub.vtt"
        threading.Thread(target=self.serve, daemon=True).start()

    def serve(self):
        while True:
            try:
                conn, _ = self.sock.accept()
            except OSError:
                return
            try:
                request = conn.recv(65536).decode("latin1")
                start = next(
                    (
                        int(line.split("=", 1)[1].split("-")[0])
                        for line in request.split("\r\n")
                        if line.lower().startswith("range:")
                    ),
                    None,
                )
                if start is None:
                    body = self.encoded
                    head = f"HTTP/1.1 200 OK\r\nContent-Length: {len(body)}\r\n"
                else:
                    body = self.encoded[start:]
                    head = (
                        f"HTTP/1.1 206 Partial Content\r\nContent-Length: {len(body)}\r\n"
                        f"Content-Range: bytes {start}-{len(self.encoded) - 1}/{len(self.encoded)}\r\n"
                    )
                conn.sendall(
                    (head + "Content-Encoding: gzip\r\nAccept-Ranges: bytes\r\nConnection: close\r\n\r\n").encode()
                    + body
                )
            except Exception:
                pass
            finally:
                conn.close()

    def close(self):
        self.sock.close()


@pytest.fixture
def encoding_server():
    server = _EncodingServer()
    yield server
    server.close()


def test_encoded_body_is_decoded_on_disk(tmp_path, encoding_server):
    save_path = tmp_path / "sub.vtt"
    list(download(url=encoding_server.url, save_path=save_path, session=rq.Session()))
    assert save_path.read_bytes() == PLAIN


def test_encoded_body_discards_a_partial_instead_of_range_resuming(tmp_path, encoding_server, monkeypatch):
    monkeypatch.setattr(importlib.import_module("unshackle.core.downloaders.requests"), "RETRY_WAIT", 0)
    save_path = tmp_path / "sub.vtt"
    save_path.with_name("sub.vtt.!dev").write_bytes(PLAIN[:5000])
    list(download(url=encoding_server.url, save_path=save_path, session=rq.Session()))
    assert save_path.read_bytes() == PLAIN


class _FakeRnetResponse:
    def __init__(self, status_code, headers, body):
        self.status_code = status_code
        self.headers = headers
        self.content_length = len(body)
        self._body = body

    def raise_for_status(self):
        pass

    def stream(self):
        yield self._body

    def close(self):
        pass


class _FakeRnetSession:
    """Duck-typed rnet stand-in: an encoded 206 for a Range request (the resume the
    fail-closed guard must reject), an encoded 200 for the clean restart."""

    def __init__(self, body):
        self.body = body
        self.range_requests = 0

    def get(self, url, stream=True, **kwargs):
        headers = kwargs.get("headers") or {}
        if any(k.lower() == "range" for k in headers):
            self.range_requests += 1
            return _FakeRnetResponse(206, {"Content-Encoding": "zstd"}, b"\x00mid-stream-slice\x00")
        return _FakeRnetResponse(200, {"Content-Encoding": "zstd"}, self.body)


def test_rnet_encoded_body_discards_a_partial_instead_of_range_resuming(tmp_path, monkeypatch):
    dl = importlib.import_module("unshackle.core.downloaders.requests")
    monkeypatch.setattr(dl, "is_rnet_session", lambda s: isinstance(s, _FakeRnetSession))
    monkeypatch.setattr(dl, "RETRY_WAIT", 0)
    session = _FakeRnetSession(PLAIN)
    save_path = tmp_path / "sub.vtt"
    save_path.with_name("sub.vtt.!dev").write_bytes(PLAIN[:5000])
    list(download(url="http://127.0.0.1:1/sub.vtt", save_path=save_path, session=session))
    # the encoded 206 must be dropped, not appended; the clean restart fetches the whole body
    assert session.range_requests == 1
    assert save_path.read_bytes() == PLAIN