test_tool_policies.py•2.35 kB
from unittest.mock import patch
from dbt_mcp.config.config import load_config
from dbt_mcp.dbt_cli.binary_type import BinaryType
from dbt_mcp.lsp.lsp_binary_manager import LspBinaryInfo
from dbt_mcp.mcp.server import create_dbt_mcp
from dbt_mcp.tools.policy import tool_policies
from dbt_mcp.tools.tool_names import ToolName
from dbt_mcp.tools.toolsets import proxied_tools, Toolset, toolsets
async def test_tool_policies_match_server_tools(env_setup):
"""Test that the ToolPolicy enum matches the tools registered in the server."""
with (
env_setup(),
patch(
"dbt_mcp.config.config.detect_binary_type", return_value=BinaryType.DBT_CORE
),
patch(
"dbt_mcp.config.config.dbt_lsp_binary_info",
return_value=LspBinaryInfo(path="/path/to/lsp", version="1.0.0"),
),
):
config = load_config()
dbt_mcp = await create_dbt_mcp(config)
# Get all tools from the server
server_tools = await dbt_mcp.list_tools()
# Manually adding SQL tools here because the server doesn't get them
# in this unit test.
server_tool_names = (
{tool.name for tool in server_tools}
| {p.value for p in proxied_tools}
# Also adding codegen because it's default off
| {t.value for t in toolsets[Toolset.DBT_CODEGEN]}
)
policy_names = {policy_name for policy_name in tool_policies}
if server_tool_names != policy_names:
raise ValueError(
f"Tool name mismatch:\n"
f"In server but not in enum: {server_tool_names - policy_names}\n"
f"In enum but not in server: {policy_names - server_tool_names}"
)
def test_tool_policies_match_tool_names():
policy_names = {policy.upper() for policy in tool_policies}
tool_names = {tool.name for tool in ToolName}
if tool_names != policy_names:
raise ValueError(
f"Tool name mismatch:\n"
f"In tool names but not in policy: {tool_names - policy_names}\n"
f"In policy but not in tool names: {policy_names - tool_names}"
)
def test_tool_policies_no_duplicates():
"""Test that there are no duplicate tool names in the policy."""
assert len(tool_policies) == len(set(tool_policies.keys()))