breakpilot-lehrer/agent-core/tests/test_heartbeat.py

"""
Tests for Heartbeat Monitoring

Tests cover:
- Heartbeat registration and updates
- Timeout detection
- Pause/resume functionality
- Status reporting
"""

import pytest
import asyncio
from datetime import datetime, timezone, timedelta
from unittest.mock import AsyncMock

import sys
sys.path.insert(0, str(__file__).rsplit('/tests/', 1)[0])

from sessions.heartbeat import HeartbeatMonitor, HeartbeatClient, HeartbeatEntry


class TestHeartbeatMonitor:
    """Tests for HeartbeatMonitor"""

    @pytest.fixture
    def monitor(self):
        """Create a heartbeat monitor"""
        return HeartbeatMonitor(
            timeout_seconds=5,
            check_interval_seconds=1,
            max_missed_beats=2
        )

    def test_register_session(self, monitor):
        """Should register session for monitoring"""
        monitor.register("session-1", "tutor-agent")

        assert "session-1" in monitor.sessions
        assert monitor.sessions["session-1"].agent_type == "tutor-agent"

    def test_beat_updates_timestamp(self, monitor):
        """Beat should update last_beat timestamp"""
        monitor.register("session-1", "agent")
        original = monitor.sessions["session-1"].last_beat

        import time
        time.sleep(0.01)

        result = monitor.beat("session-1")

        assert result is True
        assert monitor.sessions["session-1"].last_beat > original
        assert monitor.sessions["session-1"].missed_beats == 0

    def test_beat_nonexistent_session(self, monitor):
        """Beat should return False for unregistered session"""
        result = monitor.beat("nonexistent")
        assert result is False

    def test_unregister_session(self, monitor):
        """Should unregister session from monitoring"""
        monitor.register("session-1", "agent")

        result = monitor.unregister("session-1")

        assert result is True
        assert "session-1" not in monitor.sessions

    def test_pause_session(self, monitor):
        """Should pause monitoring for session"""
        monitor.register("session-1", "agent")

        result = monitor.pause("session-1")

        assert result is True
        assert "session-1" in monitor._paused_sessions

    def test_resume_session(self, monitor):
        """Should resume monitoring for paused session"""
        monitor.register("session-1", "agent")
        monitor.pause("session-1")

        result = monitor.resume("session-1")

        assert result is True
        assert "session-1" not in monitor._paused_sessions

    def test_get_status(self, monitor):
        """Should return session status"""
        monitor.register("session-1", "tutor-agent")

        status = monitor.get_status("session-1")

        assert status is not None
        assert status["session_id"] == "session-1"
        assert status["agent_type"] == "tutor-agent"
        assert status["is_healthy"] is True
        assert status["is_paused"] is False

    def test_get_status_nonexistent(self, monitor):
        """Should return None for nonexistent session"""
        status = monitor.get_status("nonexistent")
        assert status is None

    def test_get_all_status(self, monitor):
        """Should return status for all sessions"""
        monitor.register("session-1", "agent-1")
        monitor.register("session-2", "agent-2")

        all_status = monitor.get_all_status()

        assert len(all_status) == 2
        assert "session-1" in all_status
        assert "session-2" in all_status

    def test_registered_count(self, monitor):
        """Should return correct registered count"""
        assert monitor.registered_count == 0

        monitor.register("s1", "a")
        monitor.register("s2", "a")

        assert monitor.registered_count == 2

    def test_healthy_count(self, monitor):
        """Should return correct healthy count"""
        monitor.register("s1", "a")
        monitor.register("s2", "a")

        # Both should be healthy initially
        assert monitor.healthy_count == 2

        # Simulate missed beat
        monitor.sessions["s1"].missed_beats = 1
        assert monitor.healthy_count == 1


class TestHeartbeatClient:
    """Tests for HeartbeatClient"""

    @pytest.fixture
    def monitor(self):
        """Create a monitor for the client"""
        return HeartbeatMonitor(timeout_seconds=5)

    def test_client_creation(self, monitor):
        """Client should be created with correct settings"""
        client = HeartbeatClient(
            session_id="session-1",
            monitor=monitor,
            interval_seconds=2
        )

        assert client.session_id == "session-1"
        assert client.interval == 2
        assert client._running is False

    @pytest.mark.asyncio
    async def test_client_start_stop(self, monitor):
        """Client should start and stop correctly"""
        monitor.register("session-1", "agent")

        client = HeartbeatClient(
            session_id="session-1",
            monitor=monitor,
            interval_seconds=1
        )

        await client.start()
        assert client._running is True

        await asyncio.sleep(0.1)

        await client.stop()
        assert client._running is False

    @pytest.mark.asyncio
    async def test_client_context_manager(self, monitor):
        """Client should work as context manager"""
        monitor.register("session-1", "agent")

        async with HeartbeatClient("session-1", monitor, 1) as client:
            assert client._running is True

        assert client._running is False


class TestHeartbeatEntry:
    """Tests for HeartbeatEntry dataclass"""

    def test_entry_creation(self):
        """Entry should be created with correct values"""
        entry = HeartbeatEntry(
            session_id="session-1",
            agent_type="tutor-agent",
            last_beat=datetime.now(timezone.utc)
        )

        assert entry.session_id == "session-1"
        assert entry.agent_type == "tutor-agent"
        assert entry.missed_beats == 0