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

# #region Test.AgentChat.GradioApp [C:3] [TYPE Module] [SEMANTICS test,agent,gradio,chat]
# @BRIEF Tests for agent/app.py — agent_handler, _handle_resume, _extract_user_id, _save_conversation, create_chat_interface, health.
# @RELATION BINDS_TO -> [AgentChat.GradioApp]

from pathlib import Path
import sys

sys.path.insert(0, str(Path(__file__).parent.parent.parent / "src"))

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


def _make_async_iter(items):
    """Create an async iterator from a list."""
    class AsyncIter:
        def __init__(self, items):
            self._iter = iter(items)
        def __aiter__(self):
            return self
        async def __anext__(self):
            try:
                return next(self._iter)
            except StopIteration:
                raise StopAsyncIteration
    return AsyncIter(items)


def _make_agent_mock(stream_events=None, raise_on_call=False):
    """Create a properly mocked agent for async iteration."""
    agent = MagicMock()
    if raise_on_call:
        agent.astream_events = MagicMock(side_effect=stream_events if isinstance(stream_events, list) else stream_events)
    elif stream_events is not None:
        agent.astream_events = MagicMock(return_value=_make_async_iter(stream_events))
    else:
        agent.astream_events = MagicMock(return_value=_make_async_iter([]))
    return agent


@pytest.fixture(autouse=True)
def clear_locks():
    """Clear _user_locks before each test."""
    from src.agent.app import _user_locks
    _user_locks.clear()


@pytest.fixture
def mock_request():
    req = MagicMock()
    req.headers = {"authorization": ""}
    return req


# #region test_extract_user_id [C:2] [TYPE Function]
# @BRIEF Test _extract_user_id for various JWT payloads.
class TestExtractUserId:
    def test_extracts_sub(self):
        from src.agent.app import _extract_user_id
        with patch("src.agent.app.decode_token", return_value={"sub": "user-1"}):
            assert _extract_user_id("fake-jwt") == "user-1"

    def test_extracts_user_id_fallback(self):
        from src.agent.app import _extract_user_id
        with patch("src.agent.app.decode_token", return_value={"user_id": "user-2"}):
            assert _extract_user_id("fake-jwt") == "user-2"

    def test_returns_unknown_on_exception(self):
        from src.agent.app import _extract_user_id
        with patch("src.agent.app.decode_token", side_effect=Exception("bad token")):
            assert _extract_user_id("bad") == "unknown"

    def test_returns_unknown_on_empty(self):
        from src.agent.app import _extract_user_id
        assert _extract_user_id("") == "unknown"
# #endregion test_extract_user_id


# #region test_agent_handler [C:2] [TYPE Function]
# @BRIEF Test agent_handler for various scenarios.
class TestAgentHandler:
    @pytest.mark.asyncio
    async def test_handles_concurrent_send(self, mock_request):
        from src.agent.app import agent_handler, _user_locks
        _user_locks["anon_default"] = True
        results = [r async for r in agent_handler("hello", [], mock_request, None, None)]
        assert len(results) == 1
        data = json.loads(results[0])
        assert data["metadata"]["code"] == "CONCURRENT_SEND"

    @pytest.mark.asyncio
    async def test_handles_file_too_large(self, mock_request):
        from src.agent.app import agent_handler, MAX_FILE_SIZE_BYTES
        message = {"text": "analyze", "files": ["fake_path"]}
        with patch("os.path.exists", return_value=True), \
             patch("os.path.getsize", return_value=MAX_FILE_SIZE_BYTES + 1):
            results = [r async for r in agent_handler(message, [], mock_request, None, None)]
            assert len(results) == 1
            data = json.loads(results[0])
            assert data["metadata"]["code"] == "FILE_TOO_LARGE"

    @pytest.mark.asyncio
    async def test_handles_hitl_confirm(self, mock_request):
        from src.agent.app import agent_handler
        with patch("src.agent.app._handle_resume") as mock_resume:
            mock_resume.return_value = _make_async_iter([
                json.dumps({"content": "confirmed", "metadata": {"type": "confirm_resolved"}})
            ])
            with patch("src.agent.app._save_conversation", AsyncMock()):
                results = [r async for r in agent_handler("confirm", [], mock_request, "conv-1", "confirm")]
                assert len(results) == 1
                data = json.loads(results[0])
                assert data["metadata"]["type"] == "confirm_resolved"

    @pytest.mark.asyncio
    async def test_handles_hitl_deny(self, mock_request):
        from src.agent.app import agent_handler
        with patch("src.agent.app._handle_resume") as mock_resume:
            mock_resume.return_value = _make_async_iter([
                json.dumps({"content": "denied", "metadata": {"type": "confirm_resolved", "result": "denied"}})
            ])
            with patch("src.agent.app._save_conversation", AsyncMock()):
                results = [r async for r in agent_handler("deny", [], mock_request, "conv-1", "deny")]
                assert len(results) == 1
                data = json.loads(results[0])
                assert data["metadata"]["result"] == "denied"

    @pytest.mark.asyncio
    async def test_normal_send_yields_tokens(self, mock_request):
        from src.agent.app import agent_handler
        mock_chunk = MagicMock()
        mock_chunk.content = "Hello"

        mock_event_stream = [
            {"event": "on_chat_model_stream", "data": {"chunk": mock_chunk}},
            {"event": "on_chain_end", "interrupt": {}},
        ]

        agent = _make_agent_mock(mock_event_stream)
        with patch("src.agent.app.create_agent", return_value=agent), \
             patch("src.agent.app.get_all_tools", return_value=[]), \
             patch("src.agent.app._save_conversation", AsyncMock()):
            results = [r async for r in agent_handler("hello", [], mock_request, None, None)]
            assert len(results) > 0
            token_data = json.loads(results[0])
            assert token_data["metadata"]["type"] == "stream_token"

    @pytest.mark.asyncio
    async def test_tool_events_yielded(self, mock_request):
        from src.agent.app import agent_handler
        mock_event_stream = [
            {"event": "on_tool_start", "name": "test_tool", "data": {"input": {}}},
            {"event": "on_tool_end", "name": "test_tool", "data": {"output": "done"}},
        ]

        agent = _make_agent_mock(mock_event_stream)
        with patch("src.agent.app.create_agent", return_value=agent), \
             patch("src.agent.app.get_all_tools", return_value=[]), \
             patch("src.agent.app._save_conversation", AsyncMock()), \
             patch("src.agent.app.log_tool_event", AsyncMock()):
            results = [r async for r in agent_handler("hello", [], mock_request, None, None)]
            assert len(results) == 2
            assert json.loads(results[0])["metadata"]["type"] == "tool_start"
            assert json.loads(results[1])["metadata"]["type"] == "tool_end"

    @pytest.mark.asyncio
    async def test_tool_error_event(self, mock_request):
        from src.agent.app import agent_handler
        mock_event_stream = [
            {"event": "on_tool_error", "name": "test_tool", "data": {"error": "something failed"}},
        ]

        agent = _make_agent_mock(mock_event_stream)
        with patch("src.agent.app.create_agent", return_value=agent), \
             patch("src.agent.app.get_all_tools", return_value=[]), \
             patch("src.agent.app._save_conversation", AsyncMock()), \
             patch("src.agent.app.log_tool_event", AsyncMock()):
            results = [r async for r in agent_handler("hello", [], mock_request, None, None)]
            assert len(results) == 1
            assert json.loads(results[0])["metadata"]["type"] == "tool_error"

    @pytest.mark.asyncio
    async def test_output_parser_exception_retry(self, mock_request):
        from src.agent.app import agent_handler
        from langchain_core.exceptions import OutputParserException

        mock_stream_after = _make_async_iter([
            {"event": "on_chat_model_stream", "data": {"chunk": MagicMock(content="retried")}}
        ])
        call_count = [0]

        def mock_astream(*args, **kwargs):
            call_count[0] += 1
            if call_count[0] == 1:
                raise OutputParserException("bad output")
            return mock_stream_after

        agent = MagicMock()
        agent.astream_events = mock_astream

        with patch("src.agent.app.create_agent", return_value=agent), \
             patch("src.agent.app.get_all_tools", return_value=[]), \
             patch("src.agent.app._save_conversation", AsyncMock()):
            results = [r async for r in agent_handler("hello", [], mock_request, None, None)]
            assert len(results) > 0

    @pytest.mark.asyncio
    async def test_output_parser_exception_final_failure(self, mock_request):
        from src.agent.app import agent_handler
        from langchain_core.exceptions import OutputParserException

        call_count = [0]

        def mock_astream(*args, **kwargs):
            call_count[0] += 1
            raise OutputParserException(f"bad {call_count[0]}")

        agent = MagicMock()
        agent.astream_events = mock_astream

        with patch("src.agent.app.create_agent", return_value=agent), \
             patch("src.agent.app.get_all_tools", return_value=[]), \
             patch("src.agent.app._save_conversation", AsyncMock()):
            results = [r async for r in agent_handler("hello", [], mock_request, None, None)]
            assert len(results) == 1
            data = json.loads(results[0])
            assert data["metadata"]["code"] == "LLM_MALFORMED_OUTPUT"

    @pytest.mark.asyncio
    async def test_non_llm_error_saves_conversation(self, mock_request):
        from src.agent.app import agent_handler

        def mock_astream(*args, **kwargs):
            raise RuntimeError("API connection error")

        agent = MagicMock()
        agent.astream_events = mock_astream

        save_mock = AsyncMock()
        with patch("src.agent.app.create_agent", return_value=agent), \
             patch("src.agent.app.get_all_tools", return_value=[]), \
             patch("src.agent.app._save_conversation", save_mock):
            with pytest.raises(RuntimeError, match="API connection error"):
                await agent_handler("hello", [], mock_request, None, None).__anext__()
            save_mock.assert_called_once()

    @pytest.mark.asyncio
    async def test_invalid_jwt_passes_gracefully(self):
        from src.agent.app import agent_handler
        from jose import JWTError
        req = MagicMock()
        req.headers = {"authorization": "Bearer invalid.jwt"}
        mock_event_stream = [
            {"event": "on_chat_model_stream", "data": {"chunk": MagicMock(content="ok")}}
        ]
        agent = _make_agent_mock(mock_event_stream)
        with patch("src.agent.app.decode_token", side_effect=JWTError("invalid")), \
             patch("src.agent.app.create_agent", return_value=agent), \
             patch("src.agent.app.get_all_tools", return_value=[]), \
             patch("src.agent.app._save_conversation", AsyncMock()):
            results = [r async for r in agent_handler("hi", [], req, None, None)]
            assert len(results) == 1

    @pytest.mark.asyncio
    async def test_valid_user_jwt_with_audience_is_kept_for_tool_auth(self, mock_request):
        from src.agent.app import agent_handler
        from src.agent.context import get_user_jwt
        from src.core.auth.jwt import create_access_token

        token = create_access_token({"sub": "admin", "scopes": ["Admin"]})
        mock_event_stream = [
            {"event": "on_chat_model_stream", "data": {"chunk": MagicMock(content="ok")}}
        ]
        agent = _make_agent_mock(mock_event_stream)

        with patch("src.agent.app.create_agent", return_value=agent), \
             patch("src.agent.app.get_all_tools", return_value=[]), \
             patch("src.agent.app._save_conversation", AsyncMock()):
            results = [r async for r in agent_handler("hi", [], mock_request, None, None, None, token)]

        assert len(results) == 1
        assert get_user_jwt() == token
# #endregion test_agent_handler


# #region test_handle_resume [C:2] [TYPE Function]
# @BRIEF Test _handle_resume confirm and deny paths.
class TestHandleResume:
    @pytest.mark.asyncio
    async def test_confirm(self):
        from src.agent.app import _handle_resume
        mock_agent = MagicMock()
        with patch("src.agent.app.create_agent", return_value=mock_agent), \
             patch("src.agent.app.get_all_tools", return_value=[]):
            results = [r async for r in _handle_resume("conv-1", "confirm")]
            assert len(results) == 1
            data = json.loads(results[0])
            assert data["metadata"]["result"] == "confirmed"

    @pytest.mark.asyncio
    async def test_deny(self):
        from src.agent.app import _handle_resume
        mock_agent = MagicMock()
        with patch("src.agent.app.create_agent", return_value=mock_agent), \
             patch("src.agent.app.get_all_tools", return_value=[]):
            results = [r async for r in _handle_resume("conv-1", "deny")]
            assert len(results) == 1
            data = json.loads(results[0])
            assert data["metadata"]["result"] == "denied"
# #endregion test_handle_resume


# #region test_save_conversation [C:2] [TYPE Function]
# @BRIEF Test _save_conversation with various scenarios.
class TestSaveConversation:
    @pytest.mark.asyncio
    async def test_save_success(self):
        from src.agent.app import _save_conversation
        with patch("src.agent.app.httpx.AsyncClient") as mock_client:
            mock_client.return_value.__aenter__.return_value.post = AsyncMock()
            await _save_conversation("conv-1", "test message", "user-1")
            mock_client.return_value.__aenter__.return_value.post.assert_called_once()

    @pytest.mark.asyncio
    async def test_save_with_service_jwt(self):
        from src.agent.app import _save_conversation
        with patch("src.agent.app.httpx.AsyncClient") as mock_client, \
             patch("src.agent.app.os.getenv", return_value="service-token"):
            mock_client.return_value.__aenter__.return_value.post = AsyncMock()
            await _save_conversation("conv-1", "hello", "admin")
            mock_client.return_value.__aenter__.return_value.post.assert_called_once()

    @pytest.mark.asyncio
    async def test_save_failure_logged(self):
        from src.agent.app import _save_conversation
        with patch("src.agent.app.httpx.AsyncClient") as mock_client:
            mock_client.return_value.__aenter__.return_value.post.side_effect = Exception("network err")
            # Should not raise
            await _save_conversation("conv-1", "msg", "u1")

    @pytest.mark.asyncio
    async def test_save_empty_title(self):
        from src.agent.app import _save_conversation
        with patch("src.agent.app.httpx.AsyncClient") as mock_client:
            client_instance = AsyncMock()
            mock_client.return_value.__aenter__.return_value = client_instance
            await _save_conversation("conv-1", "   ", "user-1")
            call_kwargs = client_instance.post.call_args[1]
            assert call_kwargs["json"]["title"] == "Agent conversation"
# #endregion test_save_conversation


# #region test_create_chat_interface [C:2] [TYPE Function]
# @BRIEF Test create_chat_interface returns a gr.ChatInterface.
class TestCreateChatInterface:
    def test_returns_chat_interface(self):
        from src.agent.app import create_chat_interface
        with patch("src.agent.app.gr.ChatInterface") as mock_ci:
            result = create_chat_interface()
            assert result is mock_ci.return_value
# #endregion test_create_chat_interface


# #region test_health [C:2] [TYPE Function]
# @BRIEF Test health endpoint returns status ok.
class TestHealth:
    @pytest.mark.asyncio
    async def test_health_returns_ok(self):
        from src.agent.app import health
        result = await health()
        assert result["status"] == "ok"
# #endregion test_health


# #region test_file_upload_parsing [C:2] [TYPE Function]
# @BRIEF Test file upload branch — parse_upload called for valid small files.
class TestFileUploadParsing:
    @pytest.mark.asyncio
    async def test_handler_parses_uploaded_file(self, mock_request, tmp_path):
        """Valid small file triggers parse_upload and continues to stream."""
        from src.agent.app import agent_handler
        test_file = tmp_path / "test.txt"
        test_file.write_text("upload content for analysis")

        message = {"text": "analyze", "files": [str(test_file)]}

        with patch('src.agent.app.create_agent') as mock_create, \
             patch('src.agent.app.get_all_tools', return_value=[]), \
             patch('src.agent.app._save_conversation', AsyncMock()), \
             patch('src.agent.app.log_tool_event', AsyncMock()):
            agent = _make_agent_mock([
                {"event": "on_chat_model_stream", "data": {"chunk": MagicMock(content="result")}}
            ])
            mock_create.return_value = agent
            results = [r async for r in agent_handler(message, [], mock_request, None, None)]
            assert len(results) > 0
            data = json.loads(results[0])
            assert data["metadata"]["type"] == "stream_token"
# #endregion test_file_upload_parsing


# #region test_app_main_block [C:2] [TYPE Function]
# @BRIEF Test if __name__ == '__main__' block in app.py.
class TestAppMainBlock:
    def test_app_main_block(self):
        """if __name__ == '__main__' creates ChatInterface and launches."""
        import importlib.util
        from pathlib import Path

        app_path = Path(__file__).parent.parent.parent / "src" / "agent" / "app.py"
        spec = importlib.util.spec_from_file_location("__main__", str(app_path))

        mock_demo = MagicMock()
        with patch('gradio.ChatInterface') as mock_ci:
            mock_ci.return_value = mock_demo

            module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(module)

        mock_demo.launch.assert_called_once()
# #endregion test_app_main_block
# #endregion Test.AgentChat.GradioApp