headquarter/apps/api/tests/services/test_terminal_manager_multi.py

"""Unit tests for TerminalManager multi-session support."""

import asyncio
import uuid
from unittest.mock import AsyncMock, patch

import pytest

from src.services.terminal_manager import MaxSessionsExceededError, TerminalManager
from src.services.terminal_session import TerminalSession


@pytest.fixture
def manager() -> TerminalManager:
    """Provide a fresh TerminalManager instance for each test."""
    tm = TerminalManager()
    # Cancel the background idle check to avoid side effects
    if tm._idle_check_task and not tm._idle_check_task.done():
        tm._idle_check_task.cancel()
    return tm


@pytest.fixture
def mock_terminal_session(monkeypatch) -> None:
    """Monkeypatch TerminalSession.start and is_alive for unit tests."""

    async def fake_start(self, startup_command=None):
        self.last_activity = __import__("time").time()

    monkeypatch.setattr(TerminalSession, "start", fake_start)
    monkeypatch.setattr(TerminalSession, "is_alive", lambda self: True)


@pytest.fixture
def instance_id() -> uuid.UUID:
    return uuid.uuid4()


class FakeWebSocket:
    """Minimal fake WebSocket for testing attach/detach behavior."""

    def __init__(self, name: str = "ws") -> None:
        self.name = name
        self.closed = False
        self.close_code: int | None = None
        self.close_reason: str | None = None
        self._sent: list[bytes] = []

    async def close(self, code: int = 1000, reason: str = "") -> None:
        self.closed = True
        self.close_code = code
        self.close_reason = reason

    async def send_bytes(self, data: bytes) -> None:
        self._sent.append(data)


@pytest.mark.asyncio
async def test_create_session_increases_count(
    manager: TerminalManager,
    mock_terminal_session,
    instance_id: uuid.UUID,
) -> None:
    """Creating sessions increments the per-instance count."""
    assert len(manager.get_sessions_for_instance(str(instance_id))) == 0

    session1 = await manager.create_session(instance_id, "container-1")
    assert len(manager.get_sessions_for_instance(str(instance_id))) == 1
    assert session1.session_id in [
        s.session_id for s in manager.get_sessions_for_instance(str(instance_id))
    ]

    session2 = await manager.create_session(instance_id, "container-1")
    assert len(manager.get_sessions_for_instance(str(instance_id))) == 2

    # Verify sessions are distinct
    assert session1.session_id != session2.session_id


@pytest.mark.asyncio
async def test_create_session_enforces_max_5(
    manager: TerminalManager,
    mock_terminal_session,
    instance_id: uuid.UUID,
) -> None:
    """The 6th session creation raises MaxSessionsExceededError."""
    for i in range(5):
        await manager.create_session(instance_id, f"container-{i}")

    assert len(manager.get_sessions_for_instance(str(instance_id))) == 5

    with pytest.raises(MaxSessionsExceededError):
        await manager.create_session(instance_id, "container-overflow")


@pytest.mark.asyncio
async def test_get_sessions_for_instance_filters_by_instance(
    manager: TerminalManager,
    mock_terminal_session,
) -> None:
    """get_sessions_for_instance returns only sessions for the requested instance."""
    instance_a = uuid.uuid4()
    instance_b = uuid.uuid4()

    await manager.create_session(instance_a, "container-a")
    await manager.create_session(instance_a, "container-a2")
    await manager.create_session(instance_b, "container-b")

    assert len(manager.get_sessions_for_instance(str(instance_a))) == 2
    assert len(manager.get_sessions_for_instance(str(instance_b))) == 1


@pytest.mark.asyncio
async def test_close_session_removes_from_dict(
    manager: TerminalManager,
    mock_terminal_session,
    instance_id: uuid.UUID,
) -> None:
    """close_session removes the key from _sessions and marks DB closed."""
    session = await manager.create_session(instance_id, "container-1")
    session_id = session.session_id

    assert manager.get_session(str(instance_id), session_id) is not None

    with patch.object(manager, "_mark_closed_in_db", new=AsyncMock()) as mock_mark:
        await manager.close_session(str(instance_id), session_id)
        # Give the fire-and-forget task a chance to be scheduled
        await asyncio.sleep(0)

    assert manager.get_session(str(instance_id), session_id) is None
    mock_mark.assert_called_once_with(session_id)


@pytest.mark.asyncio
async def test_attach_websocket_only_closes_same_session(
    manager: TerminalManager,
    mock_terminal_session,
    instance_id: uuid.UUID,
) -> None:
    """Attaching to session A must not close WebSockets on session B."""
    session_a = await manager.create_session(instance_id, "container-1")
    session_b = await manager.create_session(instance_id, "container-1")

    ws_a1 = FakeWebSocket("ws-a1")
    ws_b1 = FakeWebSocket("ws-b1")

    # Manually attach websockets (simulate prior connections)
    session_a.attach_websocket(ws_a1)
    session_b.attach_websocket(ws_b1)

    # Now attach a new websocket to session_a
    ws_a2 = FakeWebSocket("ws-a2")
    await manager.attach_websocket(session_a, ws_a2)

    # ws_a1 should have been closed because it's on the same session
    assert ws_a1.closed is True

    # ws_b1 should NOT have been closed because it's on a different session
    assert ws_b1.closed is False

    # ws_a2 should be attached and receive buffer
    assert ws_a2 in session_a._websockets


@pytest.mark.asyncio
async def test_default_session_keyed_separately(
    manager: TerminalManager,
    mock_terminal_session,
    instance_id: uuid.UUID,
) -> None:
    """Default session uses 'default' session_id and does not collide with named sessions."""
    default_session = await manager.get_or_create_session(instance_id, "container-1")
    explicit_session = await manager.create_session(instance_id, "container-1")

    # Both should exist
    assert manager.get_session(str(instance_id), "default") is default_session
    assert (
        manager.get_session(str(instance_id), explicit_session.session_id)
        is explicit_session
    )

    # They should be different objects
    assert default_session.session_id != explicit_session.session_id


@pytest.mark.asyncio
async def test_idle_cleanup_updates_db_status(
    manager: TerminalManager,
    mock_terminal_session,
    instance_id: uuid.UUID,
) -> None:
    """Idle cleanup removes sessions from dict and calls DB update."""
    session = await manager.create_session(instance_id, "container-1")
    session_id = session.session_id

    # Make session appear idle (no websockets, old last_activity)
    session.last_activity = 0

    with patch.object(manager, "_mark_closed_in_db", new=AsyncMock()) as mock_mark:
        await manager._cleanup_idle_sessions()

    assert manager.get_session(str(instance_id), session_id) is None
    mock_mark.assert_called_once_with(session_id)