Skip to content

Commit 279353e

Browse files
authored
Merge pull request #203 from keboola/AI-1172-global-search-tests
Ai 1172 Add Global Search Tool tests
2 parents 168fec3 + e0be650 commit 279353e

File tree

8 files changed

+532
-44
lines changed

8 files changed

+532
-44
lines changed

integtests/conftest.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,7 @@ def _create_tables(storage_client: SyncStorageClient) -> list[TableDef]:
174174
ConfigDef(
175175
component_id='ex-generic-v2',
176176
configuration_id=None,
177-
internal_id='config1',
177+
internal_id='test_config1',
178178
),
179179
]
180180

File renamed without changes.

integtests/tools/components/test_tools.py

Lines changed: 1 addition & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from mcp.server.fastmcp import Context
66

77
from integtests.conftest import ConfigDef, ProjectDef
8-
from keboola_mcp_server.client import KeboolaClient, SuggestedComponent
8+
from keboola_mcp_server.client import KeboolaClient
99
from keboola_mcp_server.config import MetadataField
1010
from keboola_mcp_server.links import Link
1111
from keboola_mcp_server.tools.components.model import (
@@ -21,7 +21,6 @@
2121
add_config_row,
2222
create_config,
2323
create_sql_transformation,
24-
find_component_id,
2524
get_component,
2625
get_config,
2726
get_config_examples,
@@ -972,22 +971,6 @@ async def test_get_config_examples(mcp_context: Context, configs: list[ConfigDef
972971
assert 'parameters' in result
973972

974973

975-
@pytest.mark.asyncio
976-
async def test_find_component_id(mcp_context: Context):
977-
"""Tests that `find_component_id` returns relevant component IDs for a query."""
978-
query = 'generic extractor'
979-
generic_extractor_id = 'ex-generic-v2'
980-
981-
result = await find_component_id(query=query, ctx=mcp_context)
982-
983-
assert isinstance(result, list)
984-
assert len(result) > 0
985-
assert generic_extractor_id in [component.component_id for component in result]
986-
987-
for component in result:
988-
assert isinstance(component, SuggestedComponent)
989-
990-
991974
@pytest.mark.asyncio
992975
async def test_get_config_examples_with_invalid_component(mcp_context: Context):
993976
"""Tests that `get_config_examples` handles non-existent components properly."""

integtests/tools/test_search.py

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
import logging
2+
3+
import pytest
4+
from fastmcp import Context
5+
6+
from integtests.conftest import BucketDef, ConfigDef, TableDef
7+
from keboola_mcp_server.client import KeboolaClient, SuggestedComponent
8+
from keboola_mcp_server.tools.search import GlobalSearchOutput, find_component_id, find_ids_by_name
9+
10+
LOG = logging.getLogger(__name__)
11+
12+
13+
@pytest.mark.asyncio
14+
async def test_global_search_end_to_end(
15+
keboola_client: KeboolaClient,
16+
mcp_context: Context,
17+
buckets: list[BucketDef],
18+
tables: list[TableDef],
19+
configs: list[ConfigDef],
20+
) -> None:
21+
"""
22+
Test the global_search tool end-to-end by searching for items that exist in the test project.
23+
This verifies that the search returns expected results for buckets, tables, and configurations.
24+
"""
25+
26+
# skip this test if the global search is not available
27+
if not await keboola_client.storage_client.is_enabled('global-search'):
28+
LOG.warning('Global search is not available. Please enable it in the project settings.')
29+
pytest.skip('Global search is not available. Please enable it in the project settings.')
30+
31+
# Search for test items by name prefix 'test' which should match our test data
32+
result = await find_ids_by_name(
33+
ctx=mcp_context, name_prefixes=['test'], item_types=tuple(), limit=50, offset=0 # Search all types
34+
)
35+
36+
# Verify the result structure
37+
assert isinstance(result, GlobalSearchOutput)
38+
assert isinstance(result.counts, dict)
39+
assert isinstance(result.groups, dict)
40+
assert 'total' in result.counts
41+
42+
# Verify we found some results
43+
assert result.counts['total'] > 0, 'Should find at least some test items'
44+
45+
# Create sets of expected IDs for verification
46+
expected_bucket_ids = {bucket.bucket_id for bucket in buckets}
47+
expected_table_ids = {table.table_id for table in tables}
48+
expected_config_ids = {config.configuration_id for config in configs if config.configuration_id}
49+
50+
# Check that we can find test buckets
51+
bucket_groups = [group for group in result.groups.values() if group.type == 'bucket']
52+
assert len(bucket_groups) == 1
53+
bucket_group = bucket_groups[0]
54+
found_bucket_ids = {item.id for item in bucket_group.items}
55+
# At least some test buckets should be found
56+
assert found_bucket_ids.intersection(expected_bucket_ids), 'Should find at least one test bucket'
57+
58+
# Check that we can find test tables
59+
table_groups = [group for group in result.groups.values() if group.type == 'table']
60+
assert len(table_groups) == 1
61+
table_group = table_groups[0]
62+
found_table_ids = {item.id for item in table_group.items}
63+
# At least some test tables should be found
64+
assert found_table_ids.intersection(expected_table_ids), 'Should find at least one test table'
65+
66+
# Check that we can find test configurations
67+
config_groups = [group for group in result.groups.values() if group.type == 'configuration']
68+
assert len(config_groups) == 1
69+
config_group = config_groups[0]
70+
found_config_ids = {item.id for item in config_group.items}
71+
# At least some test configurations should be found
72+
assert found_config_ids.intersection(expected_config_ids), 'Should find at least one test configuration'
73+
74+
config_groups = [group for group in result.groups.values() if group.type == 'configuration']
75+
assert len(config_groups) == 1
76+
config_group = config_groups[0]
77+
found_config_ids = {item.id for item in config_group.items}
78+
# At least some test configurations should be found
79+
assert found_config_ids.intersection(expected_config_ids), 'Should find at least one test configuration'
80+
81+
82+
@pytest.mark.asyncio
83+
async def test_find_component_id(mcp_context: Context):
84+
"""Tests that `find_component_id` returns relevant component IDs for a query."""
85+
query = 'generic extractor'
86+
generic_extractor_id = 'ex-generic-v2'
87+
88+
result = await find_component_id(query=query, ctx=mcp_context)
89+
90+
assert isinstance(result, list)
91+
assert len(result) > 0
92+
assert all(isinstance(component, SuggestedComponent) for component in result)
93+
assert generic_extractor_id in [component.component_id for component in result]

pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
44

55
[project]
66
name = "keboola-mcp-server"
7-
version = "1.10.0"
7+
version = "1.10.2"
88
description = "MCP server for interacting with Keboola Connection"
99
readme = "README.md"
1010
requires-python = ">=3.10"

src/keboola_mcp_server/tools/components/tools.py

Lines changed: 1 addition & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from httpx import HTTPStatusError
99
from pydantic import Field
1010

11-
from keboola_mcp_server.client import JsonDict, KeboolaClient, SuggestedComponent
11+
from keboola_mcp_server.client import JsonDict, KeboolaClient
1212
from keboola_mcp_server.config import MetadataField
1313
from keboola_mcp_server.errors import tool_errors
1414
from keboola_mcp_server.links import ProjectLinksManager
@@ -56,7 +56,6 @@ def add_component_tools(mcp: KeboolaMcpServer) -> None:
5656
mcp.add_tool(FunctionTool.from_function(update_config))
5757
mcp.add_tool(FunctionTool.from_function(update_config_row))
5858
mcp.add_tool(FunctionTool.from_function(get_config_examples))
59-
mcp.add_tool(FunctionTool.from_function(find_component_id))
6059
mcp.add_tool(FunctionTool.from_function(create_sql_transformation))
6160
mcp.add_tool(FunctionTool.from_function(update_sql_transformation))
6261
mcp.add_tool(FunctionTool.from_function(list_transformations))
@@ -1022,24 +1021,3 @@ async def get_config_examples(
10221021
markdown += f'{i}. Row Configuration:\n```json\n{json.dumps(example, indent=2)}\n```\n\n'
10231022

10241023
return markdown
1025-
1026-
1027-
@tool_errors()
1028-
@with_session_state()
1029-
async def find_component_id(
1030-
ctx: Context,
1031-
query: Annotated[str, Field(description='Natural language query to find the requested component.')],
1032-
) -> list[SuggestedComponent]:
1033-
"""
1034-
Returns list of component IDs that match the given query.
1035-
1036-
USAGE:
1037-
- Use when you want to find the component for a specific purpose.
1038-
1039-
EXAMPLES:
1040-
- user_input: `I am looking for a salesforce extractor component`
1041-
- returns a list of component IDs that match the query, ordered by relevance/best match.
1042-
"""
1043-
client = KeboolaClient.from_state(ctx.session.state)
1044-
suggestion_response = await client.ai_service_client.suggest_component(query)
1045-
return suggestion_response.components

src/keboola_mcp_server/tools/search.py

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from fastmcp.tools import FunctionTool
88
from pydantic import BaseModel, Field
99

10-
from keboola_mcp_server.client import GlobalSearchResponse, ItemType, KeboolaClient
10+
from keboola_mcp_server.client import GlobalSearchResponse, ItemType, KeboolaClient, SuggestedComponent
1111
from keboola_mcp_server.errors import tool_errors
1212
from keboola_mcp_server.mcp import with_session_state
1313

@@ -20,6 +20,7 @@
2020
def add_search_tools(mcp: FastMCP) -> None:
2121
"""Add tools to the MCP server."""
2222
search_tools = [
23+
find_component_id,
2324
find_ids_by_name,
2425
]
2526
for tool in search_tools:
@@ -146,3 +147,24 @@ async def find_ids_by_name(
146147
query=joined_prefixes, types=item_types, limit=limit, offset=offset
147148
)
148149
return GlobalSearchOutput.from_api_responses(response)
150+
151+
152+
@tool_errors()
153+
@with_session_state()
154+
async def find_component_id(
155+
ctx: Context,
156+
query: Annotated[str, Field(description='Natural language query to find the requested component.')],
157+
) -> list[SuggestedComponent]:
158+
"""
159+
Returns list of component IDs that match the given query.
160+
161+
USAGE:
162+
- Use when you want to find the component for a specific purpose.
163+
164+
EXAMPLES:
165+
- user_input: `I am looking for a salesforce extractor component`
166+
- returns a list of component IDs that match the query, ordered by relevance/best match.
167+
"""
168+
client = KeboolaClient.from_state(ctx.session.state)
169+
suggestion_response = await client.ai_service_client.suggest_component(query)
170+
return suggestion_response.components

0 commit comments

Comments
 (0)