97 lines
3.5 KiB
Python
97 lines
3.5 KiB
Python
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from django.test import SimpleTestCase, override_settings
|
|
|
|
from mcp_server.tools.datasets import execute_dataset, search_datasets
|
|
|
|
|
|
class DatasetToolTests(SimpleTestCase):
|
|
@patch("mcp_server.tools.datasets.authenticated_factory_user")
|
|
@patch("mcp_server.tools.datasets.Dataset.objects.filter")
|
|
def test_search_returns_safe_catalog_fields(self, filter_mock, _user_mock):
|
|
queryset = MagicMock()
|
|
filter_mock.return_value = queryset
|
|
queryset.order_by.return_value.values.return_value.__getitem__.return_value = [
|
|
{
|
|
"code": "daily_output",
|
|
"name": "日产量",
|
|
"description": "按日统计产量",
|
|
"default_param": {"day": "2026-08-10"},
|
|
"test_param": {},
|
|
}
|
|
]
|
|
|
|
result = search_datasets(limit=10)
|
|
|
|
self.assertEqual(result["items"][0]["code"], "daily_output")
|
|
self.assertNotIn("sql_query", result["items"][0])
|
|
|
|
@patch("mcp_server.tools.datasets.cache")
|
|
@patch("mcp_server.tools.datasets.execute_rendered_dataset")
|
|
@patch("mcp_server.tools.datasets.render_dataset_sql")
|
|
@patch("mcp_server.tools.datasets.require_permission")
|
|
@patch("mcp_server.tools.datasets.authenticated_factory_user")
|
|
@patch("mcp_server.tools.datasets.Dataset.objects.get")
|
|
def test_execute_reuses_identity_permission_and_service(
|
|
self,
|
|
get_mock,
|
|
user_mock,
|
|
permission_mock,
|
|
render_mock,
|
|
execute_mock,
|
|
cache_mock,
|
|
):
|
|
item = SimpleNamespace(
|
|
code="daily_output",
|
|
name="日产量",
|
|
description="按日统计产量",
|
|
cache_seconds=10,
|
|
)
|
|
user = SimpleNamespace(id=42, belong_dept_id=7)
|
|
get_mock.return_value = item
|
|
user_mock.return_value = user
|
|
render_mock.return_value = "select 1"
|
|
cache_mock.get.return_value = None
|
|
execute_mock.return_value = (
|
|
{"data": {"ds0": [{"count": 1}]}, "data2": {}},
|
|
True,
|
|
)
|
|
|
|
result = execute_dataset("daily_output", {"day": "2026-08-10"})
|
|
|
|
permission_mock.assert_called_once_with(user, "dataset.exec")
|
|
render_query = render_mock.call_args.args[1]
|
|
self.assertEqual(render_query["r_user"], 42)
|
|
self.assertEqual(render_query["r_dept"], 7)
|
|
self.assertEqual(result["data"]["ds0"][0]["count"], 1)
|
|
self.assertNotIn("sql_query", result)
|
|
|
|
@override_settings(MCP_MAX_RESULT_BYTES=1)
|
|
@patch("mcp_server.tools.datasets.cache")
|
|
@patch("mcp_server.tools.datasets.execute_rendered_dataset")
|
|
@patch("mcp_server.tools.datasets.render_dataset_sql", return_value="")
|
|
@patch("mcp_server.tools.datasets.require_permission")
|
|
@patch("mcp_server.tools.datasets.authenticated_factory_user")
|
|
@patch("mcp_server.tools.datasets.Dataset.objects.get")
|
|
def test_execute_rejects_oversized_results(
|
|
self,
|
|
get_mock,
|
|
user_mock,
|
|
_permission_mock,
|
|
_render_mock,
|
|
execute_mock,
|
|
_cache_mock,
|
|
):
|
|
get_mock.return_value = SimpleNamespace(
|
|
code="daily_output",
|
|
name="日产量",
|
|
description="",
|
|
cache_seconds=0,
|
|
)
|
|
user_mock.return_value = SimpleNamespace(id=42, belong_dept_id=None)
|
|
execute_mock.return_value = ({"data": {"ds0": [1]}, "data2": {}}, True)
|
|
|
|
with self.assertRaisesRegex(RuntimeError, "超过 MCP 响应上限"):
|
|
execute_dataset("daily_output")
|