xiaoxia-saas/tests/unit/test_gpu_encoder.py

"""GpuEncoderClient 单元测试：mock HTTP，验证 health/sync/fallback 逻辑。"""

from __future__ import annotations

import json
import tempfile
import urllib.error
import urllib.request
from http.client import HTTPResponse
from io import BytesIO
from pathlib import Path
from unittest import mock

import pytest

from packages.shared.gpu_encoder import (
    GpuEncodeError,
    GpuEncoderClient,
    GpuHealth,
    reset_gpu_encoder_for_tests,
)


@pytest.fixture(autouse=True)
def _reset_singleton():
    reset_gpu_encoder_for_tests()
    yield
    reset_gpu_encoder_for_tests()


@pytest.fixture
def client():
    return GpuEncoderClient(
        endpoint="http://gpu.example.com:8900",
        relay_base_url="http://api.example.com",
        relay_internal_base_url="http://api-internal:8000",
        sync_timeout=60,
        health_timeout=2,
        relay_secret="test-secret",
    )


def _fake_response(status: int = 200, body: dict | bytes | None = None, headers=None):
    """Fake HTTPResponse that supports chunked read(size) used by _download_to_file."""
    if isinstance(body, dict):
        data = json.dumps(body).encode("utf-8")
    elif body is None:
        data = b""
    else:
        data = body

    # Use a real BytesIO so read(size) works for chunked downloads
    bio = BytesIO(data)

    resp = mock.MagicMock(spec=HTTPResponse)
    resp.status = status
    resp.read.side_effect = lambda n=-1: bio.read(n)
    resp.__enter__ = mock.MagicMock(return_value=resp)
    resp.__exit__ = mock.MagicMock(return_value=False)
    return resp


class TestHealthCheck:
    def test_healthy_nvenc_available(self, client):
        body = {
            "status": "healthy",
            "worker": "gpu-worker-1",
            "gpu": {"name": "Quadro P4000"},
            "nvenc": {"h264_nvenc": True, "hevc_nvenc": True},
        }
        with mock.patch("urllib.request.urlopen", return_value=_fake_response(body=body)):
            h = client.check_health()
        assert h.healthy
        assert h.nvenc_h264
        assert h.ready
        assert h.gpu_name == "Quadro P4000"

    def test_connection_error_returns_unhealthy(self, client):
        with mock.patch("urllib.request.urlopen", side_effect=urllib.error.URLError("timeout")):
            h = client.check_health()
        assert not h.healthy
        assert "health probe failed" in h.error

    def test_bad_json_returns_unhealthy(self, client):
        with mock.patch("urllib.request.urlopen", return_value=_fake_response(body=b"not json")):
            h = client.check_health()
        assert not h.healthy

    def test_nvenc_unavailable(self, client):
        body = {"status": "healthy", "gpu": {"name": "test"}, "nvenc": {"h264_nvenc": False}}
        with mock.patch("urllib.request.urlopen", return_value=_fake_response(body=body)):
            h = client.check_health()
        assert h.healthy
        assert not h.ready


class TestPostSync:
    def test_completed_job_returns_dict(self, client):
        result_body = {
            "job_id": "j1",
            "status": "completed",
            "ffmpeg_rc": 0,
            "uploaded": True,
            "duration": 5.1,
            "size": 123456,
        }
        with mock.patch("urllib.request.urlopen", return_value=_fake_response(body=result_body)) as m:
            res = client._post_sync(
                {
                    "inputs": {"in.mp4": "http://x"},
                    "ffmpeg_args": ["-i", "in.mp4"],
                    "output_url": "http://relay/k?token=s",
                    "timeout": 30,
                },
                mezzanine_path=Path("/tmp/fake.mp4"),
            )
        assert res["status"] == "completed"
        assert res["ffmpeg_rc"] == 0
        req = m.call_args[0][0]
        assert req.full_url == "http://gpu.example.com:8900/api/render/sync"

    def test_ffmpeg_failure_raises(self, client):
        body = {"status": "failed", "ffmpeg_rc": 1, "message": "Invalid data found"}
        with mock.patch("urllib.request.urlopen", return_value=_fake_response(body=body)):
            with pytest.raises(GpuEncodeError, match="rc=1"):
                client._post_sync(
                    {"inputs": {}, "ffmpeg_args": [], "output_url": "", "timeout": 10},
                    mezzanine_path=Path("/tmp/x"),
                )

    def test_http_4xx_raises(self, client):
        err = urllib.error.HTTPError(
            url="http://gpu/render/sync",
            code=422,
            msg="Unprocessable",
            hdrs={},
            fp=BytesIO(b"bad request"),
        )
        with mock.patch("urllib.request.urlopen", side_effect=err):
            with pytest.raises(GpuEncodeError, match="HTTP 422"):
                client._post_sync(
                    {"inputs": {}, "ffmpeg_args": [], "output_url": "", "timeout": 10},
                    mezzanine_path=Path("/tmp/x"),
                )


class TestRelayUrl:
    def test_put_url_uses_external_base(self, client):
        url = client._relay_put_url("abc123", "secret!")
        assert "abc123" in url
        assert "token=secret%21" in url  # urlencoded
        assert url.startswith("http://api.example.com/api/v1/internal/gpu-relay/")

    def test_internal_url_uses_internal_base(self, client):
        url = client._relay_internal_url("abc123", "s")
        assert url.startswith("http://api-internal:8000/api/v1/internal/gpu-relay/abc123")

    def test_internal_url_falls_back_to_external_when_not_set(self):
        c = GpuEncoderClient(
            endpoint="http://gpu",
            relay_base_url="http://api.example.com",
            relay_secret="s",
        )
        put = c._relay_put_url("k", "s")
        internal = c._relay_internal_url("k", "s")
        assert put.startswith("http://api.example.com/")
        # When internal not set, internal_url falls back to external base
        assert internal == put

    def test_encode_uses_different_put_and_get_urls(self, client):
        """encode_mezzanine_to_output should use external URL for PUT and internal for GET/DELETE."""
        put_url = client._relay_put_url("k", "test-secret")
        get_url = client._relay_internal_url("k", "test-secret")
        assert "api.example.com" in put_url
        assert "api-internal:8000" in get_url
        assert put_url != get_url


class TestGetRelaySecret:
    def test_explicit_secret_used(self, client):
        assert client._get_relay_secret() == "test-secret"

    def test_env_secret_used_when_not_explicit(self, monkeypatch):
        monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "from-env")
        monkeypatch.setenv("APP_ENV", "staging")
        c = GpuEncoderClient(endpoint="http://gpu", relay_base_url="http://api")
        assert c._get_relay_secret() == "from-env"

    def test_prod_without_secret_raises(self, monkeypatch):
        monkeypatch.delenv("GPU_ENCODE_RELAY_SECRET", raising=False)
        monkeypatch.setenv("APP_ENV", "production")
        c = GpuEncoderClient(endpoint="http://gpu", relay_base_url="http://api")
        with pytest.raises(GpuEncodeError, match="GPU_ENCODE_RELAY_SECRET"):
            c._get_relay_secret()


class TestDownloadToFile:
    def test_writes_file(self, client, tmp_path):
        data = b"hello" * 1000
        out = tmp_path / "out.mp4"
        with mock.patch("urllib.request.urlopen", return_value=_fake_response(body=data)):
            size = client._download_to_file("http://relay/k?token=s", out)
        assert size == len(data)
        assert out.read_bytes() == data

    def test_empty_file_raises(self, client, tmp_path):
        out = tmp_path / "out.mp4"
        with mock.patch("urllib.request.urlopen", return_value=_fake_response(body=b"")):
            with pytest.raises(GpuEncodeError, match="empty file"):
                client._download_to_file("http://relay/k", out)
        assert not out.exists()