gds-agent/mcp_server/tests/test_basic_tools.py at 35c22f4eaf4d5483eed6ecc6400807240063cd3e · brs96/gds-agent · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
import pytest
import json


@pytest.mark.asyncio
async def test_count_nodes(mcp_client):
    result = await mcp_client.call_tool("count_nodes")

    assert len(result) == 1
    result_text = result[0]["text"]
    node_count = int(result_text.strip())
    assert node_count == 302


@pytest.mark.asyncio
async def test_list_tools(mcp_client):
    """Test that all expected tools are listed."""

    # Get list of available tools
    tools = await mcp_client.list_tools()
    tool_names = [tool["name"] for tool in tools]

    # Expected tools (based on server.py imports)
    expected_tools = [
        # Basic tools
        "count_nodes",
        "get_node_properties_keys",
        "get_relationship_properties_keys",
        # Centrality algorithms
        "article_rank",
        "articulation_points",
        "betweenness_centrality",
        "bridges",
        "CELF",
        "closeness_centrality",
        "degree_centrality",
        "eigenvector_centrality",
        "pagerank",
        "harmonic_centrality",
        "HITS",
        # Community algorithms
        "conductance",
        "HDBSCAN",
        "k_core_decomposition",
        "k_1_coloring",
        "k_means_clustering",
        "label_propagation",
        "leiden",
        "local_clustering_coefficient",
        "louvain",
        "modularity_metric",
        "modularity_optimization",
        "strongly_connected_components",
        "triangle_count",
        "weakly_connected_components",
        "approximate_maximum_k_cut",
        "speaker_listener_label_propagation",
        # Path algorithms
        "find_shortest_path",
        "delta_stepping_shortest_path",
        "dijkstra_single_source_shortest_path",
        "a_star_shortest_path",
        "yens_shortest_paths",
        "minimum_weight_spanning_tree",
        "minimum_directed_steiner_tree",
        "prize_collecting_steiner_tree",
        "all_pairs_shortest_paths",
        "random_walk",
        "breadth_first_search",
        "depth_first_search",
        "bellman_ford_single_source_shortest_path",
        "longest_path",
    ]

    # Check that we have the expected tools
    assert len(tool_names) == len(expected_tools)
    for expected_tool in expected_tools:
        assert expected_tool in tool_names, (
            f"Expected tool '{expected_tool}' not found in tool list"
        )


@pytest.mark.asyncio
async def test_get_node_properties_keys(mcp_client):
    result = await mcp_client.call_tool("get_node_properties_keys")

    assert len(result) == 1
    result_text = result[0]["text"]
    properties_keys = json.loads(result_text)

    assert properties_keys == [
        "zone",
        "rail",
        "latitude",
        "name",
        "total_lines",
        "id",
        "display_name",
        "longitude",
    ]


@pytest.mark.asyncio
async def test_get_relationship_properties_keys(mcp_client):
    result = await mcp_client.call_tool("get_relationship_properties_keys")

    assert len(result) == 1
    result_text = result[0]["text"]
    properties_keys = json.loads(result_text)

    assert properties_keys == ["distance", "line", "time"]