"""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()