417 lines
16 KiB
Python
417 lines
16 KiB
Python
# #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["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"
|
|
|
|
|
|
@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 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 = 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 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 = 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 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()
|
|
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 src.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
|
|
|
|
|
|
def test_get_tools_for_query_keeps_prefetched_dashboard_prompt_small():
|
|
"""Dashboard requests with prefetched data should not send all tool schemas."""
|
|
from src.agent.tools import get_tools_for_query
|
|
|
|
tools = get_tools_for_query("Покажи доступные дашборды", prefetch_available=True)
|
|
assert [tool.name for tool in tools] == ["show_capabilities"]
|
|
|
|
|
|
def test_get_tools_for_query_selects_write_tool_by_intent():
|
|
"""Write intents should expose only the matching operational tool plus capabilities."""
|
|
from src.agent.tools import get_tools_for_query
|
|
|
|
tools = get_tools_for_query("Запусти миграцию", prefetch_available=False)
|
|
assert [tool.name for tool in tools] == ["show_capabilities", "execute_migration"]
|
|
# #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 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 = 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 src.agent.context import set_service_jwt, set_user_jwt
|
|
from src.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 src.agent.context import set_service_jwt, set_user_jwt
|
|
from src.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 src.agent.context import set_service_jwt, set_user_jwt
|
|
from src.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 src.agent.context import set_service_jwt, set_user_jwt
|
|
from src.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 src.agent.context import set_service_jwt, set_user_jwt
|
|
from src.agent.tools import run_backup
|
|
|
|
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.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 src.agent.context import set_service_jwt, set_user_jwt
|
|
from src.agent.tools import deploy_dashboard
|
|
|
|
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.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 src.agent.context import set_service_jwt, set_user_jwt
|
|
from src.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 src.agent.context import set_service_jwt, set_user_jwt
|
|
from src.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(monkeypatch):
|
|
"""_dual_auth_headers returns empty dict when no JWTs."""
|
|
monkeypatch.delenv("SERVICE_JWT", raising=False)
|
|
from src.agent.context import set_service_jwt, set_user_jwt
|
|
from src.agent.tools import _dual_auth_headers
|
|
|
|
set_service_jwt("")
|
|
set_user_jwt("")
|
|
|
|
headers = _dual_auth_headers()
|
|
assert headers == {}
|
|
# #endregion TestAgentChat.Tools.DualAuthHeaders
|
|
# #endregion TestAgentChat.Tools
|