Codex Release Bot
92 lines
2.7 KiB
Python
92 lines
2.7 KiB
Python
"""CH-AGT-001 — heartbeat → Redis → dashboard online count."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
|
|
from server.api.agent import receive_heartbeat
|
|
from server.api.dependencies import get_server_service
|
|
from server.api.schemas import AgentHeartbeat
|
|
from server.api.servers import server_stats
|
|
|
|
pytestmark = pytest.mark.chain
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_heartbeat_then_stats_online_count(
|
|
db_session,
|
|
test_server,
|
|
test_admin,
|
|
mock_request,
|
|
monkeypatch,
|
|
):
|
|
api_key = "chain-agent-key"
|
|
test_server.agent_api_key = api_key
|
|
await db_session.commit()
|
|
|
|
stored: dict[str, dict] = {}
|
|
pipeline_keys: list[str] = []
|
|
|
|
async def fake_hset(key, mapping=None):
|
|
if mapping:
|
|
stored[key] = dict(mapping)
|
|
|
|
async def fake_scan(cursor, match=None, count=200):
|
|
keys = [k for k in stored if k.startswith("heartbeat:")]
|
|
return 0, keys
|
|
|
|
async def fake_hgetall(key):
|
|
return stored.get(key, {})
|
|
|
|
class FakePipeline:
|
|
def hgetall(self, key: str) -> None:
|
|
pipeline_keys.append(key)
|
|
|
|
async def execute(self):
|
|
keys = list(pipeline_keys)
|
|
pipeline_keys.clear()
|
|
return [stored.get(k, {}) for k in keys]
|
|
|
|
def fake_pipeline():
|
|
return FakePipeline()
|
|
|
|
mock_redis = AsyncMock()
|
|
mock_redis.hset = fake_hset
|
|
mock_redis.expire = AsyncMock()
|
|
mock_redis.get = AsyncMock(return_value=None)
|
|
mock_redis.smembers = AsyncMock(return_value=set())
|
|
mock_redis.sadd = AsyncMock()
|
|
mock_redis.srem = AsyncMock()
|
|
mock_redis.delete = AsyncMock()
|
|
mock_redis.incr = AsyncMock(return_value=1)
|
|
mock_redis.scan = fake_scan
|
|
mock_redis.hgetall = fake_hgetall
|
|
mock_redis.pipeline = fake_pipeline
|
|
|
|
monkeypatch.setattr("server.api.agent.get_redis", lambda: mock_redis)
|
|
monkeypatch.setattr("server.api.servers.get_redis", lambda: mock_redis)
|
|
monkeypatch.setattr("server.utils.alert_streak.broadcast_alert", AsyncMock())
|
|
monkeypatch.setattr("server.utils.alert_streak.broadcast_recovery", AsyncMock())
|
|
monkeypatch.setattr("server.api.agent.send_telegram_system_alert", AsyncMock())
|
|
|
|
mock_request.state.db = db_session
|
|
service = await get_server_service(mock_request)
|
|
|
|
hb = await receive_heartbeat(
|
|
AgentHeartbeat(
|
|
server_id=test_server.id,
|
|
is_online=True,
|
|
system_info={"cpu_usage": 10.0, "mem_usage": 20.0, "disk_usage": 30.0},
|
|
),
|
|
api_key=api_key,
|
|
service=service,
|
|
db=db_session,
|
|
)
|
|
assert hb["status"] == "ok"
|
|
|
|
stats = await server_stats(test_admin["admin"], db_session)
|
|
assert stats["online"] >= 1
|
|
assert stats["total"] >= 1
|