ss-tools/agent/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, Mock, patch

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

import httpx
import pytest

os.environ.setdefault("AUTH_SECRET_KEY", "test-secret-key-for-jwt-testing")
os.environ["FASTAPI_URL"] = "http://test-backend:8000"
os.environ["SERVICE_JWT"] = "test-service-jwt"
os.environ["OPENAI_API_KEY"] = "sk-test-key"


def _mock_http_client(get_return=None, post_return=None, get_side_effect=None):
    """Create a mock for get_shared_http_client that returns a mock client.

    Returns (mock_client, patcher) tuple. Use as:
        mock_client, patcher = _mock_http_client(...)
        with patcher:
            ...
    """
    mock_client = AsyncMock(spec=httpx.AsyncClient)
    if get_side_effect is not None:
        mock_client.get = AsyncMock(side_effect=get_side_effect)
    elif get_return is not None:
        mock_client.get = AsyncMock(return_value=get_return)
    if post_return is not None:
        mock_client.post = AsyncMock(return_value=post_return)
    return mock_client, patch("ss_tools.agent.tools.get_shared_http_client", return_value=mock_client)


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


# #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.anyio
async def test_tool_dual_auth_headers():
    """Tools should build auth headers from ContextVar when set."""
    from ss_tools.agent.context import set_service_jwt, set_user_jwt
    from ss_tools.agent.tools import search_dashboards

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

    mock_resp = Mock(status_code=200, text='{"dashboards": [], "total": 0}')
    mock_resp.json.return_value = {"dashboards": [], "total": 0}

    mock_client, patcher = _mock_http_client(get_return=mock_resp)
    with patcher:
        await search_dashboards.ainvoke({"query": "test"})

        # Verify the HTTP request included dual-identity headers
        call_kwargs = mock_client.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"
        assert headers["Authorization"] == "Bearer service-jwt-token"
        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.anyio
async def test_tool_auth_fallback_to_env():
    """Tools should fall back to SERVICE_JWT env var when ContextVar is empty."""
    from ss_tools.agent.context import set_service_jwt, set_user_jwt
    import ss_tools.agent.tools as tools_mod
    from ss_tools.agent.tools import search_dashboards

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

    mock_resp = Mock(status_code=200, text='{"dashboards": [], "total": 0}')
    mock_resp.json.return_value = {"dashboards": [], "total": 0}

    mock_client, patcher = _mock_http_client(get_return=mock_resp)
    with patch.object(tools_mod, "FASTAPI_URL", "http://test-backend:8000"), patcher:
        await search_dashboards.ainvoke({"query": "test"})

        call_kwargs = mock_client.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.anyio
async def test_tool_http_exception_handling():
    """Tool should propagate HTTP exception as error text."""
    from ss_tools.agent.context import set_service_jwt, set_user_jwt
    from ss_tools.agent.tools import search_dashboards

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

    _, patcher = _mock_http_client(get_side_effect=Exception("Connection refused"))
    with patcher:
        # 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 ss_tools.agent.tools import get_all_tools

    tools = get_all_tools()
    tool_names = [t.name for t in tools]
    expected = {
        "show_capabilities",
        "search_dashboards",
        "get_health_summary",
        "list_environments",
        "get_task_status",
        "list_llm_providers",
        "get_llm_status",
        "create_branch",
        "commit_changes",
        "deploy_dashboard",
        "execute_migration",
        "run_backup",
        "run_llm_validation",
        "run_llm_documentation",
        "list_maintenance_events",
        "start_maintenance",
        "end_maintenance",
    }
    assert expected.issubset(set(tool_names))


def test_get_all_tools_args_schema():
    """Tools with args_schema should expose required fields."""
    from ss_tools.agent.tools import get_all_tools

    tools = get_all_tools()
    search_tool = next(t for t in tools if t.name == "search_dashboards")
    health_tool = next(t for t in tools if t.name == "get_health_summary")
    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"
    assert health_tool.args_schema is not None, "get_health_summary should have args_schema"
    assert "env_id" in health_tool.args_schema.model_fields


# ═══════════════════════════════════════════════════════════════════
#  Tool get_all — full catalog (replaces get_tools_for_query)
# ═══════════════════════════════════════════════════════════════════


def test_get_all_tools_returns_24_tools():
    """get_all_tools() returns full catalog — all 24 tools registered."""
    from ss_tools.agent.tools import get_all_tools

    tools = get_all_tools()
    tool_names = {t.name for t in tools}
    # Core tools (always present)
    assert "show_capabilities" in tool_names
    assert "search_dashboards" in tool_names
    assert "get_health_summary" in tool_names
    # Regression: minimum count — must have all 24
    assert len(tools) >= 24, f"Expected ≥24 tools, got {len(tools)}: {tool_names}"


# #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.anyio
async def test_search_dashboards_correct_url():
    """search_dashboards calls GET /api/dashboards with query params."""
    from ss_tools.agent.context import set_service_jwt, set_user_jwt
    from ss_tools.agent.tools import search_dashboards

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

    mock_resp = Mock(status_code=200, text='{"dashboards": [], "total": 0}')
    mock_resp.json.return_value = {"dashboards": [], "total": 0}

    mock_client, patcher = _mock_http_client(get_return=mock_resp)
    with patcher:
        await search_dashboards.ainvoke({"query": "dashboard-name", "env_id": "prod"})

        call_args = mock_client.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
        # Agent searches the full catalog: profile-default filter must be disabled
        # and the page size raised above the default (10) to avoid truncation.
        params = kwargs.get("params", {})
        assert params.get("page_context") == "other"
        assert params.get("page_size") == 100
        assert params.get("q") == "dashboard-name"
        assert params.get("env_id") == "prod"


@pytest.mark.anyio
async def test_search_dashboards_surfaces_profile_filter_hidden():
    """total=0 with available_total>0 must report hidden dashboards, not absence."""
    from ss_tools.agent.context import set_service_jwt, set_user_jwt
    from ss_tools.agent.tools import search_dashboards

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

    payload = {
        "dashboards": [],
        "total": 0,
        "available_total": 11,
        "effective_profile_filter": {
            "applied": True,
            "username": "admin",
            "match_logic": "owners_or_modified_by",
        },
    }
    mock_resp = Mock(status_code=200, text="")
    mock_resp.json.return_value = payload

    _, patcher = _mock_http_client(get_return=mock_resp)
    with patcher:
        result = await search_dashboards.ainvoke({"query": "", "env_id": "ss-dev"})

    assert "11" in result
    assert "hidden" in result
    assert "My Dashboards Only" in result
    assert "admin" in result
    assert "no dashboards found" not in result.lower() or "hidden" in result.lower()


@pytest.mark.anyio
async def test_search_dashboards_truncation_note():
    """When the API returns fewer items than total, the tool must say so."""
    from ss_tools.agent.context import set_service_jwt, set_user_jwt
    from ss_tools.agent.tools import search_dashboards

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

    payload = {
        "dashboards": [
            {"title": f"Dashboard {i}", "owners": ["admin"], "last_modified": "2026-01-01T00:00:00"} for i in range(10)
        ],
        "total": 15,
        "available_total": 15,
        "effective_profile_filter": {"applied": False, "username": None},
    }
    mock_resp = Mock(status_code=200, text="")
    mock_resp.json.return_value = payload

    _, patcher = _mock_http_client(get_return=mock_resp)
    with patcher:
        result = await search_dashboards.ainvoke({"query": "", "env_id": "prod"})

    assert "Found 15 dashboard(s)" in result
    assert "showing first 10 of 15" in result


@pytest.mark.anyio
async def test_search_dashboards_empty_env_still_reports_absent():
    """A genuinely empty environment must still produce the 'no dashboards' message."""
    from ss_tools.agent.context import set_service_jwt, set_user_jwt
    from ss_tools.agent.tools import search_dashboards

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

    payload = {
        "dashboards": [],
        "total": 0,
        "available_total": 0,
        "effective_profile_filter": {"applied": False, "username": None},
    }
    mock_resp = Mock(status_code=200, text="")
    mock_resp.json.return_value = payload

    _, patcher = _mock_http_client(get_return=mock_resp)
    with patcher:
        result = await search_dashboards.ainvoke({"query": "x", "env_id": "prod"})

    assert "No dashboards found matching 'x'" in result


# #endregion TestAgentChat.Tools.ToolContracts


# #region TestAgentChat.Tools.HealthSummary [C:2] [TYPE Function] [SEMANTICS test,tools,health]
# @BRIEF get_health_summary calls the correct FastAPI endpoint.


@pytest.mark.anyio
async def test_get_health_summary_calls_correct_url():
    """get_health_summary should call GET /api/health/summary."""
    from ss_tools.agent.context import set_service_jwt, set_user_jwt
    from ss_tools.agent.tools import get_health_summary

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

    mock_resp = Mock(status_code=200, text='{"status": "ok"}')

    mock_client, patcher = _mock_http_client(get_return=mock_resp)
    with patcher:
        await get_health_summary.ainvoke({"env_id": "ss-dev"})

        call_args = mock_client.get.call_args
        assert call_args is not None
        args, kwargs = call_args
        url = args[0] if args else kwargs.get("url", "")
        assert "api/health/summary" in url
        assert kwargs.get("params") == {"environment_id": "ss-dev"}


# #endregion TestAgentChat.Tools.HealthSummary


# #region TestAgentChat.Tools.ListEnvironments [C:2] [TYPE Function] [SEMANTICS test,tools,environments]
# @BRIEF list_environments calls the correct FastAPI endpoint.


@pytest.mark.anyio
async def test_list_environments_calls_correct_url():
    """list_environments should call GET /api/settings/environments."""
    from ss_tools.agent.context import set_service_jwt, set_user_jwt
    from ss_tools.agent.tools import list_environments

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

    mock_resp = Mock(status_code=200, text='["prod", "dev"]')

    mock_client, patcher = _mock_http_client(get_return=mock_resp)
    with patcher:
        await list_environments.ainvoke({})

        call_args = mock_client.get.call_args
        assert call_args is not None
        args, kwargs = call_args
        url = args[0] if args else kwargs.get("url", "")
        assert "api/settings/environments" in url


@pytest.mark.anyio
async def test_list_environments_redacts_sensitive_fields():
    """list_environments must not expose backend secrets to chat output."""
    from ss_tools.agent.context import set_service_jwt, set_user_jwt
    from ss_tools.agent.tools import list_environments

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

    mock_resp = Mock(
        status_code=200,
        text='[{"id":"prod","password":"secret-pass","api_key":"secret-key","nested":{"token":"secret-token"},"name":"ss-prod"}]',
    )

    _, patcher = _mock_http_client(get_return=mock_resp)
    with patcher:
        result = await list_environments.ainvoke({})

    assert "secret-pass" not in result
    assert "secret-key" not in result
    assert "secret-token" not in result
    assert result.count("[redacted]") == 3


# #endregion TestAgentChat.Tools.ListEnvironments


# #region TestAgentChat.Tools.TaskStatus [C:2] [TYPE Function] [SEMANTICS test,tools,task]
# @BRIEF get_task_status calls the correct FastAPI endpoint with task_id.


@pytest.mark.anyio
async def test_get_task_status_calls_correct_url():
    """get_task_status should call GET /api/tasks/{task_id}."""
    from ss_tools.agent.context import set_service_jwt, set_user_jwt
    from ss_tools.agent.tools import get_task_status

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

    mock_resp = Mock(status_code=200, text='{"status": "running"}')

    mock_client, patcher = _mock_http_client(get_return=mock_resp)
    with patcher:
        await get_task_status.ainvoke({"task_id": "task-123"})

        call_args = mock_client.get.call_args
        assert call_args is not None
        args, kwargs = call_args
        url = args[0] if args else kwargs.get("url", "")
        assert "api/tasks/task-123" in url


# #endregion TestAgentChat.Tools.TaskStatus


@pytest.mark.anyio
async def test_run_backup_posts_task_payload():
    """run_backup should create a superset-backup task through /api/tasks."""
    from ss_tools.agent.context import set_service_jwt, set_user_jwt, set_user_role
    from ss_tools.agent.tools import run_backup

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

    mock_resp = Mock(status_code=201, text='{"id": "task-1"}')

    mock_client, patcher = _mock_http_client(post_return=mock_resp)
    with patcher:
        await run_backup.ainvoke({"environment_id": "prod", "dashboard_id": 10})

        call_args = mock_client.post.call_args
        assert call_args is not None
        args, kwargs = call_args
        url = args[0] if args else kwargs.get("url", "")
        assert "api/tasks" in url
        assert kwargs["json"] == {
            "plugin_id": "superset-backup",
            "params": {"environment_id": "prod", "dashboard_ids": [10]},
        }


@pytest.mark.anyio
async def test_deploy_dashboard_posts_git_endpoint():
    """deploy_dashboard should call the native Git deploy API."""
    from ss_tools.agent.context import set_service_jwt, set_user_jwt, set_user_role
    from ss_tools.agent.tools import deploy_dashboard

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

    mock_resp = Mock(status_code=200, text='{"status": "success"}')

    mock_client, patcher = _mock_http_client(post_return=mock_resp)
    with patcher:
        await deploy_dashboard.ainvoke({"dashboard_ref": "42", "environment_id": "prod"})

        call_args = mock_client.post.call_args
        assert call_args is not None
        args, kwargs = call_args
        url = args[0] if args else kwargs.get("url", "")
        assert "api/git/repositories/42/deploy" in url
        assert kwargs["json"] == {"environment_id": "prod"}


# #region TestAgentChat.Tools.DualAuthHeaders [C:2] [TYPE Function] [SEMANTICS test,tools,auth,headers]
# @BRIEF _dual_auth_headers builds proper headers from ContextVars.


def test_dual_auth_headers_with_both_jwts():
    """_dual_auth_headers uses service auth plus user identity when both are set."""
    from ss_tools.agent.context import set_service_jwt, set_user_jwt
    from ss_tools.agent.tools import _dual_auth_headers

    set_service_jwt("svc-token")
    set_user_jwt("user-token")

    headers = _dual_auth_headers()
    assert headers.get("Authorization") == "Bearer svc-token"
    assert headers.get("X-User-JWT") == "user-token"


def test_dual_auth_headers_no_user_jwt():
    """_dual_auth_headers falls back to service Authorization when no user JWT."""
    from ss_tools.agent.context import set_service_jwt, set_user_jwt
    from ss_tools.agent.tools import _dual_auth_headers

    set_service_jwt("svc-token")
    set_user_jwt("")

    headers = _dual_auth_headers()
    assert headers.get("Authorization") == "Bearer svc-token"


def test_dual_auth_headers_no_jwts():
    """_dual_auth_headers falls back to _SERVICE_JWT when context vars are empty."""
    from ss_tools.agent.context import set_service_jwt, set_user_jwt
    from ss_tools.agent.tools import _dual_auth_headers

    set_service_jwt("")
    set_user_jwt("")

    headers = _dual_auth_headers()
    # Context vars are empty, _SERVICE_JWT is module-level constant from _config
    # (set at import time based on os.environ)
    assert "Authorization" in headers
    assert headers["Authorization"].startswith("Bearer ")


# #endregion TestAgentChat.Tools.DualAuthHeaders
# #endregion TestAgentChat.Tools