78 lines
3.0 KiB
Python
78 lines
3.0 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import datetime
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from src.services import heartbeat_report_service as heartbeat
|
|
|
|
|
|
class _FakeResponse:
|
|
def __init__(self, status_code: int, payload: dict[str, Any] | None = None) -> None:
|
|
self.status_code = status_code
|
|
self._payload = payload or {}
|
|
|
|
def json(self) -> dict[str, Any]:
|
|
return self._payload
|
|
|
|
|
|
class _FakeAsyncClient:
|
|
def __init__(self, *_args: Any, **_kwargs: Any) -> None:
|
|
pass
|
|
|
|
async def __aenter__(self) -> "_FakeAsyncClient":
|
|
return self
|
|
|
|
async def __aexit__(self, *_args: Any) -> None:
|
|
return None
|
|
|
|
async def get(self, url: str) -> _FakeResponse:
|
|
if url.startswith("http://gcp-a"):
|
|
return _FakeResponse(
|
|
200,
|
|
{"models": [{"name": "qwen3:14b"}, {"name": "bge-m3:latest"}]},
|
|
)
|
|
if url.startswith("http://gcp-b"):
|
|
return _FakeResponse(200, {"models": [{"name": "gemma3:4b"}]})
|
|
raise TimeoutError("connect failed")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_probe_ollama_reports_each_endpoint(monkeypatch) -> None:
|
|
monkeypatch.setattr(heartbeat.httpx, "AsyncClient", _FakeAsyncClient)
|
|
monkeypatch.setattr(heartbeat.settings, "OLLAMA_URL", "http://gcp-a:11434")
|
|
monkeypatch.setattr(heartbeat.settings, "OLLAMA_SECONDARY_URL", "http://gcp-b:11434")
|
|
monkeypatch.setattr(heartbeat.settings, "OLLAMA_FALLBACK_URL", "http://local-111:11434")
|
|
monkeypatch.setattr(heartbeat.settings, "OLLAMA_REQUIRED_MODELS", ["qwen3:14b", "bge-m3:latest"])
|
|
|
|
result = await heartbeat.HeartbeatReportService()._probe_ollama()
|
|
|
|
assert result["probe"].ok is True
|
|
assert result["models"] == {"qwen3:14b": True, "bge-m3:latest": True}
|
|
assert result["endpoints"]["GCP-A"].ok is True
|
|
assert result["endpoints"]["GCP-B"].ok is True
|
|
assert result["endpoints"]["111"].ok is False
|
|
|
|
|
|
def test_report_to_telegram_html_renders_ollama_endpoint_statuses() -> None:
|
|
report = heartbeat.HeartbeatReport(timestamp=datetime(2026, 5, 6, 18, 0))
|
|
report.ai_services["ollama"] = heartbeat.ProbeResult(True, "✅ 正常", 1200)
|
|
report.ai_services["nemotron"] = heartbeat.ProbeResult(True, "✅ 正常", 900)
|
|
report.ai_services["gemini"] = heartbeat.ProbeResult(True, "✅ 正常", 800)
|
|
report.ai_services["claude"] = heartbeat.ProbeResult(True, "✅ 正常", 700)
|
|
report.ollama_models = {"qwen3:14b": True}
|
|
report.ollama_endpoints = {
|
|
"GCP-A": heartbeat.ProbeResult(True, "✅ 正常", 1000),
|
|
"GCP-B": heartbeat.ProbeResult(True, "✅ 正常", 1100),
|
|
"111": heartbeat.ProbeResult(False, "❌ connect failed"),
|
|
}
|
|
report.warnings = heartbeat.HeartbeatReportService()._build_warnings(report)
|
|
|
|
text = heartbeat.report_to_telegram_html(report)
|
|
|
|
assert "GCP-A: ✅ 正常" in text
|
|
assert "GCP-B: ✅ 正常" in text
|
|
assert "111: ❌ connect failed" in text
|
|
assert "Ollama 111 異常" in "\n".join(report.warnings)
|