mirror of
https://github.com/apache/superset.git
synced 2026-08-04 13:02:41 +00:00
276 lines
9.6 KiB
Python
276 lines
9.6 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# pylint: disable=unused-argument
|
|
"""End-to-end contract tests: mock provider through the real MCP stack.
|
|
|
|
These tests exercise the full request path — REST API → orchestrator →
|
|
deterministic mock provider → in-memory FastMCP client → real MCP tool with
|
|
its middleware — with only the DAO layer and MCP authentication mocked. No
|
|
external AI API is ever called.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Generator, Iterator
|
|
from typing import Any
|
|
from unittest.mock import Mock, patch
|
|
|
|
import pytest
|
|
from enx_dev.ai_chat.approvals import RESOURCE
|
|
from pytest_mock import MockerFixture
|
|
from superset_core.common import models as core_models
|
|
|
|
AI_CHAT_E2E_APP = pytest.mark.parametrize(
|
|
"app",
|
|
[
|
|
{
|
|
"FEATURE_FLAGS": {"ENABLE_EXTENSIONS": True},
|
|
"AI_CHAT_CONFIG": {
|
|
"ENABLED": True,
|
|
"PROVIDER": "mock",
|
|
"ALLOWED_MCP_TOOLS": ["list_dashboards", "delete_dashboard"],
|
|
"REQUIRE_APPROVAL_FOR_MUTATIONS": True,
|
|
},
|
|
# Tool-level RBAC is covered by the MCP service's own suite; the
|
|
# e2e flow here focuses on the gateway contract.
|
|
"MCP_RBAC_ENABLED": False,
|
|
}
|
|
],
|
|
indirect=True,
|
|
)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def no_mcp_dev_username(app: Any) -> Generator[None, None, None]:
|
|
original = app.config.get("MCP_DEV_USERNAME")
|
|
app.config["MCP_DEV_USERNAME"] = None
|
|
yield
|
|
app.config["MCP_DEV_USERNAME"] = original
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def web_user(mocker: MockerFixture) -> Mock:
|
|
# The unit-test client is unauthenticated (authorization is patched by
|
|
# full_api_access); give the gateway a real-shaped session user.
|
|
user = Mock()
|
|
user.id = 1
|
|
user.username = "admin"
|
|
user.roles = []
|
|
g_mock = mocker.patch("enx_dev.ai_chat.api.g")
|
|
g_mock.user = user
|
|
return user
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def mock_mcp_auth() -> Iterator[Mock]:
|
|
with patch("superset.mcp_service.auth.get_user_from_request") as mock_get_user:
|
|
mock_user = Mock()
|
|
mock_user.id = 1
|
|
mock_user.username = "admin"
|
|
mock_get_user.return_value = mock_user
|
|
yield mock_get_user
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def cleanup_approvals(app: Any) -> Generator[None, None, None]:
|
|
yield
|
|
|
|
with app.app_context():
|
|
core_models.get_session().query(core_models.KeyValue).filter(
|
|
core_models.KeyValue.resource == RESOURCE
|
|
).delete()
|
|
core_models.get_session().commit()
|
|
|
|
|
|
# Extension APIs are mounted under /extensions/{publisher}/{name}.
|
|
API_BASE = "/extensions/enx-dev/ai-chat"
|
|
|
|
|
|
def _dashboard_mock() -> Mock:
|
|
dashboard = Mock()
|
|
dashboard.id = 1
|
|
dashboard.dashboard_title = "Test Dashboard"
|
|
dashboard.slug = "test-dashboard"
|
|
dashboard.url = "/dashboard/1"
|
|
dashboard.published = True
|
|
dashboard.changed_by_name = "admin"
|
|
dashboard.changed_on = None
|
|
dashboard.changed_on_humanized = None
|
|
dashboard.created_by_name = "admin"
|
|
dashboard.created_on = None
|
|
dashboard.created_on_humanized = None
|
|
dashboard.tags = []
|
|
dashboard.editors = []
|
|
dashboard.slices = []
|
|
dashboard.description = None
|
|
dashboard.css = None
|
|
dashboard.embedded = []
|
|
dashboard.charts = []
|
|
dashboard.certified_by = None
|
|
dashboard.certification_details = None
|
|
dashboard.deleted_at = None
|
|
dashboard.json_metadata = None
|
|
dashboard.is_managed_externally = False
|
|
dashboard.external_url = None
|
|
dashboard.uuid = "test-dashboard-uuid-1"
|
|
dashboard.thumbnail_url = None
|
|
dashboard._mapping = { # pylint: disable=protected-access
|
|
"id": dashboard.id,
|
|
"dashboard_title": dashboard.dashboard_title,
|
|
"slug": dashboard.slug,
|
|
"url": dashboard.url,
|
|
"published": dashboard.published,
|
|
"changed_by_name": dashboard.changed_by_name,
|
|
"changed_on": dashboard.changed_on,
|
|
"changed_on_humanized": dashboard.changed_on_humanized,
|
|
"created_by_name": dashboard.created_by_name,
|
|
"created_on": dashboard.created_on,
|
|
"created_on_humanized": dashboard.created_on_humanized,
|
|
"tags": dashboard.tags,
|
|
"editors": dashboard.editors,
|
|
"charts": [],
|
|
}
|
|
return dashboard
|
|
|
|
|
|
@AI_CHAT_E2E_APP
|
|
def test_read_only_flow_executes_real_mcp_tool(
|
|
client: Any, full_api_access: None
|
|
) -> None:
|
|
with patch(
|
|
"superset.daos.dashboard.DashboardDAO.list",
|
|
return_value=([_dashboard_mock()], 1),
|
|
):
|
|
response = client.post(
|
|
f"{API_BASE}/chat",
|
|
json={
|
|
"conversation_id": "conv_e2e_read",
|
|
"messages": [{"role": "user", "content": "list my dashboards please"}],
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
events = response.json["result"]["events"]
|
|
types = [event["type"] for event in events]
|
|
assert types == [
|
|
"tool.running",
|
|
"tool.completed",
|
|
"message.completed",
|
|
"request.completed",
|
|
]
|
|
completed = events[1]
|
|
assert completed["tool"] == "list_dashboards"
|
|
assert "Test Dashboard" in completed["result"]
|
|
# The mock provider summarized the real tool output.
|
|
assert "list_dashboards" in events[2]["content"]
|
|
|
|
|
|
@AI_CHAT_E2E_APP
|
|
def test_mutation_flow_requires_and_honors_rejection(
|
|
client: Any, full_api_access: None
|
|
) -> None:
|
|
# Step 1: the mutation is proposed, not executed.
|
|
response = client.post(
|
|
f"{API_BASE}/chat",
|
|
json={
|
|
"conversation_id": "conv_e2e_mut",
|
|
"messages": [{"role": "user", "content": "delete dashboard 42"}],
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
events = response.json["result"]["events"]
|
|
approval_event = events[-1]
|
|
assert approval_event["type"] == "tool.approval_required"
|
|
assert approval_event["tool"] == "delete_dashboard"
|
|
assert approval_event["classification"] == "destructive"
|
|
assert approval_event["arguments"] == {"request": {"identifier": 42}}
|
|
approval_id = approval_event["approval_id"]
|
|
|
|
# Step 2: rejection never executes and burns the approval.
|
|
with patch("superset.daos.dashboard.DashboardDAO.find_by_id") as mock_find:
|
|
response = client.post(
|
|
f"{API_BASE}/tool_approval",
|
|
json={
|
|
"conversation_id": "conv_e2e_mut",
|
|
"messages": [{"role": "user", "content": "delete dashboard 42"}],
|
|
"approval_id": approval_id,
|
|
"decision": "reject",
|
|
"tool_call": {
|
|
"id": approval_event["id"],
|
|
"name": "delete_dashboard",
|
|
"arguments": {"request": {"identifier": 42}},
|
|
},
|
|
},
|
|
)
|
|
mock_find.assert_not_called()
|
|
assert response.status_code == 200
|
|
types = [event["type"] for event in response.json["result"]["events"]]
|
|
assert types[0] == "tool.rejected"
|
|
assert "request.completed" in types
|
|
|
|
# Step 3: the burned approval cannot be replayed as an approval.
|
|
response = client.post(
|
|
f"{API_BASE}/tool_approval",
|
|
json={
|
|
"conversation_id": "conv_e2e_mut",
|
|
"messages": [{"role": "user", "content": "delete dashboard 42"}],
|
|
"approval_id": approval_id,
|
|
"decision": "approve",
|
|
"tool_call": {
|
|
"id": approval_event["id"],
|
|
"name": "delete_dashboard",
|
|
"arguments": {"request": {"identifier": 42}},
|
|
},
|
|
},
|
|
)
|
|
assert response.status_code == 400
|
|
assert response.json["error_code"] == "AI_CHAT_APPROVAL_EXPIRED"
|
|
|
|
|
|
@AI_CHAT_E2E_APP
|
|
def test_mutation_approval_with_tampered_arguments_rejected(
|
|
client: Any, full_api_access: None
|
|
) -> None:
|
|
response = client.post(
|
|
f"{API_BASE}/chat",
|
|
json={
|
|
"conversation_id": "conv_e2e_tamper",
|
|
"messages": [{"role": "user", "content": "delete dashboard 42"}],
|
|
},
|
|
)
|
|
approval_event = response.json["result"]["events"][-1]
|
|
assert approval_event["type"] == "tool.approval_required"
|
|
|
|
# Approving with different arguments must fail and must not execute.
|
|
with patch("superset.daos.dashboard.DashboardDAO.find_by_id") as mock_find:
|
|
response = client.post(
|
|
f"{API_BASE}/tool_approval",
|
|
json={
|
|
"conversation_id": "conv_e2e_tamper",
|
|
"messages": [{"role": "user", "content": "delete dashboard 42"}],
|
|
"approval_id": approval_event["approval_id"],
|
|
"decision": "approve",
|
|
"tool_call": {
|
|
"id": approval_event["id"],
|
|
"name": "delete_dashboard",
|
|
"arguments": {"request": {"identifier": 43}},
|
|
},
|
|
},
|
|
)
|
|
mock_find.assert_not_called()
|
|
assert response.status_code == 400
|
|
assert response.json["error_code"] == "AI_CHAT_APPROVAL_MISMATCH"
|