226 lines
7.2 KiB
Python
226 lines
7.2 KiB
Python
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from src.services import ai_control as ai_control_module
|
|
from src.services.ai_providers.gemini import GeminiProvider
|
|
from src.services.ai_rate_limiter import (
|
|
AIGenerationReservation,
|
|
get_gemini_pricing_policy,
|
|
)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _gemini_runtime_control_enabled(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("ENABLE_GEMINI", "true")
|
|
monkeypatch.setattr(
|
|
ai_control_module,
|
|
"is_provider_disabled",
|
|
AsyncMock(return_value=False),
|
|
)
|
|
|
|
|
|
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,
|
|
"thoughtsTokenCount": 7,
|
|
"totalTokenCount": 12,
|
|
},
|
|
}
|
|
|
|
|
|
@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_provider_options.return_value = {"temperature": 0.1}
|
|
|
|
provider = GeminiProvider()
|
|
provider._http_client = client
|
|
pricing = get_gemini_pricing_policy("gemini-2.5-flash-lite")
|
|
assert pricing is not None
|
|
reservation = AIGenerationReservation(
|
|
provider="gemini",
|
|
receipt_id="receipt-1",
|
|
allowed=True,
|
|
reason="allowed",
|
|
estimated_prompt_tokens=2,
|
|
estimated_completion_tokens=2048,
|
|
estimated_cost_usd=pricing.cost_usd(2, 2048),
|
|
date="2026-07-14",
|
|
model=pricing.model,
|
|
pricing_source=pricing.source,
|
|
pricing_version=pricing.version,
|
|
pricing_checked_at=pricing.checked_at,
|
|
input_usd_per_million=pricing.input_usd_per_million,
|
|
output_usd_per_million=pricing.output_usd_per_million,
|
|
)
|
|
limiter = MagicMock()
|
|
limiter.reserve_generation = AsyncMock(return_value=reservation)
|
|
limiter.finalize_generation = AsyncMock(return_value=True)
|
|
|
|
with patch(
|
|
"src.services.ai_providers.gemini.settings",
|
|
SimpleNamespace(
|
|
GEMINI_API_KEY="fake-key",
|
|
GEMINI_MODEL="gemini-2.5-flash-lite",
|
|
GEMINI_MAX_OUTPUT_TOKENS=2048,
|
|
),
|
|
), patch(
|
|
"src.services.ai_providers.gemini.get_model_registry",
|
|
return_value=registry,
|
|
), patch(
|
|
"src.services.ai_rate_limiter.get_ai_rate_limiter",
|
|
return_value=limiter,
|
|
):
|
|
context = {
|
|
"trace_id": "trace-security",
|
|
"run_id": "run-security",
|
|
"work_item_id": "work-security",
|
|
}
|
|
result = await provider.analyze("hello", context=context)
|
|
|
|
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.5-flash-lite:generateContent"
|
|
)
|
|
assert kwargs["headers"] == {"x-goog-api-key": "fake-key"}
|
|
assert "fake-key" not in url
|
|
assert "key=" not in url
|
|
limiter.reserve_generation.assert_awaited_once_with(
|
|
"gemini",
|
|
"hello",
|
|
model="gemini-2.5-flash-lite",
|
|
max_output_tokens=2048,
|
|
context=context,
|
|
)
|
|
limiter.finalize_generation.assert_awaited_once_with(
|
|
reservation,
|
|
success=True,
|
|
prompt_tokens=3,
|
|
completion_tokens=9,
|
|
)
|
|
assert result.tokens == 12
|
|
assert result.cost_usd == pricing.cost_usd(3, 9)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_authentication_rejection_writes_bounded_receipt_code() -> None:
|
|
request = httpx.Request(
|
|
"POST",
|
|
"https://generativelanguage.googleapis.com/v1beta/models/model:generateContent",
|
|
)
|
|
response = httpx.Response(status_code=401, request=request)
|
|
error = httpx.HTTPStatusError(
|
|
"authentication rejected",
|
|
request=request,
|
|
response=response,
|
|
)
|
|
client = MagicMock()
|
|
client.is_closed = False
|
|
client.post = AsyncMock(side_effect=error)
|
|
registry = MagicMock()
|
|
registry.get_provider_options.return_value = {"temperature": 0.1}
|
|
pricing = get_gemini_pricing_policy("gemini-2.5-flash-lite")
|
|
assert pricing is not None
|
|
reservation = AIGenerationReservation(
|
|
provider="gemini",
|
|
receipt_id="receipt-auth-rejected",
|
|
allowed=True,
|
|
reason="allowed",
|
|
estimated_prompt_tokens=1,
|
|
estimated_completion_tokens=1,
|
|
estimated_cost_usd=pricing.cost_usd(1, 1),
|
|
date="2026-07-14",
|
|
model=pricing.model,
|
|
pricing_source=pricing.source,
|
|
pricing_version=pricing.version,
|
|
pricing_checked_at=pricing.checked_at,
|
|
input_usd_per_million=pricing.input_usd_per_million,
|
|
output_usd_per_million=pricing.output_usd_per_million,
|
|
)
|
|
limiter = MagicMock()
|
|
limiter.reserve_generation = AsyncMock(return_value=reservation)
|
|
limiter.finalize_generation = AsyncMock(return_value=True)
|
|
provider = GeminiProvider()
|
|
provider._http_client = client
|
|
|
|
with patch(
|
|
"src.services.ai_providers.gemini.settings",
|
|
SimpleNamespace(
|
|
GEMINI_API_KEY="configured-test-value",
|
|
GEMINI_MODEL="gemini-2.5-flash-lite",
|
|
GEMINI_MAX_OUTPUT_TOKENS=1,
|
|
),
|
|
), patch(
|
|
"src.services.ai_providers.gemini.get_model_registry",
|
|
return_value=registry,
|
|
), patch(
|
|
"src.services.ai_rate_limiter.get_ai_rate_limiter",
|
|
return_value=limiter,
|
|
):
|
|
result = await provider.analyze(
|
|
"test authentication",
|
|
context={
|
|
"trace_id": "trace-auth",
|
|
"run_id": "run-auth",
|
|
"work_item_id": "work-auth",
|
|
},
|
|
)
|
|
|
|
assert result.success is False
|
|
limiter.finalize_generation.assert_awaited_once_with(
|
|
reservation,
|
|
success=False,
|
|
error_code="authentication_rejected",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_health_requires_durable_authenticated_verified_receipt() -> None:
|
|
limiter = MagicMock()
|
|
limiter.get_provider_authentication_status = AsyncMock(
|
|
return_value={
|
|
"configured": True,
|
|
"authenticated": None,
|
|
"authentication_verified": False,
|
|
}
|
|
)
|
|
provider = GeminiProvider()
|
|
|
|
with patch(
|
|
"src.services.ai_providers.gemini.settings",
|
|
SimpleNamespace(GEMINI_API_KEY="configured-test-value"),
|
|
), patch(
|
|
"src.services.ai_providers.gemini.is_provider_enabled_by_env",
|
|
return_value=True,
|
|
), patch(
|
|
"src.services.ai_rate_limiter.get_ai_rate_limiter",
|
|
return_value=limiter,
|
|
):
|
|
assert await provider.health_check() is False
|
|
limiter.get_provider_authentication_status.return_value = {
|
|
"configured": True,
|
|
"authenticated": True,
|
|
"authentication_verified": True,
|
|
}
|
|
assert await provider.health_check() is True
|