smart-support/backend/tests/unit/test_edge_cases.py

"""Edge case tests for ws_handler input validation and rate limiting."""

from __future__ import annotations

import json
from unittest.mock import AsyncMock, MagicMock

import pytest

from app.callbacks import TokenUsageCallbackHandler
from app.session_manager import SessionManager
from app.ws_handler import dispatch_message

pytestmark = pytest.mark.unit


def _make_ws() -> AsyncMock:
    ws = AsyncMock()
    ws.send_json = AsyncMock()
    return ws


def _make_graph() -> AsyncMock:
    graph = AsyncMock()

    class AsyncIterHelper:
        def __aiter__(self):
            return self

        async def __anext__(self):
            raise StopAsyncIteration

    graph.astream = MagicMock(return_value=AsyncIterHelper())
    state = MagicMock()
    state.tasks = ()
    graph.aget_state = AsyncMock(return_value=state)
    graph.intent_classifier = None
    graph.agent_registry = None
    return graph


@pytest.mark.unit
class TestEmptyMessageHandling:
    @pytest.mark.asyncio
    async def test_empty_message_content_returns_error(self) -> None:
        ws = _make_ws()
        graph = _make_graph()
        sm = SessionManager()
        cb = TokenUsageCallbackHandler()

        sm.touch("t1")
        msg = json.dumps({"type": "message", "thread_id": "t1", "content": ""})
        await dispatch_message(ws, graph, sm, cb, msg)

        call_data = ws.send_json.call_args[0][0]
        assert call_data["type"] == "error"
        msg_lower = call_data["message"].lower()
        assert "content" in msg_lower or "missing" in msg_lower

    @pytest.mark.asyncio
    async def test_whitespace_only_message_treated_as_empty(self) -> None:
        ws = _make_ws()
        graph = _make_graph()
        sm = SessionManager()
        cb = TokenUsageCallbackHandler()

        sm.touch("t1")
        msg = json.dumps({"type": "message", "thread_id": "t1", "content": "   "})
        await dispatch_message(ws, graph, sm, cb, msg)

        call_data = ws.send_json.call_args[0][0]
        assert call_data["type"] == "error"


@pytest.mark.unit
class TestOversizedMessageHandling:
    @pytest.mark.asyncio
    async def test_content_over_10000_chars_returns_error(self) -> None:
        ws = _make_ws()
        graph = _make_graph()
        sm = SessionManager()
        cb = TokenUsageCallbackHandler()

        sm.touch("t1")
        content = "x" * 10001
        msg = json.dumps({"type": "message", "thread_id": "t1", "content": content})
        await dispatch_message(ws, graph, sm, cb, msg)

        call_data = ws.send_json.call_args[0][0]
        assert call_data["type"] == "error"
        assert "too long" in call_data["message"].lower()

    @pytest.mark.asyncio
    async def test_content_exactly_10000_chars_is_accepted(self) -> None:
        ws = _make_ws()
        graph = _make_graph()
        sm = SessionManager()
        cb = TokenUsageCallbackHandler()

        sm.touch("t1")
        content = "x" * 10000
        msg = json.dumps({"type": "message", "thread_id": "t1", "content": content})
        await dispatch_message(ws, graph, sm, cb, msg)

        last_call = ws.send_json.call_args[0][0]
        # Should be processed, not an error about length
        msg_text = last_call.get("message", "").lower()
        assert last_call["type"] != "error" or "too long" not in msg_text

    @pytest.mark.asyncio
    async def test_raw_message_over_32kb_returns_error(self) -> None:
        ws = _make_ws()
        graph = _make_graph()
        sm = SessionManager()
        cb = TokenUsageCallbackHandler()

        large_msg = "x" * 40_000
        await dispatch_message(ws, graph, sm, cb, large_msg)

        call_data = ws.send_json.call_args[0][0]
        assert call_data["type"] == "error"
        assert "too large" in call_data["message"].lower()


@pytest.mark.unit
class TestInvalidJsonHandling:
    @pytest.mark.asyncio
    async def test_invalid_json_returns_error(self) -> None:
        ws = _make_ws()
        graph = _make_graph()
        sm = SessionManager()
        cb = TokenUsageCallbackHandler()

        await dispatch_message(ws, graph, sm, cb, "not valid json {{")

        call_data = ws.send_json.call_args[0][0]
        assert call_data["type"] == "error"
        assert "invalid json" in call_data["message"].lower()

    @pytest.mark.asyncio
    async def test_empty_string_returns_json_error(self) -> None:
        ws = _make_ws()
        graph = _make_graph()
        sm = SessionManager()
        cb = TokenUsageCallbackHandler()

        await dispatch_message(ws, graph, sm, cb, "")

        call_data = ws.send_json.call_args[0][0]
        assert call_data["type"] == "error"

    @pytest.mark.asyncio
    async def test_json_array_not_object_returns_error(self) -> None:
        ws = _make_ws()
        graph = _make_graph()
        sm = SessionManager()
        cb = TokenUsageCallbackHandler()

        await dispatch_message(ws, graph, sm, cb, '["not", "an", "object"]')

        call_data = ws.send_json.call_args[0][0]
        assert call_data["type"] == "error"


@pytest.mark.unit
class TestRateLimiting:
    @pytest.mark.asyncio
    async def test_rapid_fire_messages_rate_limited(self) -> None:
        ws = _make_ws()
        _make_graph()  # ensure graph factory works, not needed directly
        sm = SessionManager()
        cb = TokenUsageCallbackHandler()

        sm.touch("t1")

        # Simulate 11 rapid messages (exceeds 10 per 10 seconds limit)
        rate_limit_triggered = False
        for i in range(11):
            graph2 = _make_graph()  # fresh graph each time
            await dispatch_message(ws, graph2, sm, cb, json.dumps({
                "type": "message",
                "thread_id": "t1",
                "content": f"message {i}",
            }))
            last_call = ws.send_json.call_args[0][0]
            if last_call["type"] == "error" and "rate" in last_call.get("message", "").lower():
                rate_limit_triggered = True
                break

        assert rate_limit_triggered, "Rate limiting should trigger after 10 rapid messages"

    @pytest.mark.asyncio
    async def test_different_threads_have_separate_rate_limits(self) -> None:
        ws = _make_ws()
        sm = SessionManager()
        cb = TokenUsageCallbackHandler()

        sm.touch("t1")
        sm.touch("t2")

        # Send 5 messages on t1 and 5 on t2 -- neither should be rate limited
        for i in range(5):
            graph1 = _make_graph()
            graph2 = _make_graph()
            await dispatch_message(ws, graph1, sm, cb, json.dumps({
                "type": "message", "thread_id": "t1", "content": f"msg {i}",
            }))
            await dispatch_message(ws, graph2, sm, cb, json.dumps({
                "type": "message", "thread_id": "t2", "content": f"msg {i}",
            }))

        last_call = ws.send_json.call_args[0][0]
        assert "rate" not in last_call.get("message", "").lower()