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

# #region TestAgentChat.Tools [C:3] [TYPE Module] [SEMANTICS test,agent,tools,langchain]
# @BRIEF Tests for LangChain @tool functions — dual-identity auth, HTTP calls, tool wrapping.
# @RELATION BINDS_TO -> [AgentChat.Tools]
# @TEST_EDGE: tool_rest_call -> tool calls FastAPI with dual-identity headers
# @TEST_EDGE: tool_http_failure -> tool returns error JSON gracefully
# @TEST_EDGE: get_all_tools -> returns expected tool list
import os
from pathlib import Path
import sys
from unittest.mock import AsyncMock, patch

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

import pytest

os.environ["JWT_SECRET"] = "test-jwt-secret-key"
os.environ["FASTAPI_URL"] = "http://test-backend:8000"
os.environ["SERVICE_JWT"] = "test-service-jwt"
os.environ["OPENAI_API_KEY"] = "sk-test-key"


# #region TestAgentChat.Tools.DualAuth [C:2] [TYPE Function] [SEMANTICS test,tools,auth]
# @BRIEF Dual-identity auth headers built from ContextVar and env vars.

@pytest.mark.asyncio
async def test_tool_dual_auth_headers():
    """Tools should build auth headers from ContextVar when set."""
    from src.agent.context import set_service_jwt, set_user_jwt
    from src.agent.tools import search_dashboards

    # Set JWTs in context
    set_user_jwt("user-jwt-token")
    set_service_jwt("service-jwt-token")

    with patch("httpx.AsyncClient") as mock_client:
        mock_instance = AsyncMock()
        mock_client.return_value.__aenter__.return_value = mock_instance
        mock_instance.get.return_value.text = '{"dashboards": []}'

        await search_dashboards.ainvoke({"query": "test"})

        # Verify the HTTP request included dual-identity headers
        call_kwargs = mock_instance.get.call_args
        assert call_kwargs is not None, "HTTP GET should have been called"
        _, kwargs = call_kwargs
        headers = kwargs.get("headers", {})
        assert "Authorization" in headers, "Should include Authorization header (service JWT)"
        assert headers["Authorization"] == "Bearer service-jwt-token"
        assert "X-User-JWT" in headers, "Should include X-User-JWT header (user JWT)"
        assert headers["X-User-JWT"] == "user-jwt-token"
# #endregion TestAgentChat.Tools.DualAuth


# #region TestAgentChat.Tools.FallbackAuth [C:2] [TYPE Function] [SEMANTICS test,tools,auth,fallback]
# @BRIEF Dual-identity auth falls back to env var when ContextVar is not set.

@pytest.mark.asyncio
async def test_tool_auth_fallback_to_env():
    """Tools should fall back to SERVICE_JWT env var when ContextVar is empty."""
    from src.agent.context import set_service_jwt, set_user_jwt
    import src.agent.tools as tools_mod
    from src.agent.tools import search_dashboards

    # Clear ContextVars
    set_user_jwt("")
    set_service_jwt("")
    os.environ["SERVICE_JWT"] = "env-service-token"

    with patch.object(tools_mod, "FASTAPI_URL", "http://test-backend:8000"), \
         patch("httpx.AsyncClient") as mock_client:
        mock_instance = AsyncMock()
        mock_client.return_value.__aenter__.return_value = mock_instance
        mock_instance.get.return_value.text = '{"dashboards": []}'

        # Since tool uses os.getenv at call time, the env var will be read
        await search_dashboards.ainvoke({"query": "test"})

        call_kwargs = mock_instance.get.call_args
        assert call_kwargs is not None
        _, kwargs = call_kwargs
        headers = kwargs.get("headers", {})
        # Should use env var
        assert "Authorization" in headers
# #endregion TestAgentChat.Tools.FallbackAuth


# #region TestAgentChat.Tools.HttpFailure [C:2] [TYPE Function] [SEMANTICS test,tools,failure]
# @BRIEF Tool handles HTTP failure gracefully (returns error text, not exception).

@pytest.mark.asyncio
async def test_tool_http_exception_handling():
    """Tool should propagate HTTP exception as error text."""
    from src.agent.context import set_service_jwt, set_user_jwt
    from src.agent.tools import search_dashboards

    set_user_jwt("test-jwt")
    set_service_jwt("svc-jwt")

    with patch("httpx.AsyncClient") as mock_client:
        mock_instance = AsyncMock()
        mock_client.return_value.__aenter__.return_value = mock_instance
        mock_instance.get.side_effect = Exception("Connection refused")

        # Should propagate the exception (caller handles error)
        with pytest.raises((Exception,)):
            await search_dashboards.ainvoke({"query": "test"})
# #endregion TestAgentChat.Tools.HttpFailure


# #region TestAgentChat.Tools.GetAll [C:2] [TYPE Function] [SEMANTICS test,tools,registry]
# @BRIEF get_all_tools returns the expected list of tool functions.

def test_get_all_tools_returns_expected_list():
    """get_all_tools() should return search_dashboards, get_health_summary, etc."""
    from src.agent.tools import get_all_tools

    tools = get_all_tools()
    assert len(tools) >= 4, f"Expected at least 4 tools, got {len(tools)}"

    tool_names = [t.name for t in tools]
    assert "search_dashboards" in tool_names
    assert "get_health_summary" in tool_names
    assert "list_environments" in tool_names
    assert "get_task_status" in tool_names


def test_get_all_tools_args_schema():
    """Tools with args_schema should have SearchDashboardsInput."""
    from src.agent.tools import get_all_tools

    tools = get_all_tools()
    search_tool = next(t for t in tools if t.name == "search_dashboards")
    assert search_tool.args_schema is not None, "search_dashboards should have args_schema"
    schema_fields = search_tool.args_schema.model_fields
    assert "query" in schema_fields, "search_dashboards should have 'query' field"
    assert schema_fields["query"].is_required(), "query should be required"
# #endregion TestAgentChat.Tools.GetAll


# #region TestAgentChat.Tools.ToolContracts [C:2] [TYPE Function] [SEMANTICS test,tools,contract]
# @BRIEF Tool contracts match @POST and @PRE declared in contracts/modules.md.

@pytest.mark.asyncio
async def test_search_dashboards_correct_url():
    """search_dashboards calls GET /api/dashboards with query params."""
    from src.agent.context import set_service_jwt, set_user_jwt
    from src.agent.tools import search_dashboards

    set_user_jwt("jwt")
    set_service_jwt("svc-jwt")

    with patch("httpx.AsyncClient") as mock_client:
        mock_instance = AsyncMock()
        mock_client.return_value.__aenter__.return_value = mock_instance
        mock_instance.get.return_value.text = '{"data": []}'

        await search_dashboards.ainvoke({"query": "dashboard-name", "env_id": "prod"})

        call_args = mock_instance.get.call_args
        assert call_args is not None
        args, kwargs = call_args
        url = args[0] if args else kwargs.get("url", "")
        assert "api/dashboards" in url
# #endregion TestAgentChat.Tools.ToolContracts
# #endregion TestAgentChat.Tools