mcp-lancedb

import unittest from dotenv import load_dotenv from mcp_clickhouse import list_databases, list_tables, run_select_query, create_clickhouse_client load_dotenv() class TestClickhouseTools(unittest.TestCase): @classmethod def setUpClass(cls): """Set up the environment before tests.""" cls.client = create_clickhouse_client() # Prepare test database and table cls.test_db = "test_tool_db" cls.test_table = "test_table" cls.client.command(f"CREATE DATABASE IF NOT EXISTS {cls.test_db}") cls.client.command(f""" CREATE TABLE IF NOT EXISTS {cls.test_db}.{cls.test_table} ( id UInt32, name String ) ENGINE = MergeTree() ORDER BY id """) cls.client.command(f""" INSERT INTO {cls.test_db}.{cls.test_table} (id, name) VALUES (1, 'Alice'), (2, 'Bob') """) @classmethod def tearDownClass(cls): """Clean up the environment after tests.""" cls.client.command(f"DROP DATABASE IF EXISTS {cls.test_db}") def test_list_databases(self): """Test listing databases.""" result = list_databases() self.assertIn(self.test_db, result) def test_list_tables_without_like(self): """Test listing tables without a 'LIKE' filter.""" result = list_tables(self.test_db) self.assertIsInstance(result, list) self.assertEqual(len(result), 1) self.assertEqual(result[0]["name"], self.test_table) def test_list_tables_with_like(self): """Test listing tables with a 'LIKE' filter.""" result = list_tables(self.test_db, like=f"{self.test_table}%") self.assertIsInstance(result, list) self.assertEqual(len(result), 1) self.assertEqual(result[0]["name"], self.test_table) def test_run_select_query_success(self): """Test running a SELECT query successfully.""" query = f"SELECT * FROM {self.test_db}.{self.test_table}" result = run_select_query(query) self.assertIsInstance(result, list) self.assertEqual(len(result), 2) self.assertEqual(result[0]["id"], 1) self.assertEqual(result[0]["name"], "Alice") def test_run_select_query_failure(self): """Test running a SELECT query with an error.""" query = f"SELECT * FROM {self.test_db}.non_existent_table" result = run_select_query(query) self.assertIsInstance(result, str) self.assertIn("error running query", result) if __name__ == "__main__": unittest.main()