pyxray/tests/libs/test_xray_assets.py

from __future__ import annotations

import zipfile
from io import BytesIO
from urllib.response import addinfourl

import pytest

from pyxray.libs.xray_assets import (
    DEFAULT_VERSION,
    default_archive_name,
    default_xray_version,
    download_bytes,
    download_bytes_stream,
    ensure_xray_assets,
    latest_xray_version,
    official_archive_url,
    required_files,
)


def _zip_bytes(files: dict[str, bytes]) -> bytes:
    buffer = BytesIO()
    with zipfile.ZipFile(buffer, "w") as archive:
        for name, content in files.items():
            archive.writestr(name, content)
    return buffer.getvalue()


def test_official_archive_url_defaults_to_xray_core_v26_5_9() -> None:
    assert DEFAULT_VERSION == "v26.5.9"
    assert official_archive_url(archive_name="Xray-linux-64.zip") == (
        "https://github.com/XTLS/Xray-core/releases/download/v26.5.9/Xray-linux-64.zip"
    )


def test_default_archive_name_is_platform_specific() -> None:
    assert default_archive_name(os_name="posix", machine="x86_64") == "Xray-linux-64.zip"
    assert default_archive_name(os_name="nt", machine="AMD64") == "Xray-windows-64.zip"
    assert default_archive_name(os_name="nt", machine="ARM64") == "Xray-windows-arm64-v8a.zip"


def test_required_files_are_platform_specific() -> None:
    assert required_files(os_name="posix") == ("xray", "geoip.dat", "geosite.dat")
    assert required_files(os_name="nt") == ("xray.exe", "geoip.dat", "geosite.dat")


def test_latest_xray_version_reads_github_release_tag() -> None:
    def fetcher(url: str, timeout: float) -> str:
        assert url == "https://api.github.com/repos/XTLS/Xray-core/releases/latest"
        assert timeout == 5.0
        return '{"tag_name": "v99.1.2"}'

    assert latest_xray_version(fetcher=fetcher) == "v99.1.2"


def test_default_xray_version_falls_back_to_pinned_version(monkeypatch) -> None:  # noqa: ANN001
    monkeypatch.setattr("pyxray.libs.xray_assets._DEFAULT_VERSION_CACHE", None)
    monkeypatch.setattr(
        "pyxray.libs.xray_assets.latest_xray_version",
        lambda *, timeout: (_ for _ in ()).throw(TimeoutError("slow")),
    )

    assert default_xray_version(timeout=5.0) == DEFAULT_VERSION


def test_ensure_xray_assets_extracts_official_archive_files(tmp_path) -> None:  # noqa: ANN001
    xray_name = required_files()[0]
    archive = _zip_bytes(
        {
            xray_name: b"bin",
            "geoip.dat": b"geoip",
            "geosite.dat": b"geosite",
            "README.md": b"ignored",
        }
    )
    calls = []

    def downloader(url: str) -> bytes:
        calls.append(url)
        return archive

    result = ensure_xray_assets(tmp_path, downloader=downloader)

    assert result.ready is True
    assert result.downloaded == ("archive",)
    assert calls == [official_archive_url()]
    assert (tmp_path / xray_name).read_bytes() == b"bin"
    assert (tmp_path / "geoip.dat").read_bytes() == b"geoip"
    assert (tmp_path / "geosite.dat").read_bytes() == b"geosite"


def test_version_and_archive_url_can_be_overridden(tmp_path) -> None:  # noqa: ANN001
    archive = _zip_bytes({required_files()[0]: b"bin", "geoip.dat": b"geoip", "geosite.dat": b"geosite"})
    calls = []

    def downloader(url: str) -> bytes:
        calls.append(url)
        return archive

    ensure_xray_assets(tmp_path, version="v1.2.3", archive_url="https://mirror.invalid/xray.zip", downloader=downloader)

    assert calls == ["https://mirror.invalid/xray.zip"]


def test_dat_urls_override_archive_dat_files(tmp_path) -> None:  # noqa: ANN001
    xray_name = required_files()[0]
    archive = _zip_bytes({xray_name: b"bin", "geoip.dat": b"old-geoip", "geosite.dat": b"old-geosite"})
    payloads = {
        official_archive_url(): archive,
        "https://mirror.invalid/geoip.dat": b"new-geoip",
        "https://mirror.invalid/geosite.dat": b"new-geosite",
    }

    result = ensure_xray_assets(
        tmp_path,
        geoip_url="https://mirror.invalid/geoip.dat",
        geosite_url="https://mirror.invalid/geosite.dat",
        downloader=payloads.__getitem__,
    )

    assert result.downloaded == ("archive", "geoip.dat", "geosite.dat")
    assert (tmp_path / xray_name).read_bytes() == b"bin"
    assert (tmp_path / "geoip.dat").read_bytes() == b"new-geoip"
    assert (tmp_path / "geosite.dat").read_bytes() == b"new-geosite"


def test_existing_files_skip_download(tmp_path) -> None:  # noqa: ANN001
    xray_name = required_files()[0]
    (tmp_path / xray_name).write_bytes(b"bin")
    (tmp_path / "geoip.dat").write_bytes(b"geoip")
    (tmp_path / "geosite.dat").write_bytes(b"geosite")

    result = ensure_xray_assets(tmp_path, downloader=lambda url: pytest.fail(f"unexpected download: {url}"))

    assert result.ready is True
    assert result.downloaded == ()
    assert result.skipped == (xray_name, "geoip.dat", "geosite.dat")


def test_force_redownloads_selected_existing_file(tmp_path) -> None:  # noqa: ANN001
    (tmp_path / "geoip.dat").write_bytes(b"old")
    calls = []

    def downloader(url: str) -> bytes:
        calls.append(url)
        return b"new"

    result = ensure_xray_assets(
        tmp_path,
        target="geoip",
        force=True,
        geoip_url="https://mirror.invalid/geoip.dat",
        downloader=downloader,
    )

    assert result.downloaded == ("geoip.dat",)
    assert result.skipped == ()
    assert calls == ["https://mirror.invalid/geoip.dat"]
    assert (tmp_path / "geoip.dat").read_bytes() == b"new"


def test_missing_required_file_raises_after_bad_archive(tmp_path) -> None:  # noqa: ANN001
    archive = _zip_bytes({required_files()[0]: b"bin", "geoip.dat": b"geoip"})

    with pytest.raises(FileNotFoundError, match="geosite.dat"):
        ensure_xray_assets(tmp_path, downloader=lambda url: archive)


def test_download_bytes_stream_reports_progress(monkeypatch) -> None:  # noqa: ANN001
    payload = b"abcdef"
    events = []

    def fake_urlopen(url):  # noqa: ANN001
        return addinfourl(BytesIO(payload), {"Content-Length": str(len(payload))}, url)

    monkeypatch.setattr("pyxray.libs.xray_assets.urllib.request.urlopen", fake_urlopen)

    result = download_bytes_stream("https://example.invalid/file", lambda *event: events.append(event), chunk_size=2)

    assert result == payload
    assert events[0] == ("https://example.invalid/file", 0, 6)
    assert events[-1] == ("https://example.invalid/file", 6, 6)


def test_download_bytes_uses_explicit_proxy(monkeypatch) -> None:  # noqa: ANN001
    captured = {}

    class FakeOpener:
        def open(self, url):  # noqa: ANN001
            captured["url"] = url
            return addinfourl(BytesIO(b"ok"), {}, url)

    def fake_proxy_handler(proxies):  # noqa: ANN001
        captured["proxies"] = proxies
        return object()

    monkeypatch.setattr("pyxray.libs.xray_assets.urllib.request.ProxyHandler", fake_proxy_handler)
    monkeypatch.setattr("pyxray.libs.xray_assets.urllib.request.build_opener", lambda handler: FakeOpener())

    result = download_bytes("https://example.invalid/file", proxy_url="http://proxy.example.invalid:8080")

    assert result == b"ok"
    assert captured["url"] == "https://example.invalid/file"
    assert captured["proxies"] == {
        "http": "http://proxy.example.invalid:8080",
        "https": "http://proxy.example.invalid:8080",
    }