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
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_retrieval_augmented_llm_settings",
            "list_knowledge_banks",
            "list_llms",
            "list_retrieval_augmented_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(
        {
            "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_settings", "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("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]