"""Agent watch dual-mode — heartbeat response + probe path selection.""" from __future__ import annotations from datetime import datetime, timezone 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.domain.models import WatchPin, WatchPinSession from server.utils.watch_metrics import ( WatchProbeSample, agent_watch_heartbeat_fresh, agent_watch_sample_usable, heartbeat_watch_mode_active, parse_redis_watch_payload, PROBE_OK, ) def _heartbeat_mapping(*, watch_mode: bool = True, age_sec: float = 3.0) -> dict[str, str]: ts = datetime.now(timezone.utc).timestamp() - age_sec last = datetime.fromtimestamp(ts, tz=timezone.utc).isoformat() info = { "cpu_usage": 12.0, "mem_usage": 34.0, "disk_usage": 56.0, "load_avg": [0.2, 0.3, 0.4], } if watch_mode: info["watch_mode"] = True import json return { "is_online": "True", "last_heartbeat": last, "system_info": json.dumps(info), "agent_version": "2.1.0", } def test_heartbeat_watch_mode_active_flag(): hb = _heartbeat_mapping(watch_mode=True) assert heartbeat_watch_mode_active(hb) is True hb2 = _heartbeat_mapping(watch_mode=False) assert heartbeat_watch_mode_active(hb2) is False def test_agent_watch_heartbeat_fresh_within_window(): assert agent_watch_heartbeat_fresh(_heartbeat_mapping(age_sec=3)) is True assert agent_watch_heartbeat_fresh(_heartbeat_mapping(age_sec=30)) is False def test_agent_watch_sample_usable(): hb = _heartbeat_mapping() sample = parse_redis_watch_payload(hb) assert sample is not None assert agent_watch_sample_usable(hb, sample) is True stale = _heartbeat_mapping(age_sec=20) assert agent_watch_sample_usable(stale, sample) is False @pytest.mark.asyncio async def test_heartbeat_response_watch_active_when_pin_monitoring( db_session, test_server, test_admin, mock_request, monkeypatch, ): from datetime import timedelta from server.infrastructure.database.watch_repo import _utc_naive_now api_key = "test-agent-watch-key" test_server.agent_api_key = api_key await db_session.commit() admin_id = test_admin["admin"].id now = _utc_naive_now() sess = WatchPinSession( admin_id=admin_id, server_id=test_server.id, slot_index=0, started_at=now, ) db_session.add(sess) await db_session.flush() pin = WatchPin( admin_id=admin_id, server_id=test_server.id, slot_index=0, session_id=sess.id, pinned_at=now, expires_at=now + timedelta(minutes=30), monitoring_enabled=True, ttl_minutes=30, ttl_hours=1, ) db_session.add(pin) await db_session.commit() mock_redis = AsyncMock() mock_redis.hset = AsyncMock() mock_redis.expire = AsyncMock() mock_redis.smembers = AsyncMock(return_value=set()) mock_redis.sadd = AsyncMock() mock_redis.srem = AsyncMock() mock_redis.incr = AsyncMock(return_value=1) mock_redis.delete = AsyncMock() monkeypatch.setattr("server.api.agent.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) payload = AgentHeartbeat( server_id=test_server.id, is_online=True, system_info={"cpu_usage": 10.0, "mem_usage": 20.0, "disk_usage": 30.0, "watch_mode": True}, agent_version="2.1.0", ) result = await receive_heartbeat( payload, api_key=api_key, service=service, db=db_session, ) assert result["status"] == "ok" assert result["watch_active"] is True assert result["watch_interval_sec"] == 5 @pytest.mark.asyncio async def test_heartbeat_response_watch_inactive_without_pin( db_session, test_server, mock_request, monkeypatch, ): api_key = "test-agent-watch-key-2" test_server.agent_api_key = api_key await db_session.commit() mock_redis = AsyncMock() mock_redis.hset = AsyncMock() mock_redis.expire = AsyncMock() mock_redis.smembers = AsyncMock(return_value=set()) mock_redis.incr = AsyncMock(return_value=1) mock_redis.delete = AsyncMock() monkeypatch.setattr("server.api.agent.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) payload = AgentHeartbeat( server_id=test_server.id, is_online=True, system_info={"cpu_usage": 1.0}, agent_version="2.1.0", ) result = await receive_heartbeat( payload, api_key=api_key, service=service, db=db_session, ) assert result["watch_active"] is False