541a32a9cb
Some checks failed
CD Pipeline / workflow-shape (push) Successful in 0s
CD Pipeline / cancel-stale-cd (push) Has been skipped
CD Pipeline / tests (push) Successful in 2m38s
CD Pipeline / build-and-deploy (push) Failing after 24m19s
CD Pipeline / post-deploy-checks (push) Has been skipped
313 lines
10 KiB
Python
313 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
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,
|
|
)
|
|
|
|
|
|
def _gemini_canary_run_id(percent: int) -> tuple[str, int]:
|
|
for index in range(1000):
|
|
run_id = f"gemini-canary-{index}"
|
|
bucket = int(hashlib.sha256(run_id.encode()).hexdigest()[:8], 16) % 100
|
|
if bucket < percent:
|
|
return run_id, bucket
|
|
raise AssertionError("missing deterministic canary bucket")
|
|
|
|
|
|
@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_shadow_and_unselected_canary_perform_no_external_call() -> None:
|
|
provider = GeminiProvider()
|
|
client = MagicMock()
|
|
client.is_closed = False
|
|
client.post = AsyncMock()
|
|
provider._http_client = client
|
|
|
|
with patch(
|
|
"src.services.ai_providers.gemini.settings",
|
|
SimpleNamespace(
|
|
GEMINI_API_KEY="configured-test-value",
|
|
GEMINI_EXECUTION_MODE="shadow",
|
|
GEMINI_CANARY_PERCENT=0,
|
|
),
|
|
):
|
|
shadow = await provider.analyze(
|
|
"synthetic",
|
|
context={"data_classification": "sanitized", "run_id": "run-shadow"},
|
|
)
|
|
|
|
selected_run, selected_bucket = _gemini_canary_run_id(5)
|
|
unselected_run = next(
|
|
f"gemini-unselected-{index}"
|
|
for index in range(1000)
|
|
if int(
|
|
hashlib.sha256(f"gemini-unselected-{index}".encode()).hexdigest()[:8],
|
|
16,
|
|
)
|
|
% 100
|
|
>= 5
|
|
)
|
|
with patch(
|
|
"src.services.ai_providers.gemini.settings",
|
|
SimpleNamespace(
|
|
GEMINI_API_KEY="configured-test-value",
|
|
GEMINI_EXECUTION_MODE="canary",
|
|
GEMINI_CANARY_PERCENT=5,
|
|
),
|
|
):
|
|
unselected = await provider.analyze(
|
|
"synthetic",
|
|
context={
|
|
"data_classification": "sanitized",
|
|
"run_id": unselected_run,
|
|
},
|
|
)
|
|
allowed, reason, mode, bucket = provider._execution_mode_gate(
|
|
{"run_id": selected_run}
|
|
)
|
|
|
|
assert shadow.error == "gemini_shadow_metadata_only"
|
|
assert unselected.error == "gemini_canary_not_selected"
|
|
assert (allowed, reason, mode, bucket) == (
|
|
True,
|
|
"gemini_canary_selected",
|
|
"canary",
|
|
selected_bucket,
|
|
)
|
|
client.post.assert_not_called()
|
|
|
|
|
|
@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,
|
|
GEMINI_EXECUTION_MODE="production",
|
|
GEMINI_CANARY_PERCENT=0,
|
|
),
|
|
), 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",
|
|
"data_classification": "sanitized",
|
|
}
|
|
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(
|
|
"provider body SECRET_PROVIDER_PAYLOAD_MUST_NOT_BE_LOGGED",
|
|
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,
|
|
GEMINI_EXECUTION_MODE="production",
|
|
GEMINI_CANARY_PERCENT=0,
|
|
),
|
|
), 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,
|
|
), patch("src.services.ai_providers.gemini.logger.warning") as warning:
|
|
result = await provider.analyze(
|
|
"test authentication",
|
|
context={
|
|
"trace_id": "trace-auth",
|
|
"run_id": "run-auth",
|
|
"work_item_id": "work-auth",
|
|
"data_classification": "sanitized",
|
|
},
|
|
)
|
|
|
|
assert result.success is False
|
|
assert result.error == "authentication_rejected"
|
|
limiter.finalize_generation.assert_awaited_once_with(
|
|
reservation,
|
|
success=False,
|
|
error_code="authentication_rejected",
|
|
)
|
|
assert "SECRET_PROVIDER_PAYLOAD_MUST_NOT_BE_LOGGED" not in str(warning.call_args)
|
|
assert warning.call_args.kwargs["error_code"] == "authentication_rejected"
|
|
assert warning.call_args.kwargs["error_type"] == "HTTPStatusError"
|
|
assert "error" not in warning.call_args.kwargs
|
|
|
|
|
|
@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",
|
|
GEMINI_EXECUTION_MODE="production",
|
|
GEMINI_CANARY_PERCENT=0,
|
|
),
|
|
), 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
|