Nexus/tests/test_agent_watch_mode.py

"""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