refactor(agent): extract agent+shared into standalone packages with full GRACE semantic markup
- Move agent code from backend/src/agent/ to agent/src/ss_tools/agent/ - Extract shared stdlib-only utilities to shared/src/ss_tools/shared/ - Add #region/#endregion contracts to all ~140 functions (INV_1 compliance) - Update docker files, entrypoint, build scripts for new package layout - Backend now imports ss_tools.shared._llm_health (no gradio/langchain deps) - Add specs for 036-039 feature plans
This commit is contained in:
446
agent/tests/test_agent/test_langchain_tools.py
Normal file
446
agent/tests/test_agent/test_langchain_tools.py
Normal file
@@ -0,0 +1,446 @@
|
||||
# #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 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"
|
||||
|
||||
|
||||
@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")
|
||||
|
||||
with patch("httpx.AsyncClient") as mock_client:
|
||||
mock_instance = AsyncMock()
|
||||
mock_client.return_value.__aenter__.return_value = mock_instance
|
||||
mock_instance.get.return_value = Mock(
|
||||
status_code=200,
|
||||
text='{"dashboards": [], "total": 0}',
|
||||
)
|
||||
mock_instance.get.return_value.json.return_value = {"dashboards": [], "total": 0}
|
||||
|
||||
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"
|
||||
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"
|
||||
|
||||
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 = Mock(
|
||||
status_code=200,
|
||||
text='{"dashboards": [], "total": 0}',
|
||||
)
|
||||
mock_instance.get.return_value.json.return_value = {"dashboards": [], "total": 0}
|
||||
|
||||
# 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.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")
|
||||
|
||||
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 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")
|
||||
|
||||
with patch("httpx.AsyncClient") as mock_client:
|
||||
mock_instance = AsyncMock()
|
||||
mock_client.return_value.__aenter__.return_value = mock_instance
|
||||
mock_instance.get.return_value = Mock(
|
||||
status_code=200,
|
||||
text='{"dashboards": [], "total": 0}',
|
||||
)
|
||||
mock_instance.get.return_value.json.return_value = {"dashboards": [], "total": 0}
|
||||
|
||||
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
|
||||
|
||||
|
||||
# #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")
|
||||
|
||||
with patch("httpx.AsyncClient") as mock_client:
|
||||
mock_instance = AsyncMock()
|
||||
mock_client.return_value.__aenter__.return_value = mock_instance
|
||||
mock_instance.get.return_value.status_code = 200
|
||||
mock_instance.get.return_value.text = '{"status": "ok"}'
|
||||
|
||||
await get_health_summary.ainvoke({"env_id": "ss-dev"})
|
||||
|
||||
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/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")
|
||||
|
||||
with patch("httpx.AsyncClient") as mock_client:
|
||||
mock_instance = AsyncMock()
|
||||
mock_client.return_value.__aenter__.return_value = mock_instance
|
||||
mock_instance.get.return_value.status_code = 200
|
||||
mock_instance.get.return_value.text = '["prod", "dev"]'
|
||||
|
||||
await list_environments.ainvoke({})
|
||||
|
||||
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/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")
|
||||
|
||||
with patch("httpx.AsyncClient") as mock_client:
|
||||
mock_instance = AsyncMock()
|
||||
mock_client.return_value.__aenter__.return_value = mock_instance
|
||||
mock_instance.get.return_value.status_code = 200
|
||||
mock_instance.get.return_value.text = '[{"id":"prod","password":"secret-pass","api_key":"secret-key","nested":{"token":"secret-token"},"name":"ss-prod"}]'
|
||||
|
||||
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")
|
||||
|
||||
with patch("httpx.AsyncClient") as mock_client:
|
||||
mock_instance = AsyncMock()
|
||||
mock_client.return_value.__aenter__.return_value = mock_instance
|
||||
mock_instance.get.return_value.status_code = 200
|
||||
mock_instance.get.return_value.text = '{"status": "running"}'
|
||||
|
||||
await get_task_status.ainvoke({"task_id": "task-123"})
|
||||
|
||||
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/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")
|
||||
|
||||
with patch("httpx.AsyncClient") as mock_client:
|
||||
mock_instance = AsyncMock()
|
||||
mock_client.return_value.__aenter__.return_value = mock_instance
|
||||
mock_instance.post.return_value = Mock(status_code=201, text='{"id": "task-1"}')
|
||||
|
||||
await run_backup.ainvoke({"environment_id": "prod", "dashboard_id": 10})
|
||||
|
||||
call_args = mock_instance.post.call_args
|
||||
assert call_args is not None
|
||||
args, kwargs = call_args
|
||||
assert "api/tasks" in args[0]
|
||||
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")
|
||||
|
||||
with patch("httpx.AsyncClient") as mock_client:
|
||||
mock_instance = AsyncMock()
|
||||
mock_client.return_value.__aenter__.return_value = mock_instance
|
||||
mock_instance.post.return_value = Mock(status_code=200, text='{"status": "success"}')
|
||||
|
||||
await deploy_dashboard.ainvoke({"dashboard_ref": "42", "environment_id": "prod"})
|
||||
|
||||
call_args = mock_instance.post.call_args
|
||||
assert call_args is not None
|
||||
args, kwargs = call_args
|
||||
assert "api/git/repositories/42/deploy" in args[0]
|
||||
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
|
||||
Reference in New Issue
Block a user