Files
awoooi/apps/api/tests/test_gemini_provider_security.py
Your Name d845d53257
All checks were successful
CD Pipeline / build-and-deploy (push) Successful in 15m5s
fix(security): keep Gemini key out of request URLs
2026-04-29 22:56:12 +08:00

53 lines
1.6 KiB
Python

from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from src.services.ai_providers.gemini import GeminiProvider
class _FakeGeminiResponse:
def raise_for_status(self) -> None:
return None
def json(self) -> dict:
return {
"candidates": [{"content": {"parts": [{"text": "{\"ok\": true}"}]}}],
"usageMetadata": {
"promptTokenCount": 3,
"candidatesTokenCount": 2,
"totalTokenCount": 5,
},
}
@pytest.mark.asyncio
async def test_gemini_provider_uses_header_auth_not_query_string():
client = MagicMock()
client.is_closed = False
client.post = AsyncMock(return_value=_FakeGeminiResponse())
registry = MagicMock()
registry.get_model.return_value = "gemini-2.0-flash"
registry.get_provider_options.return_value = {"temperature": 0.1}
provider = GeminiProvider()
provider._http_client = client
with patch(
"src.services.ai_providers.gemini.settings",
SimpleNamespace(GEMINI_API_KEY="fake-key"),
), patch("src.services.ai_providers.gemini.get_model_registry", return_value=registry):
result = await provider.analyze("hello")
assert result.success is True
client.post.assert_awaited_once()
url = client.post.await_args.args[0]
kwargs = client.post.await_args.kwargs
assert url == "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent"
assert kwargs["headers"] == {"x-goog-api-key": "fake-key"}
assert "fake-key" not in url
assert "key=" not in url