Files
awoooi/apps/api/tests/test_ollama_provider_endpoints.py
ogt 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
feat(sre): enforce typed controlled automation
2026-07-16 17:32:01 +08:00

272 lines
9.0 KiB
Python

from __future__ import annotations
import hashlib
from typing import Any
from unittest.mock import AsyncMock
import pytest
from src.services.ai_providers import ollama as ollama_module
from src.services.ai_providers.ollama import OllamaGcpBProvider, OllamaProvider
class _FakeRegistry:
def get_model(self, provider: str, use_case: str) -> str:
return "qwen3:14b"
def get_provider_options(self, provider: str) -> dict[str, Any]:
return {"num_predict": 32, "temperature": 0.1, "top_p": 0.9}
class _FakeResponse:
status_code = 200
def raise_for_status(self) -> None:
return None
def json(self) -> dict[str, Any]:
return {
"response": '{"summary":"ok"}',
"eval_count": 4,
"prompt_eval_count": 3,
}
class _FakeClient:
def __init__(self) -> None:
self.posted_urls: list[str] = []
self.posted_payloads: list[dict[str, Any]] = []
self.checked_urls: list[str] = []
async def post(self, url: str, **kwargs: Any) -> _FakeResponse:
self.posted_urls.append(url)
self.posted_payloads.append(kwargs.get("json", {}))
return _FakeResponse()
async def get(self, url: str, **kwargs: Any) -> _FakeResponse:
self.checked_urls.append(url)
return _FakeResponse()
def _synthetic_context(prompt: str, **overrides: Any) -> dict[str, Any]:
return {
"data_classification": "sanitized",
"trace_id": "trace-test",
"run_id": "run-test",
"work_item_id": "AIA-SRE-013",
"synthetic_validation": True,
"synthetic_prompt_sha256": hashlib.sha256(prompt.encode()).hexdigest(),
"cloud_transport_receipts": [
{"provider": "ollama_gcp_a", "receipt_id": "verified-a"},
{"provider": "ollama_gcp_b", "receipt_id": "verified-b"},
],
"infrastructure_remediation_write_allowed": False,
**overrides,
}
async def _injected_transport_verifier(**kwargs: Any) -> str | None:
receipts = (kwargs.get("context") or {}).get("cloud_transport_receipts") or []
if not any(item.get("provider") == kwargs["provider"] for item in receipts):
return "gcp_cloud_transport_receipt_missing"
return None
@pytest.mark.asyncio
async def test_ollama_gcp_b_analyze_uses_secondary_url(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(ollama_module, "get_model_registry", lambda: _FakeRegistry())
monkeypatch.setattr(ollama_module.settings, "OLLAMA_URL", "http://primary:11435")
monkeypatch.setattr(
ollama_module.settings,
"OLLAMA_SECONDARY_URL",
"http://secondary:11436",
)
monkeypatch.setattr(ollama_module.settings, "ALERT_OLLAMA_MODEL", "qwen3:14b")
client = _FakeClient()
provider = OllamaGcpBProvider(
transport_receipt_verifier=_injected_transport_verifier
)
async def _get_client() -> _FakeClient:
return client
monkeypatch.setattr(provider, "_get_client", _get_client)
result = await provider.analyze(
"diagnose",
context=_synthetic_context("diagnose", task_type="diagnose"),
)
assert result.success is True
assert result.provider == "ollama_gcp_b"
assert client.posted_urls == ["http://secondary:11436/api/generate"]
assert client.posted_payloads[0]["model"] == "qwen3:14b"
@pytest.mark.asyncio
async def test_ollama_gcp_a_allows_heavy_diagnose_model(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(ollama_module, "get_model_registry", lambda: _FakeRegistry())
monkeypatch.setattr(ollama_module.settings, "OLLAMA_URL", "http://primary:11435")
monkeypatch.setattr(ollama_module.settings, "OLLAMA_SECONDARY_URL", "http://secondary:11436")
monkeypatch.setattr(ollama_module.settings, "ALERT_OLLAMA_MODEL", "qwen3:14b")
client = _FakeClient()
provider = OllamaProvider(transport_receipt_verifier=_injected_transport_verifier)
async def _get_client() -> _FakeClient:
return client
monkeypatch.setattr(provider, "_get_client", _get_client)
result = await provider.analyze(
"diagnose",
context=_synthetic_context("diagnose", task_type="diagnose"),
)
assert result.success is True
assert client.posted_urls == ["http://primary:11435/api/generate"]
assert client.posted_payloads[0]["model"] == "qwen3:14b"
@pytest.mark.asyncio
async def test_ollama_gcp_a_coerces_non_diagnosis_heavy_model_to_health_model(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(ollama_module, "get_model_registry", lambda: _FakeRegistry())
monkeypatch.setattr(ollama_module.settings, "OLLAMA_URL", "http://primary:11435")
monkeypatch.setattr(ollama_module.settings, "OLLAMA_SECONDARY_URL", "http://secondary:11436")
monkeypatch.setattr(ollama_module.settings, "OLLAMA_HEALTH_CHECK_MODEL", "gemma3:4b")
client = _FakeClient()
provider = OllamaProvider(transport_receipt_verifier=_injected_transport_verifier)
async def _get_client() -> _FakeClient:
return client
monkeypatch.setattr(provider, "_get_client", _get_client)
result = await provider.analyze(
"background summary",
context=_synthetic_context("background summary", task_type="summary"),
)
assert result.success is True
assert client.posted_urls == ["http://primary:11435/api/generate"]
assert client.posted_payloads[0]["model"] == "gemma3:4b"
@pytest.mark.asyncio
async def test_ollama_gcp_a_can_explicitly_allow_heavy_model(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(ollama_module, "get_model_registry", lambda: _FakeRegistry())
monkeypatch.setattr(ollama_module.settings, "OLLAMA_URL", "http://primary:11435")
monkeypatch.setattr(ollama_module.settings, "OLLAMA_SECONDARY_URL", "http://secondary:11436")
monkeypatch.setattr(ollama_module.settings, "ALERT_OLLAMA_MODEL", "qwen3:14b")
client = _FakeClient()
provider = OllamaProvider(transport_receipt_verifier=_injected_transport_verifier)
async def _get_client() -> _FakeClient:
return client
monkeypatch.setattr(provider, "_get_client", _get_client)
result = await provider.analyze(
"deep diagnose",
context=_synthetic_context(
"deep diagnose",
task_type="diagnose",
allow_gcp_heavy_model=True,
),
)
assert result.success is True
assert client.posted_payloads[0]["model"] == "qwen3:14b"
@pytest.mark.asyncio
async def test_gcp_ollama_rejects_unreceipted_prompt_before_network_call(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(ollama_module.settings, "OLLAMA_URL", "http://primary:11435")
monkeypatch.setattr(
ollama_module.settings,
"OLLAMA_SECONDARY_URL",
"http://secondary:11436",
)
client = _FakeClient()
provider = OllamaProvider(transport_receipt_verifier=_injected_transport_verifier)
async def _get_client() -> _FakeClient:
return client
monkeypatch.setattr(provider, "_get_client", _get_client)
result = await provider.analyze(
"raw operational log",
context={"data_classification": "sanitized"},
)
assert result.success is False
assert result.error == "gcp_cloud_transport_receipt_missing"
assert result.audit_metadata["generation_attempted"] is False
assert client.posted_urls == []
@pytest.mark.asyncio
async def test_gcp_ollama_blocks_tool_loop_before_tool_or_network_execution(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(ollama_module.settings, "OLLAMA_URL", "http://primary:11435")
monkeypatch.setattr(
ollama_module.settings,
"OLLAMA_SECONDARY_URL",
"http://secondary:11436",
)
client = _FakeClient()
provider = OllamaProvider(transport_receipt_verifier=_injected_transport_verifier)
tool_executor = AsyncMock()
async def _get_client() -> _FakeClient:
return client
monkeypatch.setattr(provider, "_get_client", _get_client)
prompt = "sanitized investigation"
result = await provider.analyze_with_tools(
prompt,
available_tools=[object()], # type: ignore[list-item]
tool_executor=tool_executor,
context=_synthetic_context(prompt, task_type="diagnose"),
)
assert result.success is False
assert result.error == "gcp_tool_loop_blocked_unverified_tool_results"
assert result.audit_metadata == {
"generation_attempted": False,
"tool_execution_attempted": False,
"transport_boundary": "public_http_degraded",
}
tool_executor.assert_not_awaited()
assert client.posted_urls == []
@pytest.mark.asyncio
async def test_ollama_base_provider_health_uses_endpoint_hook(
monkeypatch: pytest.MonkeyPatch,
) -> None:
client = _FakeClient()
provider = OllamaProvider()
monkeypatch.setattr(provider, "_endpoint_url", lambda: "http://primary:11435")
async def _get_client() -> _FakeClient:
return client
monkeypatch.setattr(provider, "_get_client", _get_client)
assert await provider.health_check() is True
assert client.checked_urls == ["http://primary:11435/api/tags"]