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
457 lines
14 KiB
Python
457 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import hashlib
|
|
import json
|
|
from datetime import UTC, datetime, timedelta
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from src.services import ai_router as ai_router_module
|
|
from src.services.ai_providers.interfaces import AIResult
|
|
from src.services.ai_router import AIProviderRegistry, AIRouterExecutor
|
|
from src.services.openclaw import OpenClawService
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _explicitly_enable_gemini_environment(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setenv("ENABLE_GEMINI", "true")
|
|
|
|
|
|
class _FakeRedis:
|
|
def __init__(
|
|
self,
|
|
cached_provider: str,
|
|
*,
|
|
gemini_disabled: bool = False,
|
|
durable_receipts: dict[str, dict[str, Any]] | None = None,
|
|
) -> None:
|
|
self.cached_provider = cached_provider
|
|
self.gemini_disabled = gemini_disabled
|
|
self.durable_receipts = durable_receipts or {}
|
|
self.set_calls: list[tuple[str, str, int | None]] = []
|
|
|
|
async def get(self, key: str) -> str | None:
|
|
if key.startswith("ai:control:disabled:"):
|
|
return "1" if self.gemini_disabled else "0"
|
|
if key.startswith("ai:ollama:unavailable_receipt:"):
|
|
receipt_id = key.removeprefix("ai:ollama:unavailable_receipt:")
|
|
receipt = self.durable_receipts.get(receipt_id)
|
|
return json.dumps(receipt) if receipt else None
|
|
return json.dumps({
|
|
"response": '{"provider":"stale"}',
|
|
"provider": self.cached_provider,
|
|
})
|
|
|
|
async def set(self, key: str, value: str, ex: int | None = None) -> None:
|
|
self.set_calls.append((key, value, ex))
|
|
|
|
async def eval(self, *_args: object) -> list[object]:
|
|
if self.gemini_disabled:
|
|
return [0, "paired_paid_state_not_enabled", ""]
|
|
run_id = str(_args[-1] or "")
|
|
if not run_id:
|
|
return [0, "run_owned_canary_not_admitted", ""]
|
|
return [
|
|
1,
|
|
"run_owned_canary_admitted",
|
|
f"{run_id}.attempt-fixture",
|
|
]
|
|
|
|
|
|
class _FakeProvider:
|
|
name = "ollama_gcp_a"
|
|
privacy_level = "local"
|
|
is_enabled = True
|
|
capabilities = {"rca", "chat"}
|
|
|
|
def __init__(self) -> None:
|
|
self.calls = 0
|
|
|
|
async def analyze(self, prompt: str, context: dict[str, Any] | None = None) -> AIResult:
|
|
self.calls += 1
|
|
return AIResult(
|
|
raw_response='{"provider":"fresh_ollama"}',
|
|
success=True,
|
|
provider=self.name,
|
|
)
|
|
|
|
|
|
class _FailingLocalProvider:
|
|
privacy_level = "local"
|
|
is_enabled = True
|
|
capabilities = {"rca", "chat"}
|
|
|
|
def __init__(self, name: str) -> None:
|
|
self.name = name
|
|
self.calls = 0
|
|
|
|
async def analyze(self, prompt: str, context: dict[str, Any] | None = None) -> AIResult:
|
|
self.calls += 1
|
|
return AIResult(raw_response="", success=False, provider=self.name, error="forced failure")
|
|
|
|
|
|
class _CloudProvider:
|
|
name = "gemini"
|
|
privacy_level = "cloud"
|
|
is_enabled = True
|
|
capabilities = {"rca", "chat"}
|
|
|
|
def __init__(self) -> None:
|
|
self.calls = 0
|
|
|
|
async def analyze(self, prompt: str, context: dict[str, Any] | None = None) -> AIResult:
|
|
self.calls += 1
|
|
return AIResult(raw_response='{"provider":"gemini"}', success=True, provider=self.name)
|
|
|
|
|
|
class _CloudOllamaProvider(_CloudProvider):
|
|
def __init__(self, name: str) -> None:
|
|
super().__init__()
|
|
self.name = name
|
|
|
|
|
|
class _SuccessfulLocalProvider(_FailingLocalProvider):
|
|
async def analyze(
|
|
self,
|
|
prompt: str,
|
|
context: dict[str, Any] | None = None,
|
|
) -> AIResult:
|
|
self.calls += 1
|
|
return AIResult(
|
|
raw_response='{"provider":"ollama_local"}',
|
|
success=True,
|
|
provider=self.name,
|
|
)
|
|
|
|
|
|
class _SlowLocalProvider(_FailingLocalProvider):
|
|
async def analyze(
|
|
self,
|
|
prompt: str,
|
|
context: dict[str, Any] | None = None,
|
|
) -> AIResult:
|
|
self.calls += 1
|
|
await asyncio.sleep(1)
|
|
return AIResult(
|
|
raw_response='{"provider":"late"}',
|
|
success=True,
|
|
provider=self.name,
|
|
)
|
|
|
|
|
|
def _paid_alert_context(**overrides: Any) -> dict[str, Any]:
|
|
service = object.__new__(OpenClawService)
|
|
context = service._prepare_paid_cloud_route_context(
|
|
{
|
|
"alert_type": "HostHighCpuLoad",
|
|
"source": "test-alert-producer",
|
|
"target_resource": "node-110",
|
|
**overrides,
|
|
}
|
|
)
|
|
assert context is not None
|
|
context.setdefault("run_id", "run-cache-policy-1")
|
|
return context
|
|
|
|
|
|
def _synthetic_cloud_context(prompt: str, **overrides: Any) -> dict[str, Any]:
|
|
return {
|
|
"intent_hint": "diagnose",
|
|
"synthetic_validation": True,
|
|
"data_classification": "sanitized",
|
|
"synthetic_prompt_sha256": hashlib.sha256(prompt.encode()).hexdigest(),
|
|
"infrastructure_remediation_write_allowed": False,
|
|
"run_id": "AIA-SRE-013-test",
|
|
**overrides,
|
|
}
|
|
|
|
|
|
def _durable_unavailable_receipt(provider: str, run_id: str) -> dict[str, Any]:
|
|
observed_at = datetime.now(UTC)
|
|
return {
|
|
"schema_version": "ollama_unavailable_receipt_v1",
|
|
"receipt_id": f"ollama-unavailable:{run_id}:{provider}",
|
|
"source": "ollama_failover_manager",
|
|
"provider": provider,
|
|
"run_id": run_id,
|
|
"status": "bounded_unavailable",
|
|
"check_completed": True,
|
|
"reason": "endpoint_unreachable",
|
|
"observed_at": observed_at.isoformat(),
|
|
"expires_at": (observed_at + timedelta(seconds=120)).isoformat(),
|
|
"durable_write_ack": True,
|
|
"verifier_status": "verified",
|
|
"verified_by": "ollama_failover_manager.health_probe",
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_executor_skips_cached_cloud_provider_when_ollama_lane_is_required(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
fake_redis = _FakeRedis(cached_provider="gemini")
|
|
fake_provider = _FakeProvider()
|
|
registry = AIProviderRegistry()
|
|
registry.register(fake_provider)
|
|
|
|
monkeypatch.setattr(ai_router_module._settings, "MOCK_MODE", False)
|
|
monkeypatch.setattr("src.core.redis_client.get_redis", lambda: fake_redis)
|
|
|
|
result = await AIRouterExecutor(registry).execute(
|
|
prompt="diagnose alert",
|
|
provider_order=["ollama_gcp_a", "ollama_gcp_b", "ollama_local", "gemini"],
|
|
context=_paid_alert_context(intent_hint="diagnose"),
|
|
)
|
|
|
|
assert result.provider == "ollama_gcp_a"
|
|
assert result.raw_response == '{"provider":"fresh_ollama"}'
|
|
assert fake_provider.calls == 1
|
|
# GCP Ollama is a remote sanitized-cloud hop; its response is deliberately
|
|
# non-cacheable even though it is the first free provider.
|
|
assert fake_redis.set_calls == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_executor_applies_gemini_runtime_disable_before_cache_lookup(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
fake_redis = _FakeRedis(cached_provider="gemini", gemini_disabled=True)
|
|
fake_provider = _FakeProvider()
|
|
registry = AIProviderRegistry()
|
|
registry.register(fake_provider)
|
|
|
|
monkeypatch.setattr(ai_router_module._settings, "MOCK_MODE", False)
|
|
monkeypatch.setattr("src.core.redis_client.get_redis", lambda: fake_redis)
|
|
|
|
result = await AIRouterExecutor(registry).execute(
|
|
prompt="non-alert summary",
|
|
provider_order=["ollama_gcp_a", "gemini"],
|
|
context=_synthetic_cloud_context(
|
|
"non-alert summary",
|
|
intent_hint="summary",
|
|
),
|
|
)
|
|
|
|
assert result.provider == "ollama_gcp_a"
|
|
assert result.from_cache is False
|
|
assert fake_provider.calls == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_executor_applies_gemini_environment_disable_before_cache_lookup(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
fake_redis = _FakeRedis(cached_provider="gemini")
|
|
fake_provider = _FakeProvider()
|
|
registry = AIProviderRegistry()
|
|
registry.register(fake_provider)
|
|
|
|
monkeypatch.delenv("ENABLE_GEMINI", raising=False)
|
|
monkeypatch.setattr(ai_router_module._settings, "MOCK_MODE", False)
|
|
monkeypatch.setattr("src.core.redis_client.get_redis", lambda: fake_redis)
|
|
|
|
result = await AIRouterExecutor(registry).execute(
|
|
prompt="non-alert summary",
|
|
provider_order=["ollama_gcp_a", "gemini"],
|
|
context=_synthetic_cloud_context(
|
|
"non-alert summary",
|
|
intent_hint="summary",
|
|
),
|
|
)
|
|
|
|
assert result.provider == "ollama_gcp_a"
|
|
assert result.from_cache is False
|
|
assert fake_provider.calls == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_executor_allows_cached_ollama_provider_for_ollama_lane(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
fake_redis = _FakeRedis(cached_provider="ollama")
|
|
fake_provider = _FakeProvider()
|
|
registry = AIProviderRegistry()
|
|
registry.register(fake_provider)
|
|
|
|
monkeypatch.setattr(ai_router_module._settings, "MOCK_MODE", False)
|
|
monkeypatch.setattr("src.core.redis_client.get_redis", lambda: fake_redis)
|
|
|
|
result = await AIRouterExecutor(registry).execute(
|
|
prompt="diagnose alert",
|
|
provider_order=["ollama_gcp_a", "ollama_gcp_b", "ollama_local"],
|
|
context={"intent_hint": "diagnose", "alert_type": "HostHighCpuLoad"},
|
|
)
|
|
|
|
assert result.provider == "ollama"
|
|
assert result.from_cache is True
|
|
assert fake_provider.calls == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_alert_cloud_backup_waits_for_real_ollama_local_attempt(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
fake_redis = _FakeRedis(cached_provider="none")
|
|
gcp_a = _FailingLocalProvider("ollama_gcp_a")
|
|
gemini = _CloudProvider()
|
|
registry = AIProviderRegistry()
|
|
registry.register(gcp_a)
|
|
registry.register(gemini)
|
|
|
|
monkeypatch.setattr(ai_router_module._settings, "MOCK_MODE", False)
|
|
monkeypatch.setattr("src.core.redis_client.get_redis", lambda: fake_redis)
|
|
|
|
result = await AIRouterExecutor(registry).execute(
|
|
prompt="diagnose alert",
|
|
provider_order=["ollama_gcp_a", "gemini"],
|
|
context=_paid_alert_context(
|
|
intent_hint="diagnose",
|
|
alert_requires_ollama_before_cloud=True,
|
|
),
|
|
)
|
|
|
|
assert result.success is False
|
|
assert gcp_a.calls == 1
|
|
assert gemini.calls == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_require_local_skips_both_gcp_ollama_hops_and_uses_host111(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
fake_redis = _FakeRedis(cached_provider="none")
|
|
gcp_a = _CloudOllamaProvider("ollama_gcp_a")
|
|
gcp_b = _CloudOllamaProvider("ollama_gcp_b")
|
|
local = _SuccessfulLocalProvider("ollama_local")
|
|
registry = AIProviderRegistry()
|
|
registry.register(gcp_a)
|
|
registry.register(gcp_b)
|
|
registry.register(local)
|
|
|
|
monkeypatch.setattr(ai_router_module._settings, "MOCK_MODE", False)
|
|
monkeypatch.setattr("src.core.redis_client.get_redis", lambda: fake_redis)
|
|
|
|
result = await AIRouterExecutor(registry).execute(
|
|
prompt="local-only diagnosis",
|
|
provider_order=["ollama_gcp_a", "ollama_gcp_b", "ollama_local"],
|
|
context={"intent_hint": "diagnose"},
|
|
require_local=True,
|
|
)
|
|
|
|
assert result.success is True
|
|
assert result.provider == "ollama_local"
|
|
assert gcp_a.calls == 0
|
|
assert gcp_b.calls == 0
|
|
assert local.calls == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_alert_cloud_backup_allowed_after_ollama_local_attempt(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
fake_redis = _FakeRedis(cached_provider="none")
|
|
gcp_a = _FailingLocalProvider("ollama_gcp_a")
|
|
gcp_b = _FailingLocalProvider("ollama_gcp_b")
|
|
local = _FailingLocalProvider("ollama_local")
|
|
gemini = _CloudProvider()
|
|
registry = AIProviderRegistry()
|
|
registry.register(gcp_a)
|
|
registry.register(gcp_b)
|
|
registry.register(local)
|
|
registry.register(gemini)
|
|
|
|
monkeypatch.setattr(ai_router_module._settings, "MOCK_MODE", False)
|
|
monkeypatch.setattr("src.core.redis_client.get_redis", lambda: fake_redis)
|
|
|
|
result = await AIRouterExecutor(registry).execute(
|
|
prompt="diagnose alert",
|
|
provider_order=["ollama_gcp_a", "ollama_local", "gemini"],
|
|
context=_paid_alert_context(
|
|
intent_hint="diagnose",
|
|
alert_requires_ollama_before_cloud=True,
|
|
),
|
|
)
|
|
|
|
assert result.success is True
|
|
assert result.provider == "gemini"
|
|
assert gcp_a.calls == 1
|
|
assert gcp_b.calls == 1
|
|
assert local.calls == 1
|
|
assert gemini.calls == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_alert_provider_timeout_allows_next_fallback(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
durable_receipts = {
|
|
receipt["receipt_id"]: receipt
|
|
for receipt in (
|
|
_durable_unavailable_receipt(
|
|
"ollama_gcp_a",
|
|
"run-cache-policy-1",
|
|
),
|
|
_durable_unavailable_receipt(
|
|
"ollama_gcp_b",
|
|
"run-cache-policy-1",
|
|
),
|
|
)
|
|
}
|
|
fake_redis = _FakeRedis(
|
|
cached_provider="none",
|
|
durable_receipts=durable_receipts,
|
|
)
|
|
local = _SlowLocalProvider("ollama_local")
|
|
gemini = _CloudProvider()
|
|
registry = AIProviderRegistry()
|
|
registry.register(local)
|
|
registry.register(gemini)
|
|
|
|
monkeypatch.setattr(ai_router_module._settings, "MOCK_MODE", False)
|
|
monkeypatch.setattr("src.core.redis_client.get_redis", lambda: fake_redis)
|
|
|
|
result = await AIRouterExecutor(registry).execute(
|
|
prompt="diagnose alert",
|
|
provider_order=["ollama_local", "gemini"],
|
|
context=_paid_alert_context(
|
|
intent_hint="alert_triage",
|
|
alert_requires_ollama_before_cloud=True,
|
|
provider_timeout_seconds=0.01,
|
|
ollama_unavailable_receipts=list(durable_receipts.values()),
|
|
),
|
|
)
|
|
|
|
assert result.success is True
|
|
assert result.provider == "gemini"
|
|
assert local.calls == 1
|
|
assert gemini.calls == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_provider_unsuccessful_results_count_toward_circuit_breaker(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
fake_redis = _FakeRedis(cached_provider="none")
|
|
provider = _FailingLocalProvider("ollama_gcp_a")
|
|
registry = AIProviderRegistry()
|
|
registry.register(provider)
|
|
executor = AIRouterExecutor(registry)
|
|
|
|
monkeypatch.setattr(ai_router_module._settings, "MOCK_MODE", False)
|
|
monkeypatch.setattr("src.core.redis_client.get_redis", lambda: fake_redis)
|
|
|
|
for index in range(6):
|
|
result = await executor.execute(
|
|
prompt=f"failure {index}",
|
|
provider_order=["ollama_gcp_a"],
|
|
context=_synthetic_cloud_context(f"failure {index}"),
|
|
)
|
|
assert result.success is False
|
|
|
|
assert provider.calls == 5
|
|
assert executor._get_circuit_breaker("ollama_gcp_a").is_open() is True
|