diff --git a/tests/unit/test_cli/test_tool_commands.py b/tests/unit/test_cli/test_tool_commands.py index 913fe416..9c204614 100644 --- a/tests/unit/test_cli/test_tool_commands.py +++ b/tests/unit/test_cli/test_tool_commands.py @@ -55,6 +55,9 @@ class TestSetToolStructuredQuery: mcp_tool=None, collection="sales_data", template=None, + schema_name=None, + index_name=None, + limit=None, arguments=[], group=None, state=None, @@ -92,6 +95,9 @@ class TestSetToolStructuredQuery: mcp_tool=None, collection=None, # No collection specified template=None, + schema_name=None, + index_name=None, + limit=None, arguments=[], group=None, state=None, @@ -132,6 +138,9 @@ class TestSetToolStructuredQuery: mcp_tool=None, collection='sales_data', template=None, + schema_name=None, + index_name=None, + limit=None, arguments=[], group=None, state=None, @@ -201,6 +210,144 @@ class TestSetToolStructuredQuery: assert 'Exception:' in printed_output or 'invalid choice:' in printed_output.lower() +class TestSetToolRowEmbeddingsQuery: + """Test the set_tool function with row-embeddings-query type.""" + + @patch('trustgraph.cli.set_tool.Api') + def test_set_row_embeddings_query_tool_full(self, mock_api_class, mock_api, capsys): + """Test setting a row-embeddings-query tool with all parameters.""" + mock_api_class.return_value, mock_config = mock_api + mock_config.get.return_value = [] + + set_tool( + url="http://test.com", + id="customer_search", + name="find_customer", + description="Find customers by name using semantic search", + type="row-embeddings-query", + mcp_tool=None, + collection="sales", + template=None, + schema_name="customers", + index_name="full_name", + limit=20, + arguments=[], + group=None, + state=None, + applicable_states=None + ) + + captured = capsys.readouterr() + assert "Tool set." in captured.out + + # Verify the tool was stored correctly + call_args = mock_config.put.call_args[0][0] + assert len(call_args) == 1 + config_value = call_args[0] + assert config_value.type == "tool" + assert config_value.key == "customer_search" + + stored_tool = json.loads(config_value.value) + assert stored_tool["name"] == "find_customer" + assert stored_tool["type"] == "row-embeddings-query" + assert stored_tool["collection"] == "sales" + assert stored_tool["schema-name"] == "customers" + assert stored_tool["index-name"] == "full_name" + assert stored_tool["limit"] == 20 + + @patch('trustgraph.cli.set_tool.Api') + def test_set_row_embeddings_query_tool_minimal(self, mock_api_class, mock_api, capsys): + """Test setting row-embeddings-query tool with minimal parameters.""" + mock_api_class.return_value, mock_config = mock_api + mock_config.get.return_value = [] + + set_tool( + url="http://test.com", + id="product_search", + name="find_product", + description="Find products using semantic search", + type="row-embeddings-query", + mcp_tool=None, + collection=None, + template=None, + schema_name="products", + index_name=None, # No index filter + limit=None, # Use default + arguments=[], + group=None, + state=None, + applicable_states=None + ) + + captured = capsys.readouterr() + assert "Tool set." in captured.out + + call_args = mock_config.put.call_args[0][0] + stored_tool = json.loads(call_args[0].value) + assert stored_tool["type"] == "row-embeddings-query" + assert stored_tool["schema-name"] == "products" + assert "index-name" not in stored_tool # Should not be included if None + assert "limit" not in stored_tool # Should not be included if None + assert "collection" not in stored_tool # Should not be included if None + + def test_set_main_row_embeddings_query_with_all_options(self): + """Test set main() with row-embeddings-query tool type and all options.""" + test_args = [ + 'tg-set-tool', + '--id', 'customer_search', + '--name', 'find_customer', + '--type', 'row-embeddings-query', + '--description', 'Find customers by name', + '--schema-name', 'customers', + '--collection', 'sales', + '--index-name', 'full_name', + '--limit', '25', + '--api-url', 'http://custom.com' + ] + + with patch('sys.argv', test_args), \ + patch('trustgraph.cli.set_tool.set_tool') as mock_set: + + set_main() + + mock_set.assert_called_once_with( + url='http://custom.com', + id='customer_search', + name='find_customer', + description='Find customers by name', + type='row-embeddings-query', + mcp_tool=None, + collection='sales', + template=None, + schema_name='customers', + index_name='full_name', + limit=25, + arguments=[], + group=None, + state=None, + applicable_states=None, + token=None + ) + + def test_valid_types_includes_row_embeddings_query(self): + """Test that 'row-embeddings-query' is included in valid tool types.""" + test_args = [ + 'tg-set-tool', + '--id', 'test_tool', + '--name', 'test_tool', + '--type', 'row-embeddings-query', + '--description', 'Test tool', + '--schema-name', 'test_schema' + ] + + with patch('sys.argv', test_args), \ + patch('trustgraph.cli.set_tool.set_tool') as mock_set: + + # Should not raise an exception about invalid type + set_main() + mock_set.assert_called_once() + + class TestShowToolsStructuredQuery: """Test the show_tools function with structured-query tools.""" @@ -259,9 +406,9 @@ class TestShowToolsStructuredQuery: @patch('trustgraph.cli.show_tools.Api') def test_show_mixed_tool_types(self, mock_api_class, mock_api, capsys): - """Test displaying multiple tool types including structured-query.""" + """Test displaying multiple tool types including structured-query and row-embeddings-query.""" mock_api_class.return_value, mock_config = mock_api - + tools = [ { "name": "ask_knowledge", @@ -270,37 +417,47 @@ class TestShowToolsStructuredQuery: "collection": "docs" }, { - "name": "query_data", + "name": "query_data", "description": "Query structured data", "type": "structured-query", "collection": "sales" }, + { + "name": "find_customer", + "description": "Find customers by semantic search", + "type": "row-embeddings-query", + "schema-name": "customers", + "collection": "crm" + }, { "name": "complete_text", "description": "Generate text", "type": "text-completion" } ] - + config_values = [ ConfigValue(type="tool", key=f"tool_{i}", value=json.dumps(tool)) for i, tool in enumerate(tools) ] mock_config.get_values.return_value = config_values - + show_config("http://test.com") - + captured = capsys.readouterr() output = captured.out - + # All tool types should be displayed assert "knowledge-query" in output - assert "structured-query" in output + assert "structured-query" in output + assert "row-embeddings-query" in output assert "text-completion" in output - + # Collections should be shown for appropriate tools assert "docs" in output # knowledge-query collection assert "sales" in output # structured-query collection + assert "crm" in output # row-embeddings-query collection + assert "customers" in output # row-embeddings-query schema-name def test_show_main_parses_args_correctly(self): """Test that show main() parses arguments correctly.""" @@ -317,6 +474,76 @@ class TestShowToolsStructuredQuery: mock_show.assert_called_once_with(url='http://custom.com', token=None) +class TestShowToolsRowEmbeddingsQuery: + """Test the show_tools function with row-embeddings-query tools.""" + + @patch('trustgraph.cli.show_tools.Api') + def test_show_row_embeddings_query_tool_full(self, mock_api_class, mock_api, capsys): + """Test displaying a row-embeddings-query tool with all fields.""" + mock_api_class.return_value, mock_config = mock_api + + tool_config = { + "name": "find_customer", + "description": "Find customers by name using semantic search", + "type": "row-embeddings-query", + "collection": "sales", + "schema-name": "customers", + "index-name": "full_name", + "limit": 20 + } + + config_value = ConfigValue( + type="tool", + key="customer_search", + value=json.dumps(tool_config) + ) + mock_config.get_values.return_value = [config_value] + + show_config("http://test.com") + + captured = capsys.readouterr() + output = captured.out + + # Check that tool information is displayed + assert "customer_search" in output + assert "find_customer" in output + assert "row-embeddings-query" in output + assert "sales" in output # Collection + assert "customers" in output # Schema name + assert "full_name" in output # Index name + assert "20" in output # Limit + + @patch('trustgraph.cli.show_tools.Api') + def test_show_row_embeddings_query_tool_minimal(self, mock_api_class, mock_api, capsys): + """Test displaying row-embeddings-query tool with minimal fields.""" + mock_api_class.return_value, mock_config = mock_api + + tool_config = { + "name": "find_product", + "description": "Find products using semantic search", + "type": "row-embeddings-query", + "schema-name": "products" + # No collection, index-name, or limit + } + + config_value = ConfigValue( + type="tool", + key="product_search", + value=json.dumps(tool_config) + ) + mock_config.get_values.return_value = [config_value] + + show_config("http://test.com") + + captured = capsys.readouterr() + output = captured.out + + # Should display the tool with schema-name + assert "product_search" in output + assert "row-embeddings-query" in output + assert "products" in output # Schema name + + class TestStructuredQueryToolValidation: """Test validation specific to structured-query tools."""