xiaoxia-agent
66dc73c8d3
- Update tests to use _relay_put_url/_relay_internal_url instead of removed _relay_url - Fix _fake_response to support chunked read(size) used by _download_to_file - Fix regex match for error message (rc=1 not ffmpeg_rc=1) - Add tests for internal/external URL separation and fallback behavior
212 lines
7.7 KiB
Python
212 lines
7.7 KiB
Python
"""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()
|