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