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
from pathlib import Path
import sys

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

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

import jwt

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


@pytest.fixture
def anyio_backend():
    return "asyncio"


@pytest.fixture(autouse=True)
def mock_save_conversation():
    with patch("src.agent.app.save_conversation", new_callable=AsyncMock):
        yield


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


def _empty_agent_state():
    state = MagicMock(spec_set=["next"])
    state.next = ()
    return state


# #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.anyio
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.langgraph_setup.create_agent"):
        # 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.AuthGraceful [C:2] [TYPE Function] [SEMANTICS test,handler,auth]
# @BRIEF Missing or invalid JWT does NOT reject — Gradio handler forwards to graph (auth at tool layer).
# @TEST_EDGE missing_auth, invalid_token
# @RATIONALE Per design, @gradio/client does not forward Authorization headers, so the Gradio handler
#            does NOT enforce JWT. Missing/invalid JWT falls back to anonymous context.
#            Tool-level auth is enforced via SERVICE_JWT + X-User-JWT dual identity pattern.
@pytest.mark.anyio
async def test_handler_missing_auth_continues_gracefully():
    """Missing authorization header does NOT yield UNAUTHORIZED — handler continues."""
    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}

    # Patch create_agent to prevent LLM call — handler should proceed without JWT
    with patch("src.agent.app.create_agent") as mock_create:
        mock_graph = AsyncMock()
        async def _empty_stream(*_args, **_kwargs):
            """Empty async generator — yields nothing."""
            return
            yield  # make it a generator
        mock_graph.astream_events = _empty_stream
        mock_graph.aget_state = AsyncMock(return_value=_empty_agent_state())
        mock_create.return_value = mock_graph

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

        # No UNAUTHORIZED error — handler proceeds with empty user_jwt
        for r in results:
            import json
            parsed = json.loads(r) if isinstance(r, str) else r
            meta = parsed.get("metadata", {})
            assert meta.get("code") != "UNAUTHORIZED", \
                "Handler should not reject missing auth — JWT optional at Gradio layer"


@pytest.mark.anyio
async def test_handler_invalid_jwt_continues_gracefully():
    """Invalid JWT does NOT yield UNAUTHORIZED — handler continues with fallback context."""
    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}

    with patch("src.agent.app.create_agent") as mock_create:
        mock_graph = AsyncMock()
        async def _empty_stream(*_args, **_kwargs):
            return
            yield
        mock_graph.astream_events = _empty_stream
        mock_graph.aget_state = AsyncMock(return_value=_empty_agent_state())
        mock_create.return_value = mock_graph

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

        for r in results:
            import json
            parsed = json.loads(r) if isinstance(r, str) else r
            meta = parsed.get("metadata", {})
            assert meta.get("code") != "UNAUTHORIZED", \
                "Handler should not reject invalid JWT — invalid token ignored at Gradio layer"
# #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.anyio
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_graph.aget_state = AsyncMock(return_value=_empty_agent_state())
        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.anyio
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}

    # Patch create_agent in _confirmation because handle_resume (now in _confirmation.py)
    # imports create_agent directly from langgraph_setup
    with patch("src.agent._confirmation.create_agent") as mock_create:
        mock_graph = MagicMock()
        async def _empty_stream(*_args, **_kwargs):
            return
            yield
        mock_graph.astream_events = _empty_stream
        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"
        assert mock_create.call_args.kwargs["interrupt_before"] == []
# #endregion TestAgentChat.Handler.ResumeConfirm


# #region TestAgentChat.Handler.ResumeDeny [C:2] [TYPE Function] [SEMANTICS test,handler,deny]
@pytest.mark.anyio
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._confirmation.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