Files
ss-tools/backend/tests/agent/test_superset_tools.py

562 lines
27 KiB
Python

# #region Test.Agent.SupersetTools [C:3] [TYPE Module] [SEMANTICS test,agent,tools,superset,sql,audit,dashboard,dataset]
# @BRIEF Integration tests for new Superset agent tools: SQL, explore DB, audit, create/copy dashboard, create dataset, format SQL.
# @RELATION BINDS_TO -> [AgentChat.Tools]
# @RELATION BINDS_TO -> [AgentChat.ToolResolver]
# @TEST_EDGE: safe_sql -> superset_execute_sql with SELECT passes
# @TEST_EDGE: dangerous_sql -> superset_execute_sql with DROP returns error
# @TEST_EDGE: external_fail -> HTTP error returns error string
from pathlib import Path
import sys
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "src"))
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import httpx
# ── Helpers ──
def _mock_httpx_response(status_code=200, json_data=None, text=""):
"""Build a mock httpx.Response with given status, JSON, and text."""
resp = MagicMock(spec=httpx.Response)
resp.status_code = status_code
resp.text = text
resp.json = MagicMock(return_value=json_data or {})
return resp
# ═══════════════════════════════════════════════════════════════════
# Tool name registration
# ═══════════════════════════════════════════════════════════════════
class TestToolRegistration:
"""Test that new tools are registered in get_all_tools() and classification sets."""
def test_all_new_tools_registered(self):
"""All 7 new tools appear in get_all_tools()."""
with patch("src.agent.tools.logger", MagicMock()):
from src.agent.tools import get_all_tools
tools = get_all_tools()
tool_names = {t.name for t in tools}
expected = {
"superset_execute_sql",
"superset_explore_database",
"superset_audit_permissions",
"superset_create_dashboard",
"superset_copy_dashboard",
"superset_create_dataset",
"superset_format_sql",
}
assert expected.issubset(tool_names), f"Missing: {expected - tool_names}"
@pytest.mark.skip(reason="_SAFE_AGENT_TOOLS removed during 035 refactoring — needs test rewrite for new tool registry API")
def test_tool_resolver_sets_contain_new_tools(self):
"""Classification sets in _tool_resolver.py contain the new tools."""
with patch("src.agent.tools.logger", MagicMock()):
from src.agent._tool_resolver import (
_SAFE_AGENT_TOOLS,
_GUARDED_AGENT_TOOLS,
_FAST_CONFIRM_TOOLS,
)
assert "superset_execute_sql" in _SAFE_AGENT_TOOLS
assert "superset_explore_database" in _SAFE_AGENT_TOOLS
assert "superset_audit_permissions" in _SAFE_AGENT_TOOLS
assert "superset_format_sql" in _SAFE_AGENT_TOOLS
assert "superset_create_dashboard" in _GUARDED_AGENT_TOOLS
assert "superset_copy_dashboard" in _GUARDED_AGENT_TOOLS
assert "superset_create_dataset" in _GUARDED_AGENT_TOOLS
assert "superset_explore_database" in _FAST_CONFIRM_TOOLS
assert "superset_audit_permissions" in _FAST_CONFIRM_TOOLS
assert "superset_format_sql" in _FAST_CONFIRM_TOOLS
# ═══════════════════════════════════════════════════════════════════
# Intent matching (get_tools_for_query)
# ═══════════════════════════════════════════════════════════════════
class TestIntentMatching:
"""Test get_tools_for_query matches intent keywords for new tools."""
pytestmark = pytest.mark.skip(reason="get_tools_for_query removed during 035 refactoring — needs test rewrite for new tool pipeline API")
def _get_for_query(self, query):
with patch("src.agent.tools.logger", MagicMock()):
from src.agent.tools import get_tools_for_query
return {t.name for t in get_tools_for_query(query)}
def test_sql_intent(self):
tools = self._get_for_query("run a SQL query on users")
assert "superset_execute_sql" in tools
def test_schema_intent(self):
tools = self._get_for_query("list all schemas in the database")
assert "superset_explore_database" in tools
def test_table_intent(self):
tools = self._get_for_query("show me the tables in public schema")
assert "superset_explore_database" in tools
def test_audit_intent(self):
tools = self._get_for_query("audit user permissions access rights")
assert "superset_audit_permissions" in tools
def test_create_dashboard_intent(self):
tools = self._get_for_query("create a new dashboard called Sales")
assert "superset_create_dashboard" in tools
def test_copy_dashboard_intent(self):
tools = self._get_for_query("copy dashboard 42")
assert "superset_copy_dashboard" in tools
def test_create_dataset_intent(self):
tools = self._get_for_query("create a new dataset for orders table")
assert "superset_create_dataset" in tools
def test_format_sql_intent(self):
# The intent matcher looks for exact phrase "format sql" in text
tools = self._get_for_query("please format sql for me")
assert "superset_format_sql" in tools
def test_no_match_returns_defaults(self):
tools = self._get_for_query("hello how are you")
assert "show_capabilities" in tools
assert "search_dashboards" in tools
# ═══════════════════════════════════════════════════════════════════
# superset_execute_sql
# ═══════════════════════════════════════════════════════════════════
class TestSupersetExecuteSql:
"""Test superset_execute_sql agent tool."""
@pytest.fixture(autouse=True)
def setup_mocks(self):
with patch("src.agent.tools.logger", MagicMock()), \
patch("src.agent.tools._dual_auth_headers", return_value={"Authorization": "Bearer test"}), \
patch("src.agent.tools.get_user_jwt", return_value="user-jwt"), \
patch("src.agent.tools.get_service_jwt", return_value="svc-jwt"):
yield
@pytest.mark.asyncio
async def test_safe_sql_execution(self):
"""Safe SELECT returns API result."""
mock_resp = _mock_httpx_response(200, text='{"status": "success", "data": []}')
with patch("httpx.AsyncClient.post", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_execute_sql
result = await superset_execute_sql.ainvoke({
"environment_id": "prod",
"database_id": 1,
"sql": "SELECT 1",
})
assert "success" in result
@pytest.mark.asyncio
async def test_dangerous_sql_returns_error(self):
"""Dangerous SQL returns error from API (guard is server-side)."""
mock_resp = _mock_httpx_response(200, text='{"error": "Dangerous SQL rejected"}')
with patch("httpx.AsyncClient.post", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_execute_sql
result = await superset_execute_sql.ainvoke({
"environment_id": "prod",
"database_id": 1,
"sql": "DROP TABLE users",
})
assert "Dangerous" in result
@pytest.mark.asyncio
async def test_optional_schema(self):
"""Schema parameter is passed when provided."""
mock_resp = _mock_httpx_response(200, text="{}")
with patch("httpx.AsyncClient.post") as mock_post:
mock_post.return_value = mock_resp
from src.agent.tools import superset_execute_sql
await superset_execute_sql.ainvoke({
"environment_id": "prod",
"database_id": 1,
"sql": "SELECT 1",
"query_schema": "public",
})
assert "schema" in mock_post.call_args[1]["params"]
@pytest.mark.asyncio
async def test_http_error(self):
"""HTTP error returns error string."""
mock_resp = _mock_httpx_response(500, text="Internal Server Error")
with patch("httpx.AsyncClient.post", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_execute_sql
result = await superset_execute_sql.ainvoke({
"environment_id": "prod",
"database_id": 1,
"sql": "SELECT 1",
})
assert "Error 500" in result
# ═══════════════════════════════════════════════════════════════════
# superset_explore_database
# ═══════════════════════════════════════════════════════════════════
class TestSupersetExploreDatabase:
"""Test superset_explore_database agent tool."""
@pytest.fixture(autouse=True)
def setup_mocks(self):
with patch("src.agent.tools.logger", MagicMock()), \
patch("src.agent.tools._dual_auth_headers", return_value={"Authorization": "Bearer test"}), \
patch("src.agent.tools.get_user_jwt", return_value="user-jwt"), \
patch("src.agent.tools.get_service_jwt", return_value="svc-jwt"):
yield
@pytest.mark.asyncio
async def test_list_schemas(self):
"""Action 'schemas' GETs /api/agent/superset/databases/{id}/schemas."""
mock_resp = _mock_httpx_response(200, text='["public", "analytics"]')
with patch("httpx.AsyncClient.get", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_explore_database
result = await superset_explore_database.ainvoke({
"environment_id": "prod",
"database_id": 1,
"action": "schemas",
})
assert "public" in result
@pytest.mark.asyncio
async def test_list_tables(self):
"""Action 'tables' GETs /api/agent/superset/databases/{id}/tables."""
mock_resp = _mock_httpx_response(200, text="table1, table2")
with patch("httpx.AsyncClient.get", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_explore_database
result = await superset_explore_database.ainvoke({
"environment_id": "prod",
"database_id": 1,
"action": "tables",
"schema_name": "public",
})
assert "table" in result
@pytest.mark.asyncio
async def test_table_metadata(self):
"""Action 'table_metadata' GETs metadata endpoint."""
mock_resp = _mock_httpx_response(200, text='{"columns": []}')
with patch("httpx.AsyncClient.get", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_explore_database
result = await superset_explore_database.ainvoke({
"environment_id": "prod",
"database_id": 1,
"action": "table_metadata",
"table_name": "users",
})
assert "columns" in result
@pytest.mark.asyncio
async def test_select_star(self):
"""Action 'select_star' GETs select_star endpoint."""
mock_resp = _mock_httpx_response(200, text='"SELECT * FROM users"')
with patch("httpx.AsyncClient.get", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_explore_database
result = await superset_explore_database.ainvoke({
"environment_id": "prod",
"database_id": 1,
"action": "select_star",
"table_name": "users",
})
assert "SELECT" in result
@pytest.mark.asyncio
async def test_tables_missing_schema_returns_error(self):
"""Missing schema_name for 'tables' returns error without HTTP call."""
with patch("httpx.AsyncClient.get") as mock_get:
from src.agent.tools import superset_explore_database
result = await superset_explore_database.ainvoke({
"environment_id": "prod",
"database_id": 1,
"action": "tables",
})
assert "Error" in result
mock_get.assert_not_called()
@pytest.mark.asyncio
async def test_table_metadata_missing_table_name_returns_error(self):
"""Missing table_name for 'table_metadata' returns error without HTTP call."""
with patch("httpx.AsyncClient.get") as mock_get:
from src.agent.tools import superset_explore_database
result = await superset_explore_database.ainvoke({
"environment_id": "prod",
"database_id": 1,
"action": "table_metadata",
})
assert "Error" in result
mock_get.assert_not_called()
@pytest.mark.asyncio
async def test_unknown_action_returns_error(self):
"""Unknown action returns error without HTTP call."""
with patch("httpx.AsyncClient.get") as mock_get:
from src.agent.tools import superset_explore_database
result = await superset_explore_database.ainvoke({
"environment_id": "prod",
"database_id": 1,
"action": "invalid_action",
})
assert "unknown action" in result.lower()
mock_get.assert_not_called()
# ═══════════════════════════════════════════════════════════════════
# superset_audit_permissions
# ═══════════════════════════════════════════════════════════════════
class TestSupersetAuditPermissions:
"""Test superset_audit_permissions agent tool."""
@pytest.fixture(autouse=True)
def setup_mocks(self):
with patch("src.agent.tools.logger", MagicMock()), \
patch("src.agent.tools._dual_auth_headers", return_value={"Authorization": "Bearer test"}), \
patch("src.agent.tools.get_user_jwt", return_value="user-jwt"), \
patch("src.agent.tools.get_service_jwt", return_value="svc-jwt"):
yield
@pytest.mark.asyncio
async def test_basic_audit(self):
"""Permissions audit GETs /api/agent/superset/audit/permissions."""
mock_resp = _mock_httpx_response(200, text='{"users": [], "total_users": 0}')
with patch("httpx.AsyncClient.get", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_audit_permissions
result = await superset_audit_permissions.ainvoke({
"environment_id": "prod",
})
assert "total_users" in result
@pytest.mark.asyncio
async def test_audit_with_filters(self):
"""Username filter and include_admin are passed as query params."""
mock_resp = _mock_httpx_response(200, text="{}")
with patch("httpx.AsyncClient.get") as mock_get:
mock_get.return_value = mock_resp
from src.agent.tools import superset_audit_permissions
await superset_audit_permissions.ainvoke({
"environment_id": "prod",
"username_filter": "alice",
"include_admin": True,
})
call_params = mock_get.call_args[1]["params"]
assert call_params["username_filter"] == "alice"
assert call_params["include_admin"] == "true"
@pytest.mark.asyncio
async def test_audit_without_filters(self):
"""Without filters, no additional params beyond environment_id."""
mock_resp = _mock_httpx_response(200, text="{}")
with patch("httpx.AsyncClient.get") as mock_get:
mock_get.return_value = mock_resp
from src.agent.tools import superset_audit_permissions
await superset_audit_permissions.ainvoke({
"environment_id": "prod",
})
call_params = mock_get.call_args[1]["params"]
assert "username_filter" not in call_params
assert "include_admin" not in call_params
# ═══════════════════════════════════════════════════════════════════
# superset_create_dashboard
# ═══════════════════════════════════════════════════════════════════
class TestSupersetCreateDashboard:
"""Test superset_create_dashboard agent tool."""
@pytest.fixture(autouse=True)
def setup_mocks(self):
with patch("src.agent.tools.logger", MagicMock()), \
patch("src.agent.tools._dual_auth_headers", return_value={"Authorization": "Bearer test"}), \
patch("src.agent.tools.get_user_jwt", return_value="user-jwt"), \
patch("src.agent.tools.get_service_jwt", return_value="svc-jwt"):
yield
@pytest.mark.asyncio
async def test_create_dashboard(self):
"""Creates dashboard via POST to /api/agent/superset/dashboards."""
mock_resp = _mock_httpx_response(201, text='{"id": 1, "dashboard_title": "Sales"}')
with patch("httpx.AsyncClient.post", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_create_dashboard
result = await superset_create_dashboard.ainvoke({
"environment_id": "prod",
"dashboard_title": "Sales",
"slug": "sales",
"published": True,
})
assert "Sales" in result
@pytest.mark.asyncio
async def test_create_dashboard_minimal(self):
"""Minimal create with just title."""
mock_resp = _mock_httpx_response(201, text='{"id": 2}')
with patch("httpx.AsyncClient.post") as mock_post:
mock_post.return_value = mock_resp
from src.agent.tools import superset_create_dashboard
await superset_create_dashboard.ainvoke({
"environment_id": "prod",
"dashboard_title": "Minimal",
})
call_params = mock_post.call_args[1]["params"]
assert call_params["dashboard_title"] == "Minimal"
assert call_params["published"] == "false"
# ═══════════════════════════════════════════════════════════════════
# superset_copy_dashboard
# ═══════════════════════════════════════════════════════════════════
class TestSupersetCopyDashboard:
"""Test superset_copy_dashboard agent tool."""
@pytest.fixture(autouse=True)
def setup_mocks(self):
with patch("src.agent.tools.logger", MagicMock()), \
patch("src.agent.tools._dual_auth_headers", return_value={"Authorization": "Bearer test"}), \
patch("src.agent.tools.get_user_jwt", return_value="user-jwt"), \
patch("src.agent.tools.get_service_jwt", return_value="svc-jwt"):
yield
@pytest.mark.asyncio
async def test_copy_dashboard(self):
"""Copies dashboard via POST to /api/agent/superset/dashboards/{id}/copy."""
mock_resp = _mock_httpx_response(201, text='{"id": 2, "result": "copied"}')
with patch("httpx.AsyncClient.post", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_copy_dashboard
result = await superset_copy_dashboard.ainvoke({
"environment_id": "prod",
"dashboard_id": 1,
"dashboard_title": "Copy of Sales",
})
assert "copied" in result
@pytest.mark.asyncio
async def test_copy_dashboard_without_title(self):
"""Copy without optional title."""
mock_resp = _mock_httpx_response(201, text="{}")
with patch("httpx.AsyncClient.post") as mock_post:
mock_post.return_value = mock_resp
from src.agent.tools import superset_copy_dashboard
await superset_copy_dashboard.ainvoke({
"environment_id": "prod",
"dashboard_id": 1,
})
call_params = mock_post.call_args[1]["params"]
assert "dashboard_title" not in call_params
# ═══════════════════════════════════════════════════════════════════
# superset_create_dataset
# ═══════════════════════════════════════════════════════════════════
class TestSupersetCreateDataset:
"""Test superset_create_dataset agent tool."""
@pytest.fixture(autouse=True)
def setup_mocks(self):
with patch("src.agent.tools.logger", MagicMock()), \
patch("src.agent.tools._dual_auth_headers", return_value={"Authorization": "Bearer test"}), \
patch("src.agent.tools.get_user_jwt", return_value="user-jwt"), \
patch("src.agent.tools.get_service_jwt", return_value="svc-jwt"):
yield
@pytest.mark.asyncio
async def test_create_dataset(self):
"""Creates dataset via POST to /api/agent/superset/datasets."""
mock_resp = _mock_httpx_response(201, text='{"id": 5, "table_name": "orders"}')
with patch("httpx.AsyncClient.post", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_create_dataset
result = await superset_create_dataset.ainvoke({
"environment_id": "prod",
"table_name": "orders",
"database": 1,
"schema_name": "public",
})
assert "orders" in result
@pytest.mark.asyncio
async def test_create_dataset_minimal(self):
"""Minimal create without schema."""
mock_resp = _mock_httpx_response(201, text='{"id": 6}')
with patch("httpx.AsyncClient.post") as mock_post:
mock_post.return_value = mock_resp
from src.agent.tools import superset_create_dataset
await superset_create_dataset.ainvoke({
"environment_id": "prod",
"table_name": "users",
"database": 1,
})
call_params = mock_post.call_args[1]["params"]
assert "schema_name" not in call_params
# ═══════════════════════════════════════════════════════════════════
# superset_format_sql
# ═══════════════════════════════════════════════════════════════════
class TestSupersetFormatSql:
"""Test superset_format_sql agent tool."""
@pytest.fixture(autouse=True)
def setup_mocks(self):
with patch("src.agent.tools.logger", MagicMock()), \
patch("src.agent.tools._dual_auth_headers", return_value={"Authorization": "Bearer test"}), \
patch("src.agent.tools.get_user_jwt", return_value="user-jwt"), \
patch("src.agent.tools.get_service_jwt", return_value="svc-jwt"):
yield
@pytest.mark.asyncio
async def test_format_sql(self):
"""Formats SQL via POST to /api/agent/superset/sqllab/format."""
mock_resp = _mock_httpx_response(200, text="SELECT\n 1\nFROM\n users\n")
with patch("httpx.AsyncClient.post", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_format_sql
result = await superset_format_sql.ainvoke({
"environment_id": "prod",
"sql": "SELECT 1 FROM users",
})
assert "SELECT" in result
@pytest.mark.asyncio
async def test_format_sql_error(self):
"""HTTP error returns error string."""
mock_resp = _mock_httpx_response(500, text="Internal Error")
with patch("httpx.AsyncClient.post", AsyncMock(return_value=mock_resp)):
from src.agent.tools import superset_format_sql
result = await superset_format_sql.ainvoke({
"environment_id": "prod",
"sql": "INVALID",
})
assert "Error 500" in result
# ═══════════════════════════════════════════════════════════════════
# Tool Resolver inference
# ═══════════════════════════════════════════════════════════════════
class TestToolResolverInference:
"""Test _tool_resolver.py inference with new Superset tool patterns."""
pytestmark = pytest.mark.skip(reason="infer_tool_from_text removed during 035 refactoring — needs test rewrite for new resolver API")
def test_infer_from_text_sql(self):
"""SQL-related query infers superset_execute_sql."""
with patch("src.agent.tools.logger", MagicMock()):
from src.agent._tool_resolver import infer_tool_from_text
# The resolver doesn't currently have SQL inference — verify fallback
# but ensure it's extended if needed in the future
result = infer_tool_from_text("sql select query users")
# No specific SQL inference yet — returns None or base tool
# This test documents current behavior; update if resolver is extended
assert result is None or isinstance(result, str)
# #endregion Test.Agent.SupersetTools