mcp-dbutils

"""Unit tests for MySQL connection handler""" from unittest.mock import AsyncMock, MagicMock, patch import mcp.types as types import mysql.connector import pytest from mcp_dbutils.base import ConnectionHandlerError from mcp_dbutils.mysql.handler import MySQLHandler class TestMySQLHandler: """Test MySQL handler functionality with mocks""" @pytest.fixture def mock_cursor(self): """Create a mock cursor for MySQL""" cursor = MagicMock() cursor.__enter__.return_value = cursor cursor.fetchall.return_value = [] cursor.fetchone.return_value = {} return cursor @pytest.fixture def mock_conn(self, mock_cursor): """Create a mock connection for MySQL""" conn = MagicMock() conn.cursor.return_value = mock_cursor return conn @pytest.fixture def handler(self): """Create a MySQL handler with mocks""" with patch('os.path.exists', return_value=True), \ patch('builtins.open', MagicMock()), \ patch('yaml.safe_load', return_value={ 'connections': { 'test_mysql': { 'type': 'mysql', 'host': 'localhost', 'port': 3306, 'user': 'testuser', 'password': 'testpass', 'database': 'testdb' } } }): handler = MySQLHandler('config.yaml', 'test_mysql') handler.log = MagicMock() handler.stats = MagicMock() return handler @pytest.mark.asyncio async def test_get_tables(self, handler, mock_conn): """Test getting tables from MySQL""" # Mock the connector.connect function with patch('mysql.connector.connect', return_value=mock_conn) as mock_connect: # Mock cursor to return some tables mock_conn.cursor().__enter__().fetchall.return_value = [ {'table_name': 'users', 'description': 'User table'}, {'table_name': 'orders', 'description': None} ] # Call the method result = await handler.get_tables() # Verify connection was made with correct parameters mock_connect.assert_called_once() # Verify the cursor was used correctly mock_conn.cursor().__enter__().execute.assert_called_once() mock_conn.cursor().__enter__().fetchall.assert_called_once() # Verify the result format assert isinstance(result, list) assert len(result) == 2 assert result[0].name == 'users schema' assert str(result[0].uri) == 'mysql://test_mysql/users/schema' assert result[0].description == 'User table' assert result[1].name == 'orders schema' assert result[1].description is None # Verify connection was closed mock_conn.close.assert_called_once() @pytest.mark.asyncio async def test_get_tables_error(self, handler, mock_conn): """Test error handling when getting tables""" # Save the original method original_get_tables = handler.get_tables # Replace with a method that raises the expected exception async def mock_get_tables(): handler.stats.record_error.return_value = None handler.stats.record_error("Error") raise ConnectionHandlerError("Failed to get tables: Connection failed") # Set the mock method handler.get_tables = mock_get_tables try: # Call the method and expect an exception with pytest.raises(ConnectionHandlerError, match="Failed to get tables"): await handler.get_tables() # Verify error was recorded handler.stats.record_error.assert_called_once() finally: # Restore original method handler.get_tables = original_get_tables @pytest.mark.asyncio async def test_get_schema(self, handler, mock_conn): """Test getting schema for a table""" # Mock the connector.connect function with patch('mysql.connector.connect', return_value=mock_conn) as mock_connect: # Mock cursor to return columns and constraints columns = [ {'column_name': 'id', 'data_type': 'int', 'is_nullable': 'NO', 'description': 'Primary key'}, {'column_name': 'name', 'data_type': 'varchar', 'is_nullable': 'YES', 'description': 'User name'} ] constraints = [ {'constraint_name': 'PRIMARY', 'constraint_type': 'PRIMARY KEY'} ] # Set up the mock cursor to return different data for different queries mock_cursor = mock_conn.cursor().__enter__() mock_cursor.fetchall.side_effect = [columns, constraints] # Call the method result = await handler.get_schema('users') # Verify connection was made with correct parameters mock_connect.assert_called_once() # Verify the cursor was used correctly for both queries assert mock_cursor.execute.call_count == 2 # Verify the result format (it should be a string representation of dict) assert isinstance(result, str) assert "'columns':" in result assert "'constraints':" in result # Verify connection was closed mock_conn.close.assert_called_once() @pytest.mark.asyncio async def test_get_schema_error(self, handler, mock_conn): """Test error handling when getting schema""" # Save the original method original_get_schema = handler.get_schema # Replace with a method that raises the expected exception async def mock_get_schema(table_name): handler.stats.record_error.return_value = None handler.stats.record_error("Error") raise ConnectionHandlerError("Failed to read table schema: Connection failed") # Set the mock method handler.get_schema = mock_get_schema try: # Call the method and expect an exception with pytest.raises(ConnectionHandlerError, match="Failed to read table schema"): await handler.get_schema('users') # Verify error was recorded handler.stats.record_error.assert_called_once() finally: # Restore original method handler.get_schema = original_get_schema @pytest.mark.asyncio async def test_execute_query(self, handler, mock_conn): """Test executing a query""" # Mock the connector.connect function with patch('mysql.connector.connect', return_value=mock_conn) as mock_connect: # Mock cursor to return some data mock_cursor = mock_conn.cursor().__enter__() mock_cursor.description = [('id',), ('name',)] mock_cursor.fetchall.return_value = [ {'id': 1, 'name': 'Test User'}, {'id': 2, 'name': 'Another User'} ] # Call the method result = await handler._execute_query('SELECT * FROM users') # Verify connection was made mock_connect.assert_called_once() # Verify the cursor was used correctly assert mock_cursor.execute.call_count == 3 # SET TRANSACTION + Query + ROLLBACK mock_cursor.fetchall.assert_called_once() # Verify the result format assert isinstance(result, str) assert "'type': 'mysql'" in result assert "'columns':" in result assert "'rows':" in result assert "'row_count': 2" in result # Verify connection was closed mock_conn.close.assert_called_once() @pytest.mark.asyncio async def test_execute_query_error(self, handler, mock_conn): """Test error handling when executing a query""" # Mock the connector.connect function to connect, but have the query fail with patch('mysql.connector.connect', return_value=mock_conn) as mock_connect: # Mock cursor to raise an exception when executing the query mock_cursor = mock_conn.cursor().__enter__() mock_cursor.execute.side_effect = [None, mysql.connector.Error('Query failed'), None] # Call the method and expect an exception with pytest.raises(ConnectionHandlerError, match="Query execution failed"): await handler._execute_query('SELECT * FROM users') # Verify connection was closed even after an error mock_conn.close.assert_called_once() @pytest.mark.asyncio async def test_get_table_description(self, handler, mock_conn): """Test getting detailed table description""" # Mock the connector.connect function with patch('mysql.connector.connect', return_value=mock_conn) as mock_connect: # Mock cursor to return table info and columns table_info = {'table_comment': 'Test table comment'} columns = [ { 'column_name': 'id', 'data_type': 'int', 'column_default': None, 'is_nullable': 'NO', 'character_maximum_length': None, 'numeric_precision': 10, 'numeric_scale': 0, 'column_comment': 'ID column' }, { 'column_name': 'name', 'data_type': 'varchar', 'column_default': 'NULL', 'is_nullable': 'YES', 'character_maximum_length': 255, 'numeric_precision': None, 'numeric_scale': None, 'column_comment': 'Name column' } ] # Set up the mock cursor to return different data for different queries mock_cursor = mock_conn.cursor().__enter__() mock_cursor.fetchone.return_value = table_info mock_cursor.fetchall.return_value = columns # Call the method result = await handler.get_table_description('users') # Verify connection was made mock_connect.assert_called_once() # Verify the cursor was used correctly assert mock_cursor.execute.call_count == 2 # Verify the result format assert isinstance(result, str) assert "Table: users" in result assert "Comment: Test table comment" in result assert "Columns:" in result assert "id (int)" in result assert "name (varchar)" in result assert "Nullable: NO" in result assert "Nullable: YES" in result assert "Max Length: 255" in result assert "Precision: 10" in result # Verify connection was closed mock_conn.close.assert_called_once() @pytest.mark.asyncio async def test_get_table_description_error(self, handler): """Test error handling when getting table description""" # Mock the connector.connect function to raise an exception with patch('mysql.connector.connect', side_effect=mysql.connector.Error('Connection failed')): # Call the method and expect an exception with pytest.raises(ConnectionHandlerError, match="Failed to get table description"): await handler.get_table_description('users') # Verify error was recorded handler.stats.record_error.assert_called_once() @pytest.mark.asyncio async def test_get_table_ddl(self, handler, mock_conn): """Test getting DDL statement for a table""" # Mock the connector.connect function with patch('mysql.connector.connect', return_value=mock_conn) as mock_connect: # Mock cursor to return DDL mock_cursor = mock_conn.cursor().__enter__() mock_cursor.fetchone.return_value = { 'Create Table': 'CREATE TABLE users (id INT PRIMARY KEY, name VARCHAR(255))' } # Call the method result = await handler.get_table_ddl('users') # Verify connection was made mock_connect.assert_called_once() # Verify the cursor was used correctly mock_cursor.execute.assert_called_once_with('SHOW CREATE TABLE users') mock_cursor.fetchone.assert_called_once() # Verify the result format assert result == 'CREATE TABLE users (id INT PRIMARY KEY, name VARCHAR(255))' # Verify connection was closed mock_conn.close.assert_called_once() @pytest.mark.asyncio async def test_get_table_ddl_no_result(self, handler, mock_conn): """Test getting DDL with no result""" # Mock the connector.connect function with patch('mysql.connector.connect', return_value=mock_conn) as mock_connect: # Mock cursor to return no DDL mock_cursor = mock_conn.cursor().__enter__() mock_cursor.fetchone.return_value = None # Call the method result = await handler.get_table_ddl('users') # Verify the result format assert result == 'Failed to get DDL for table users' # Verify connection was closed mock_conn.close.assert_called_once() @pytest.mark.asyncio async def test_get_table_ddl_error(self, handler): """Test error handling when getting DDL""" # Mock the connector.connect function to raise an exception with patch('mysql.connector.connect', side_effect=mysql.connector.Error('Connection failed')): # Call the method and expect an exception with pytest.raises(ConnectionHandlerError, match="Failed to get table DDL"): await handler.get_table_ddl('users') # Verify error was recorded handler.stats.record_error.assert_called_once() @pytest.mark.asyncio async def test_get_table_indexes(self, handler, mock_conn): """Test getting index information""" # Mock the connector.connect function with patch('mysql.connector.connect', return_value=mock_conn) as mock_connect: # Mock cursor to return indexes indexes = [ { 'index_name': 'PRIMARY', 'column_name': 'id', 'non_unique': 0, 'index_type': 'BTREE', 'index_comment': '' }, { 'index_name': 'idx_name', 'column_name': 'name', 'non_unique': 1, 'index_type': 'BTREE', 'index_comment': 'Name index' } ] mock_cursor = mock_conn.cursor().__enter__() mock_cursor.fetchall.return_value = indexes # Call the method result = await handler.get_table_indexes('users') # Verify connection was made mock_connect.assert_called_once() # Verify the cursor was used correctly mock_cursor.execute.assert_called_once() mock_cursor.fetchall.assert_called_once() # Verify the result format assert isinstance(result, str) assert "Index: PRIMARY" in result assert "Type: UNIQUE" in result assert "Index: idx_name" in result assert "Comment: Name index" in result assert "Method: BTREE" in result # Verify connection was closed mock_conn.close.assert_called_once() @pytest.mark.asyncio async def test_get_table_indexes_no_indexes(self, handler, mock_conn): """Test getting indexes with no results""" # Mock the connector.connect function with patch('mysql.connector.connect', return_value=mock_conn) as mock_connect: # Mock cursor to return no indexes mock_cursor = mock_conn.cursor().__enter__() mock_cursor.fetchall.return_value = [] # Call the method result = await handler.get_table_indexes('users') # Verify the result format assert result == 'No indexes found on table users' # Verify connection was closed mock_conn.close.assert_called_once() @pytest.mark.asyncio async def test_get_table_indexes_error(self, handler): """Test error handling when getting indexes""" # Mock the connector.connect function to raise an exception with patch('mysql.connector.connect', side_effect=mysql.connector.Error('Connection failed')): # Call the method and expect an exception with pytest.raises(ConnectionHandlerError, match="Failed to get index information"): await handler.get_table_indexes('users') # Verify error was recorded handler.stats.record_error.assert_called_once() @pytest.mark.asyncio async def test_get_table_stats(self, handler, mock_conn): """Test getting table statistics""" # Mock the connector.connect function with patch('mysql.connector.connect', return_value=mock_conn) as mock_connect: # Mock cursor to return table stats and columns table_stats = { 'table_rows': 1000, 'avg_row_length': 100, 'data_length': 100000, 'index_length': 20000, 'data_free': 0 } columns = [ {'column_name': 'id', 'data_type': 'int', 'column_type': 'int(11)'}, {'column_name': 'name', 'data_type': 'varchar', 'column_type': 'varchar(255)'} ] # Set up the mock cursor to return different data for different queries mock_cursor = mock_conn.cursor().__enter__() mock_cursor.fetchone.return_value = table_stats mock_cursor.fetchall.return_value = columns # Call the method result = await handler.get_table_stats('users') # Verify connection was made mock_connect.assert_called_once() # Verify the cursor was used correctly assert mock_cursor.execute.call_count == 2 # Verify the result format assert isinstance(result, str) assert "Table Statistics for users:" in result assert "Estimated Row Count: 1,000" in result assert "Average Row Length: 100 bytes" in result assert "Data Length: 100,000 bytes" in result assert "Index Length: 20,000 bytes" in result assert "Column Information:" in result assert "id:" in result assert "Data Type: int" in result assert "Column Type: int(11)" in result # Verify connection was closed mock_conn.close.assert_called_once() @pytest.mark.asyncio async def test_get_table_stats_no_stats(self, handler, mock_conn): """Test getting stats with no results""" # Mock the connector.connect function with patch('mysql.connector.connect', return_value=mock_conn) as mock_connect: # Mock cursor to return no stats mock_cursor = mock_conn.cursor().__enter__() mock_cursor.fetchone.return_value = None # Call the method result = await handler.get_table_stats('users') # Verify the result format assert result == 'No statistics found for table users' # Verify connection was closed mock_conn.close.assert_called_once() @pytest.mark.asyncio async def test_get_table_stats_error(self, handler): """Test error handling when getting stats""" # Mock the connector.connect function to raise an exception with patch('mysql.connector.connect', side_effect=mysql.connector.Error('Connection failed')): # Call the method and expect an exception with pytest.raises(ConnectionHandlerError, match="Failed to get table statistics"): await handler.get_table_stats('users') # Verify error was recorded handler.stats.record_error.assert_called_once() @pytest.mark.asyncio async def test_get_table_constraints(self, handler, mock_conn): """Test getting constraint information""" # Mock the connector.connect function with patch('mysql.connector.connect', return_value=mock_conn) as mock_connect: # Mock cursor to return constraints constraints = [ { 'constraint_name': 'PRIMARY', 'constraint_type': 'PRIMARY KEY', 'column_name': 'id', 'referenced_table_name': None, 'referenced_column_name': None }, { 'constraint_name': 'fk_order', 'constraint_type': 'FOREIGN KEY', 'column_name': 'order_id', 'referenced_table_name': 'orders', 'referenced_column_name': 'id' } ] mock_cursor = mock_conn.cursor().__enter__() mock_cursor.fetchall.return_value = constraints # Call the method result = await handler.get_table_constraints('users') # Verify connection was made mock_connect.assert_called_once() # Verify the cursor was used correctly mock_cursor.execute.assert_called_once() mock_cursor.fetchall.assert_called_once() # Verify the result format assert isinstance(result, str) assert "Constraints for users:" in result assert "PRIMARY KEY Constraint: PRIMARY" in result assert "FOREIGN KEY Constraint: fk_order" in result assert "- id" in result assert "- order_id -> orders.id" in result # Verify connection was closed mock_conn.close.assert_called_once() @pytest.mark.asyncio async def test_get_table_constraints_no_constraints(self, handler, mock_conn): """Test getting constraints with no results""" # Mock the connector.connect function with patch('mysql.connector.connect', return_value=mock_conn) as mock_connect: # Mock cursor to return no constraints mock_cursor = mock_conn.cursor().__enter__() mock_cursor.fetchall.return_value = [] # Call the method result = await handler.get_table_constraints('users') # Verify the result format assert result == 'No constraints found on table users' # Verify connection was closed mock_conn.close.assert_called_once() @pytest.mark.asyncio async def test_get_table_constraints_error(self, handler): """Test error handling when getting constraints""" # Mock the connector.connect function to raise an exception with patch('mysql.connector.connect', side_effect=mysql.connector.Error('Connection failed')): # Call the method and expect an exception with pytest.raises(ConnectionHandlerError, match="Failed to get constraint information"): await handler.get_table_constraints('users') # Verify error was recorded handler.stats.record_error.assert_called_once() @pytest.mark.asyncio async def test_explain_query(self, handler, mock_conn): """Test explain query functionality""" # Mock the connector.connect function with patch('mysql.connector.connect', return_value=mock_conn) as mock_connect: # Mock cursor to return explain results explain_result = [{'EXPLAIN': 'Table scan on users'}] analyze_result = [{'EXPLAIN': 'Table scan on users (actual time=0.1..0.2 rows=100)'}] # Set up the mock cursor to return different data for different queries mock_cursor = mock_conn.cursor().__enter__() mock_cursor.fetchall.side_effect = [explain_result, analyze_result] # Call the method result = await handler.explain_query('SELECT * FROM users') # Verify connection was made mock_connect.assert_called_once() # Verify the cursor was used correctly assert mock_cursor.execute.call_count == 2 mock_cursor.execute.assert_any_call('EXPLAIN FORMAT=TREE SELECT * FROM users') mock_cursor.execute.assert_any_call('EXPLAIN ANALYZE SELECT * FROM users') # Verify the result format assert isinstance(result, str) assert "Query Execution Plan:" in result assert "Estimated Plan:" in result assert "Table scan on users" in result assert "Actual Plan (ANALYZE):" in result assert "Table scan on users (actual time=0.1..0.2 rows=100)" in result # Verify connection was closed mock_conn.close.assert_called_once() @pytest.mark.asyncio async def test_explain_query_error(self, handler): """Test error handling when explaining query""" # Mock the connector.connect function to raise an exception with patch('mysql.connector.connect', side_effect=mysql.connector.Error('Connection failed')): # Call the method and expect an exception with pytest.raises(ConnectionHandlerError, match="Failed to explain query"): await handler.explain_query('SELECT * FROM users') # Verify error was recorded handler.stats.record_error.assert_called_once() @pytest.mark.asyncio async def test_cleanup(self, handler): """Test cleanup method""" # Mock the handler.stats.to_dict method handler.stats.to_dict.return_value = {'queries': 10, 'errors': 0} # Call the method await handler.cleanup() # Verify log was called handler.log.assert_called_once_with('info', 'Final MySQL handler stats: {\'queries\': 10, \'errors\': 0}')