tests/test_tool_surface.py

# Copyright 2026 Dataiku SAS
#
# Licensed 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.

"""Contract tests for the fixed MCP tool surface."""

import asyncio
import json
from pathlib import Path
import re
import runpy
from collections import defaultdict

import pytest

import dataiku_mcp


EXPECTED_TOOLS_BY_MODULE = {
    "agent_reviews": frozenset(
        {
            "get_agent_review",
            "get_agent_review_run_results",
            "list_agent_review_runs",
            "list_agent_review_tests",
            "list_agent_reviews",
        }
    ),
    "agents": frozenset(
        {
            "get_agent_settings",
            "get_agent_tool_settings",
            "list_agent_tools",
            "list_agent_versions",
            "list_agents",
        }
    ),
    "cobuild": frozenset(
        {
            "answer_cobuild_confirmation",
            "answer_cobuild_question",
            "get_cobuild_turn_status",
            "list_cobuild_conversations",
            "send_cobuild_message",
            "start_cobuild_conversation",
        }
    ),
    "code_environments": frozenset(
        {
            "create_code_env",
            "delete_code_env",
            "list_code_envs",
            "update_code_env",
        }
    ),
    "connections": frozenset(
        {"get_connection_info", "list_connections", "test_connection"}
    ),
    "container_exec": frozenset({"set_container_exec_config"}),
    "cross_project_sharing": frozenset({"list_shared_objects"}),
    "dashboards": frozenset({"get_dashboard_settings", "list_dashboards"}),
    "data_collections": frozenset(
        {"list_data_collection_objects", "list_data_collections"}
    ),
    "data_quality": frozenset(
        {
            "get_data_quality_rule",
            "get_data_quality_rule_history",
            "get_data_quality_rule_results",
            "get_data_quality_status",
            "list_data_quality_rules",
        }
    ),
    "datasets": frozenset(
        {
            "create_upload_dataset",
            "export_dataset",
            "get_dataset_column_descriptions",
            "get_dataset_info",
            "get_dataset_metrics",
            "get_dataset_profile",
            "get_dataset_sample",
            "list_datasets",
        }
    ),
    "evaluation_stores": frozenset(
        {"get_evaluation_store_details", "list_evaluation_stores"}
    ),
    "flow": frozenset(
        {
            "get_flow_graph",
            "get_flow_object_metadata",
            "list_flow_zones",
        }
    ),
    "general_settings": frozenset(
        {"list_container_exec_configs", "list_spark_configs"}
    ),
    "groups": frozenset(
        {"create_group", "delete_group", "list_groups", "update_group"}
    ),
    "insights": frozenset({"get_insight_settings", "list_insights"}),
    "instances": frozenset(
        {
            "configure_instance",
            "delete_instance",
            "get_current_instance",
            "list_instances",
            "switch_instance",
        }
    ),
    "jobs": frozenset(
        {
            "abort_job",
            "build_datasets",
            "get_future_status",
            "get_job_log",
            "get_job_status",
            "list_jobs",
            "run_recipe",
            "wait_for_job",
        }
    ),
    "licensing": frozenset({"get_licensing_status"}),
    "llms_and_knowledge_banks": frozenset(
        {
            "get_knowledge_bank_settings",
            "get_llm_info",
            "get_rag_llm_settings",
            "list_knowledge_banks",
            "list_llms",
            "list_rag_llms",
            "search_knowledge_bank",
        }
    ),
    "machine_learning.analyses": frozenset(
        {
            "get_ml_analysis_settings",
            "get_ml_analysis_summary",
            "get_ml_model_details",
            "list_ml_analyses",
            "list_ml_analysis_models",
        }
    ),
    "machine_learning.saved_models": frozenset(
        {
            "get_saved_model_version_details",
            "list_saved_model_versions",
            "list_saved_models",
        }
    ),
    "managed_folders": frozenset(
        {
            "create_managed_folder",
            "get_managed_folder_contents",
            "get_managed_folder_info",
            "list_managed_folders",
            "upload_file_to_managed_folder",
        }
    ),
    "plugins": frozenset(
        {
            "install_plugin",
            "ensure_plugin_code_env",
            "delete_plugin",
            "list_plugins",
            "list_plugin_usages",
            "update_plugin",
        }
    ),
    "project_folders": frozenset(
        {
            "create_project_folder",
            "delete_project_folder",
            "get_project_folder",
            "list_project_folders",
            "move_project_to_folder",
        }
    ),
    "project_libraries": frozenset(
        {
            "list_project_library",
            "read_project_library_file",
            "search_project_library",
            "validate_project_library_file",
            "write_project_library_file",
        }
    ),
    "projects": frozenset(
        {
            "count_projects",
            "create_project",
            "get_project_metadata",
            "get_project_settings",
            "get_project_variables",
            "list_projects",
            "set_project_variables",
            "update_project_settings",
        }
    ),
    "recipes": frozenset({"get_recipe_settings", "list_recipes"}),
    "scenarios": frozenset(
        {
            "get_scenario_run_history",
            "get_scenario_settings",
            "list_messaging_channels",
            "list_scenarios",
            "run_scenario",
        }
    ),
    "semantic_models": frozenset(
        {"get_semantic_model_version_config", "list_semantic_models"}
    ),
    "users": frozenset({"create_user", "delete_user", "list_users", "update_user"}),
    "webapps": frozenset({"get_webapp_settings", "get_webapp_state", "list_webapps"}),
    "wikis": frozenset({"get_wiki_article", "list_wiki_articles"}),
}


@pytest.fixture(scope="module")
def registered_tools_by_module() -> dict[str, frozenset[str]]:
    tools_by_module = defaultdict(set)
    for tool in asyncio.run(dataiku_mcp.mcp.list_tools()):
        module = tool.fn.__module__.removeprefix("dataiku_mcp.tools.")
        tools_by_module[module].add(tool.name)
    return {
        module: frozenset(tool_names) for module, tool_names in tools_by_module.items()
    }


def test_registered_tool_modules_are_fixed(registered_tools_by_module):
    assert registered_tools_by_module.keys() == EXPECTED_TOOLS_BY_MODULE.keys()


@pytest.mark.parametrize("distribution", ["stdio", "http"])
def test_registered_tool_names_fit_bedrock_limit(
    registered_tools_by_module, tmp_path, distribution
):
    root = Path(__file__).resolve().parents[1]
    if distribution == "http":
        renderer = runpy.run_path(str(root / "scripts/build_http_plugin.py"))
        root = renderer["render"](
            tmp_path / "http-plugin",
            endpoint="https://mcp.customer.example/mcp",
            archive=False,
        )
    manifest = json.loads((root / ".claude-plugin/plugin.json").read_text())
    codex = json.loads((root / ".codex-plugin/plugin.json").read_text())
    bundled_mcp = json.loads((root / codex["mcpServers"]).read_text())
    servers = bundled_mcp["mcpServers"] | manifest["mcpServers"]
    full_names = [
        re.sub(
            r"[^a-zA-Z0-9_-]",
            "_",
            f"mcp__plugin_{manifest['name']}_{server}__{name}",
        )
        for server in servers
        for names in registered_tools_by_module.values()
        for name in sorted(names)
    ]
    overlong_names = {name: len(name) for name in full_names if len(name) > 64}
    assert not overlong_names, (
        f"Claude Code tool names exceed Bedrock's 64-character limit: {overlong_names}"
    )


@pytest.mark.parametrize("module", EXPECTED_TOOLS_BY_MODULE)
def test_registered_tools_are_fixed_for_module(module, registered_tools_by_module):
    assert registered_tools_by_module[module] == EXPECTED_TOOLS_BY_MODULE[module]