xiaoxia-saas/tests/unit/test_gpu_relay.py

"""gpu_relay API 路由单元测试：覆盖 helper 函数 + PUT/GET/HEAD/DELETE handler。"""
from __future__ import annotations

import os
from pathlib import Path
from unittest import mock

import pytest
from fastapi import HTTPException

from apps.api.app.api.routes import gpu_relay


# ── _relay_dir ────────────────────────────────────────────────────────
class TestRelayDir:
    def test_default_dir(self, tmp_path, monkeypatch):
        monkeypatch.delenv("GENERATED_FILES_DIR", raising=False)
        monkeypatch.delenv("GPU_ENCODE_RELAY_DIR", raising=False)
        # 用 tmp_path 作 base
        monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path))
        p = gpu_relay._relay_dir()
        assert p == tmp_path / "gpu_relay"
        assert p.exists()

    def test_custom_subdir(self, tmp_path, monkeypatch):
        monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path))
        monkeypatch.setenv("GPU_ENCODE_RELAY_DIR", "custom_relay")
        p = gpu_relay._relay_dir()
        assert p == tmp_path / "custom_relay"
        assert p.exists()


# ── _secret ──────────────────────────────────────────────────────────
class TestSecret:
    def test_explicit_secret_returned(self, monkeypatch):
        monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "topsecret")
        gpu_relay._DEFAULT_SECRET_LOGGED = False
        assert gpu_relay._secret() == "topsecret"

    def test_prod_without_secret_raises(self, monkeypatch):
        monkeypatch.delenv("GPU_ENCODE_RELAY_SECRET", raising=False)
        monkeypatch.setenv("APP_ENV", "production")
        with pytest.raises(RuntimeError, match="GPU_ENCODE_RELAY_SECRET must be set"):
            gpu_relay._secret()

    def test_dev_without_secret_generates_ephemeral(self, monkeypatch, caplog):
        monkeypatch.delenv("GPU_ENCODE_RELAY_SECRET", raising=False)
        monkeypatch.setenv("APP_ENV", "development")
        gpu_relay._DEFAULT_SECRET_LOGGED = False
        with caplog.at_level("WARNING"):
            secret = gpu_relay._secret()
        assert len(secret) > 16
        assert "ephemeral dev token" in caplog.text
        # 第二次调用不再 log（_DEFAULT_SECRET_LOGGED=True）
        before = len(caplog.records)
        secret2 = gpu_relay._secret()
        assert secret2 == secret
        assert len(caplog.records) == before
        # 清理
        monkeypatch.delenv("GPU_ENCODE_RELAY_SECRET", raising=False)


# ── _safe_key ────────────────────────────────────────────────────────
class TestSafeKey:
    @pytest.mark.parametrize("bad", ["", "../etc", "a/b", "a\\b", ".", "..", "a b", "a%b"])
    def test_invalid_keys_rejected(self, bad):
        with pytest.raises(HTTPException) as ei:
            gpu_relay._safe_key(bad)
        assert ei.value.status_code == 400

    @pytest.mark.parametrize("good", ["abc123", "ABC-Def_01", "a" * 32])
    def test_valid_keys_accepted(self, good):
        assert gpu_relay._safe_key(good) == good

    def test_strips_whitespace(self):
        assert gpu_relay._safe_key("  abc  ") == "abc"


# ── _check_token ─────────────────────────────────────────────────────
class TestCheckToken:
    def test_missing_token_401(self):
        with pytest.raises(HTTPException) as ei:
            gpu_relay._check_token(None)
        assert ei.value.status_code == 401

    def test_wrong_token_401(self, monkeypatch):
        monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "correct")
        with pytest.raises(HTTPException) as ei:
            gpu_relay._check_token("wrong")
        assert ei.value.status_code == 401

    def test_correct_token_passes(self, monkeypatch):
        monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "correct")
        assert gpu_relay._check_token("correct") is None


# ── build_relay_* helpers ───────────────────────────────────────────
class TestBuildRelayUrls:
    def test_put_url(self):
        url = gpu_relay.build_relay_put_url("http://api.example.com/", "k1", "s")
        assert url == "http://api.example.com/api/v1/internal/gpu-relay/k1?token=s"

    def test_get_url_same_as_put(self):
        assert gpu_relay.build_relay_get_url("http://api", "k", "s") == \
               gpu_relay.build_relay_put_url("http://api", "k", "s")

    def test_generate_key_is_hex(self):
        k = gpu_relay.generate_key()
        assert len(k) == 32
        int(k, 16)  # valid hex


# ── PUT endpoint ────────────────────────────────────────────────────
@pytest.mark.asyncio
class TestPutObject:
    async def test_put_writes_file(self, tmp_path, monkeypatch):
        monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "s")
        monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path))
        # async request.stream 模拟
        async def _stream():
            yield b"chunk1"
            yield b"chunk2"
        req = mock.MagicMock(); req.stream = _stream
        resp = await gpu_relay.put_object(key="abc123", request=req, token="s")
        assert resp["ok"] is True
        assert resp["size"] == len(b"chunk1") + len(b"chunk2")
        p = tmp_path / "gpu_relay" / "abc123"
        assert p.read_bytes() == b"chunk1chunk2"
        # .part 临时文件应已 rename
        assert not p.with_suffix(p.suffix + ".part").exists()

    async def test_put_invalid_key_400(self, tmp_path, monkeypatch):
        monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "s")
        monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path))
        req = mock.MagicMock()
        with pytest.raises(HTTPException) as ei:
            await gpu_relay.put_object(key="../bad", request=req, token="s")
        assert ei.value.status_code == 400

    async def test_put_bad_token_401(self, tmp_path, monkeypatch):
        monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "correct")
        monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path))
        req = mock.MagicMock()
        with pytest.raises(HTTPException) as ei:
            await gpu_relay.put_object(key="abc", request=req, token="wrong")
        assert ei.value.status_code == 401

    async def test_put_write_error_cleans_tmp(self, tmp_path, monkeypatch):
        monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "s")
        monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path))
        async def _bad_stream():
            yield b"x"
            raise OSError("disk full")
        req = mock.MagicMock(); req.stream = _bad_stream
        with pytest.raises(HTTPException) as ei:
            await gpu_relay.put_object(key="abc", request=req, token="s")
        assert ei.value.status_code == 500
        # tmp 文件被清理
        part = tmp_path / "gpu_relay" / "abc.part"
        assert not part.exists()


# ── GET endpoint ────────────────────────────────────────────────────
@pytest.mark.asyncio
class TestGetObject:
    async def test_get_missing_404(self, tmp_path, monkeypatch):
        monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "s")
        monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path))
        with pytest.raises(HTTPException) as ei:
            await gpu_relay.get_object(key="nope", token="s")
        assert ei.value.status_code == 404

    async def test_get_returns_file(self, tmp_path, monkeypatch):
        monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "s")
        monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path))
        p = tmp_path / "gpu_relay" / "exist"
        p.parent.mkdir(parents=True, exist_ok=True)
        p.write_bytes(b"viddata")
        resp = await gpu_relay.get_object(key="exist", token="s")
        assert resp.media_type == "video/mp4"


# ── HEAD endpoint ───────────────────────────────────────────────────
@pytest.mark.asyncio
class TestHeadObject:
    async def test_head_missing_404(self, tmp_path, monkeypatch):
        monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "s")
        monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path))
        resp = await gpu_relay.head_object(key="nope", token="s")
        assert resp.status_code == 404

    async def test_head_returns_content_length(self, tmp_path, monkeypatch):
        monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "s")
        monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path))
        p = tmp_path / "gpu_relay" / "k"; p.parent.mkdir(parents=True, exist_ok=True); p.write_bytes(b"12345")
        resp = await gpu_relay.head_object(key="k", token="s")
        assert resp.status_code == 200
        assert resp.headers["Content-Length"] == "5"


# ── DELETE endpoint ──────────────────────────────────────────────────
@pytest.mark.asyncio
class TestDeleteObject:
    async def test_delete_existing(self, tmp_path, monkeypatch):
        monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "s")
        monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path))
        p = tmp_path / "gpu_relay" / "k"; p.parent.mkdir(parents=True, exist_ok=True); p.write_bytes(b"x")
        resp = await gpu_relay.delete_object(key="k", token="s")
        assert resp["ok"] is True
        assert not p.exists()

    async def test_delete_missing_is_noop(self, tmp_path, monkeypatch):
        monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "s")
        monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path))
        # 不存在时不应 404，返回 ok
        resp = await gpu_relay.delete_object(key="nope", token="s")
        assert resp["ok"] is True

    async def test_delete_unlink_error_500(self, tmp_path, monkeypatch):
        monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "s")
        monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path))
        p = tmp_path / "gpu_relay" / "k"; p.parent.mkdir(parents=True, exist_ok=True); p.write_bytes(b"x")
        with mock.patch.object(Path, "unlink", side_effect=OSError("perm denied")):
            with pytest.raises(HTTPException) as ei:
                await gpu_relay.delete_object(key="k", token="s")
        assert ei.value.status_code == 500