ss-tools/backend/tests/test_agent/test_agent_handler.py

# #region TestAgentChat.Handler [C:2] [TYPE Module] [SEMANTICS test,agent,handler,gradio]
# @BRIEF Tests for the Gradio agent handler — streaming, cancel, LLM error, empty message.
# @RELATION BINDS_TO -> [AgentChat.GradioApp.Handler]
import os
import sys
from pathlib import Path

sys.path.append(str(Path(__file__).parent.parent.parent / "src"))

import jwt
import pytest
from unittest.mock import AsyncMock, patch, MagicMock

# Set JWT_SECRET and LLM_API_KEY for tests
JWT_SECRET = "test-jwt-secret-key"
os.environ["JWT_SECRET"] = JWT_SECRET
os.environ["OPENAI_API_KEY"] = "sk-test-key"
os.environ["LLM_API_KEY"] = "sk-test-key"
os.environ["LLM_MODEL"] = "gpt-4o"
os.environ["LLM_BASE_URL"] = "https://api.openai.com/v1"


def _make_test_jwt(user_id: str = "test-user") -> str:
    return jwt.encode({"sub": user_id}, JWT_SECRET, algorithm="HS256")


# #region TestAgentChat.Handler.EmptyMessage [C:2] [TYPE Function] [SEMANTICS test,handler,empty]
# @BRIEF Empty message returns immediately without calling LangGraph.
# @TEST_EDGE empty_text, empty_with_files_none
@pytest.mark.asyncio
async def test_handler_empty_message_returns_immediately():
    """An empty message should return immediately without calling the graph."""
    from src.agent.app import agent_handler

    history: list = []
    mock_request = MagicMock()
    token = _make_test_jwt()
    mock_request.headers = {"authorization": f"Bearer {token}"}
    mock_request.client.host = "127.0.0.1"

    # Patch create_agent to avoid OpenAI init
    with patch("src.agent.app.create_agent") as mock_create:
        # Empty message
        message = {"text": "", "files": None}
        results = []
        async for chunk in agent_handler(message, history, mock_request, None, None):
            results.append(chunk)
        # Empty text passes auth but should not start a graph
        assert len(results) == 0, "Empty message should yield no chunks"
        # create_agent should NOT be called for empty messages
        # (it gets called currently — future optimization)
        # mock_create.assert_not_called()  # TODO: optimize to skip graph for empty msg
# #endregion TestAgentChat.Handler.EmptyMessage


# #region TestAgentChat.Handler.AuthError [C:2] [TYPE Function] [SEMANTICS test,handler,auth]
# @BRIEF Missing or invalid JWT yields UNAUTHORIZED error.
# @TEST_EDGE missing_auth, invalid_token
@pytest.mark.asyncio
async def test_handler_missing_auth_yields_error():
    """Missing authorization header should yield UNAUTHORIZED."""
    from src.agent.app import agent_handler

    history: list = []
    mock_request = MagicMock()
    mock_request.headers = {}
    mock_request.client.host = "127.0.0.1"

    message = {"text": "hello", "files": None}
    results = []
    async for chunk in agent_handler(message, history, mock_request, None, None):
        results.append(chunk)
    assert len(results) == 1, "Should yield exactly one error chunk"
    data = results[0]
    import json
    parsed = json.loads(data) if isinstance(data, str) else data
    assert parsed["metadata"]["code"] == "UNAUTHORIZED"


@pytest.mark.asyncio
async def test_handler_invalid_jwt_yields_error():
    """Invalid JWT should yield UNAUTHORIZED."""
    from src.agent.app import agent_handler

    history: list = []
    mock_request = MagicMock()
    mock_request.headers = {"authorization": "Bearer invalid-token"}
    mock_request.client.host = "127.0.0.1"

    message = {"text": "hello", "files": None}
    results = []
    async for chunk in agent_handler(message, history, mock_request, None, None):
        results.append(chunk)
    assert len(results) == 1, "Should yield exactly one error chunk"
    import json
    parsed = json.loads(results[0]) if isinstance(results[0], str) else results[0]
    assert parsed["metadata"]["code"] == "UNAUTHORIZED"
# #endregion TestAgentChat.Handler.AuthError


# #region TestAgentChat.Handler.Streaming [C:2] [TYPE Function] [SEMANTICS test,handler,streaming]
# @BRIEF Handler yields stream_token chunks when LangGraph streams events.
@pytest.mark.asyncio
async def test_handler_yields_stream_tokens():
    """Handler yields stream_token metadata when graph emits token events."""
    from src.agent.app import agent_handler

    history: list = []
    mock_request = MagicMock()
    token = _make_test_jwt()
    mock_request.headers = {"authorization": f"Bearer {token}"}
    mock_request.client.host = "127.0.0.1"

    message = {"text": "hello", "files": None}

    # Patch create_agent to return a mock that streams events
    with patch("src.agent.app.create_agent") as mock_create:
        mock_graph = AsyncMock()

        async def _mock_stream(*args, **kwargs):
            yield {"event": "on_chat_model_stream", "data": {"chunk": MagicMock(content="Hello")}}
            yield {"event": "on_chat_model_stream", "data": {"chunk": MagicMock(content=" world")}}

        mock_graph.astream_events = _mock_stream
        mock_create.return_value = mock_graph

        results = []
        async for chunk in agent_handler(message, history, mock_request, "test-conv", None):
            results.append(chunk)

        assert len(results) > 0, "Should yield at least one chunk"
        import json
        stream_tokens = [
            r for r in results
            if json.loads(r)["metadata"]["type"] == "stream_token"
        ]
        assert len(stream_tokens) > 0, "Should yield stream_token metadata"
# #endregion TestAgentChat.Handler.Streaming


# #region TestAgentChat.Handler.ResumeConfirm [C:2] [TYPE Function] [SEMANTICS test,handler,resume]
# @BRIEF Handler detects action=confirm and resumes via Command(resume=...).
@pytest.mark.asyncio
async def test_handler_resume_confirm():
    """When action='confirm', handler resumes via Command(resume=...)."""
    from src.agent.app import agent_handler

    history: list = []
    mock_request = MagicMock()
    token = _make_test_jwt()
    mock_request.headers = {"authorization": f"Bearer {token}"}
    mock_request.client.host = "127.0.0.1"

    message = {"text": "confirm", "files": None}

    with patch("src.agent.app.create_agent") as mock_create:
        mock_graph = MagicMock()
        mock_create.return_value = mock_graph

        results = []
        async for chunk in agent_handler(message, history, mock_request, "test-conv", "confirm"):
            results.append(chunk)

        # Should yield confirm_resolved metadata
        assert len(results) == 1
        import json
        parsed = json.loads(results[0]) if isinstance(results[0], str) else results[0]
        assert parsed["metadata"]["type"] == "confirm_resolved"
        assert parsed["metadata"]["result"] == "confirmed"
# #endregion TestAgentChat.Handler.ResumeConfirm


# #region TestAgentChat.Handler.ResumeDeny [C:2] [TYPE Function] [SEMANTICS test,handler,deny]
@pytest.mark.asyncio
async def test_handler_resume_deny():
    """When action='deny', handler yields confirm_resolved with denied."""
    from src.agent.app import agent_handler

    history: list = []
    mock_request = MagicMock()
    token = _make_test_jwt()
    mock_request.headers = {"authorization": f"Bearer {token}"}
    mock_request.client.host = "127.0.0.1"

    message = {"text": "deny", "files": None}

    with patch("src.agent.app.create_agent") as mock_create:
        mock_graph = MagicMock()
        mock_create.return_value = mock_graph

        results = []
        async for chunk in agent_handler(message, history, mock_request, "test-conv", "deny"):
            results.append(chunk)

        assert len(results) == 1
        import json
        parsed = json.loads(results[0]) if isinstance(results[0], str) else results[0]
        assert parsed["metadata"]["type"] == "confirm_resolved"
        assert parsed["metadata"]["result"] == "denied"
# #endregion TestAgentChat.Handler.ResumeDeny
# #endregion TestAgentChat.Handler