232 lines
9.0 KiB
Python
232 lines
9.0 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
|
|
from src.services import ai_control as ai_control_module
|
|
from src.services import ai_router as ai_router_module
|
|
from src.services import openclaw as openclaw_module
|
|
from src.services.ai_router import AIProviderEnum
|
|
from src.services.intent_classifier import IntentType
|
|
from src.services.openclaw import OpenClawService
|
|
|
|
|
|
@dataclass
|
|
class _FakeEndpoint:
|
|
provider_name: str
|
|
url: str = "http://example.test"
|
|
|
|
|
|
class _FakeRoute:
|
|
def all_endpoints_in_order(self) -> list[_FakeEndpoint]:
|
|
return [
|
|
_FakeEndpoint("ollama_gcp_a"),
|
|
_FakeEndpoint("ollama_gcp_b"),
|
|
_FakeEndpoint("ollama_local"),
|
|
_FakeEndpoint("gemini", ""),
|
|
]
|
|
|
|
|
|
class _FakeFailoverManager:
|
|
def __init__(self) -> None:
|
|
self.task_types: list[str] = []
|
|
|
|
async def select_provider(self, task_type: str = "general") -> _FakeRoute:
|
|
self.task_types.append(task_type)
|
|
return _FakeRoute()
|
|
|
|
|
|
class _UnorderedFailoverManager:
|
|
async def select_provider(self, task_type: str = "general") -> SimpleNamespace:
|
|
return SimpleNamespace(
|
|
all_endpoints_in_order=lambda: [
|
|
_FakeEndpoint("ollama_local"),
|
|
_FakeEndpoint("gemini"),
|
|
_FakeEndpoint("ollama_gcp_b"),
|
|
_FakeEndpoint("ollama_gcp_a"),
|
|
],
|
|
)
|
|
|
|
|
|
class _FakeRouter:
|
|
async def route(self, prompt: str, context: dict[str, Any]) -> SimpleNamespace:
|
|
return SimpleNamespace(
|
|
selected_provider=AIProviderEnum.GEMINI,
|
|
fallback_chain=[
|
|
(AIProviderEnum.CLAUDE, "claude"),
|
|
(AIProviderEnum.OLLAMA, "qwen2.5:7b-instruct"),
|
|
],
|
|
intent=IntentType.DIAGNOSE,
|
|
routing_reason="high complexity would normally prefer cloud",
|
|
)
|
|
|
|
|
|
class _FakeExecutor:
|
|
def __init__(self) -> None:
|
|
self.provider_order: list[str] | None = None
|
|
|
|
async def execute(
|
|
self,
|
|
*,
|
|
prompt: str,
|
|
provider_order: list[str],
|
|
context: dict[str, Any],
|
|
cache_ttl: int,
|
|
require_local: bool,
|
|
) -> SimpleNamespace:
|
|
self.provider_order = provider_order
|
|
return SimpleNamespace(
|
|
raw_response='{"root_cause":"ok","suggested_action":"NO_ACTION"}',
|
|
provider=provider_order[0],
|
|
success=True,
|
|
tokens=42,
|
|
cost_usd=0.0,
|
|
latency_ms=10.0,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_alert_context_uses_ollama_lane_then_gemini_backup(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
fake_executor = _FakeExecutor()
|
|
fake_failover = _FakeFailoverManager()
|
|
|
|
monkeypatch.setattr(openclaw_module.settings, "USE_AI_ROUTER", True)
|
|
monkeypatch.setattr(openclaw_module.settings, "MOCK_MODE", False)
|
|
monkeypatch.setattr(openclaw_module.settings, "ALERT_AI_ALLOW_CLOUD_FALLBACK", True)
|
|
monkeypatch.setattr(openclaw_module.settings, "ALERT_AI_ENFORCE_OLLAMA_FIRST", True)
|
|
monkeypatch.setattr(ai_control_module, "get_ai_router_enabled", AsyncMock(return_value=None))
|
|
monkeypatch.setattr(ai_control_module, "get_primary_provider", AsyncMock(return_value=None))
|
|
monkeypatch.setattr(ai_control_module, "is_provider_disabled", AsyncMock(return_value=False))
|
|
monkeypatch.setattr(ai_router_module, "get_ai_router", lambda: _FakeRouter())
|
|
monkeypatch.setattr(ai_router_module, "get_ai_executor", lambda: fake_executor)
|
|
monkeypatch.setattr(openclaw_module, "get_ollama_failover_manager", lambda: fake_failover)
|
|
|
|
service = object.__new__(OpenClawService)
|
|
result = await service._call_with_fallback(
|
|
"diagnose alert",
|
|
alert_context={
|
|
"incident_id": "INC-1",
|
|
"alertname": "HostHighCpuLoad",
|
|
"target_resource": "node-exporter-110",
|
|
},
|
|
)
|
|
|
|
assert result == (
|
|
'{"root_cause":"ok","suggested_action":"NO_ACTION"}',
|
|
"ollama_gcp_a",
|
|
True,
|
|
42,
|
|
0.0,
|
|
)
|
|
assert fake_executor.provider_order == ["ollama_gcp_a", "ollama_gcp_b", "ollama_local", "gemini"]
|
|
assert fake_failover.task_types == ["diagnose"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_alert_context_can_disable_cloud_backup_for_cost_stop(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
fake_executor = _FakeExecutor()
|
|
fake_failover = _FakeFailoverManager()
|
|
|
|
monkeypatch.setattr(openclaw_module.settings, "USE_AI_ROUTER", True)
|
|
monkeypatch.setattr(openclaw_module.settings, "MOCK_MODE", False)
|
|
monkeypatch.setattr(openclaw_module.settings, "ALERT_AI_ALLOW_CLOUD_FALLBACK", False)
|
|
monkeypatch.setattr(openclaw_module.settings, "ALERT_AI_ENFORCE_OLLAMA_FIRST", True)
|
|
monkeypatch.setattr(ai_control_module, "get_ai_router_enabled", AsyncMock(return_value=None))
|
|
monkeypatch.setattr(ai_control_module, "get_primary_provider", AsyncMock(return_value=None))
|
|
monkeypatch.setattr(ai_control_module, "is_provider_disabled", AsyncMock(return_value=False))
|
|
monkeypatch.setattr(ai_router_module, "get_ai_router", lambda: _FakeRouter())
|
|
monkeypatch.setattr(ai_router_module, "get_ai_executor", lambda: fake_executor)
|
|
monkeypatch.setattr(openclaw_module, "get_ollama_failover_manager", lambda: fake_failover)
|
|
|
|
service = object.__new__(OpenClawService)
|
|
await service._call_with_fallback(
|
|
"diagnose alert",
|
|
alert_context={"incident_id": "INC-1", "alertname": "HostHighCpuLoad"},
|
|
)
|
|
|
|
assert fake_executor.provider_order == ["ollama_gcp_a", "ollama_gcp_b", "ollama_local"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_non_alert_context_keeps_router_cloud_order(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
fake_executor = _FakeExecutor()
|
|
|
|
monkeypatch.setattr(openclaw_module.settings, "USE_AI_ROUTER", True)
|
|
monkeypatch.setattr(openclaw_module.settings, "MOCK_MODE", False)
|
|
monkeypatch.setattr(openclaw_module.settings, "ALERT_AI_ALLOW_CLOUD_FALLBACK", False)
|
|
monkeypatch.setattr(openclaw_module.settings, "ALERT_AI_ENFORCE_OLLAMA_FIRST", True)
|
|
monkeypatch.setattr(ai_control_module, "get_ai_router_enabled", AsyncMock(return_value=None))
|
|
monkeypatch.setattr(ai_control_module, "get_primary_provider", AsyncMock(return_value=None))
|
|
monkeypatch.setattr(ai_control_module, "is_provider_disabled", AsyncMock(return_value=False))
|
|
monkeypatch.setattr(ai_router_module, "get_ai_router", lambda: _FakeRouter())
|
|
monkeypatch.setattr(ai_router_module, "get_ai_executor", lambda: fake_executor)
|
|
|
|
service = object.__new__(OpenClawService)
|
|
await service._call_with_fallback("general question", alert_context={"intent_hint": "query"})
|
|
|
|
assert fake_executor.provider_order == ["gemini", "claude", "ollama"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_alert_context_uses_gcp_a_gcp_b_then_111_order(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
fake_failover = _FakeFailoverManager()
|
|
|
|
monkeypatch.setattr(openclaw_module.settings, "ALERT_AI_ALLOW_CLOUD_FALLBACK", False)
|
|
monkeypatch.setattr(openclaw_module.settings, "ALERT_AI_ENFORCE_OLLAMA_FIRST", True)
|
|
monkeypatch.setattr(openclaw_module, "get_ollama_failover_manager", lambda: fake_failover)
|
|
|
|
service = object.__new__(OpenClawService)
|
|
provider_order = await service._resolve_alert_provider_order(
|
|
task_type="diagnose",
|
|
alert_context={"incident_id": "INC-1", "alertname": "HostHighCpuLoad"},
|
|
)
|
|
|
|
assert provider_order == ["ollama_gcp_a", "ollama_gcp_b", "ollama_local"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_alert_context_sorts_ollama_lane_and_drops_cloud_providers(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setattr(openclaw_module.settings, "ALERT_AI_ALLOW_CLOUD_FALLBACK", False)
|
|
monkeypatch.setattr(openclaw_module.settings, "ALERT_AI_ENFORCE_OLLAMA_FIRST", True)
|
|
monkeypatch.setattr(openclaw_module, "get_ollama_failover_manager", lambda: _UnorderedFailoverManager())
|
|
|
|
service = object.__new__(OpenClawService)
|
|
provider_order = await service._resolve_alert_provider_order(
|
|
task_type="diagnose",
|
|
alert_context={"incident_id": "INC-1", "alertname": "HostHighCpuLoad"},
|
|
)
|
|
|
|
assert provider_order == ["ollama_gcp_a", "ollama_gcp_b", "ollama_local"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_alert_context_sorts_ollama_lane_before_gemini_backup(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setattr(openclaw_module.settings, "ALERT_AI_ALLOW_CLOUD_FALLBACK", True)
|
|
monkeypatch.setattr(openclaw_module.settings, "ALERT_AI_ENFORCE_OLLAMA_FIRST", True)
|
|
monkeypatch.setattr(openclaw_module, "get_ollama_failover_manager", lambda: _UnorderedFailoverManager())
|
|
|
|
service = object.__new__(OpenClawService)
|
|
provider_order = await service._resolve_alert_provider_order(
|
|
task_type="diagnose",
|
|
alert_context={"incident_id": "INC-1", "alertname": "HostHighCpuLoad"},
|
|
cloud_provider_order=["claude", "gemini", "ollama"],
|
|
)
|
|
|
|
assert provider_order == ["ollama_gcp_a", "ollama_gcp_b", "ollama_local", "gemini"]
|