Supabase MCP Server

from unittest.mock import AsyncMock, MagicMock import pytest from supabase_mcp.exceptions import SafetyError from supabase_mcp.services.database.query_manager import QueryManager from supabase_mcp.services.database.sql.loader import SQLLoader from supabase_mcp.services.database.sql.validator import ( QueryValidationResults, SQLQueryCategory, SQLQueryCommand, SQLValidator, ValidatedStatement, ) from supabase_mcp.services.safety.models import ClientType, OperationRiskLevel @pytest.mark.asyncio(loop_scope="module") class TestQueryManager: """Tests for the Query Manager.""" @pytest.mark.unit async def test_query_execution(self, mock_query_manager: QueryManager): """Test query execution through the Query Manager.""" query_manager = mock_query_manager # Ensure validator and safety_manager are proper mocks query_manager.validator = MagicMock() query_manager.safety_manager = MagicMock() # Create a mock validation result for a SELECT query validated_statement = ValidatedStatement( category=SQLQueryCategory.DQL, command=SQLQueryCommand.SELECT, risk_level=OperationRiskLevel.LOW, query="SELECT * FROM users", needs_migration=False, object_type="TABLE", schema_name="public", ) validation_result = QueryValidationResults( statements=[validated_statement], highest_risk_level=OperationRiskLevel.LOW, has_transaction_control=False, original_query="SELECT * FROM users", ) # Make the validator return our mock validation result query_manager.validator.validate_query.return_value = validation_result # Make the db_client return a mock query result mock_query_result = MagicMock() query_manager.db_client.execute_query = AsyncMock(return_value=mock_query_result) # Execute a query query = "SELECT * FROM users" result = await query_manager.handle_query(query) # Verify the validator was called with the query query_manager.validator.validate_query.assert_called_once_with(query) # Verify the db_client was called with the validation result query_manager.db_client.execute_query.assert_called_once_with(validation_result, False) # Verify the result is what we expect assert result == mock_query_result @pytest.mark.asyncio @pytest.mark.unit async def test_safety_validation_blocks_dangerous_query(self, mock_query_manager: QueryManager): """Test that the safety validation blocks dangerous queries.""" # Create a query manager with the mock dependencies query_manager = mock_query_manager # Ensure validator and safety_manager are proper mocks query_manager.validator = MagicMock() query_manager.safety_manager = MagicMock() # Create a mock validation result for a DROP TABLE query validated_statement = ValidatedStatement( category=SQLQueryCategory.DDL, command=SQLQueryCommand.DROP, risk_level=OperationRiskLevel.EXTREME, query="DROP TABLE users", needs_migration=False, object_type="TABLE", schema_name="public", ) validation_result = QueryValidationResults( statements=[validated_statement], highest_risk_level=OperationRiskLevel.EXTREME, has_transaction_control=False, original_query="DROP TABLE users", ) # Make the validator return our mock validation result query_manager.validator.validate_query.return_value = validation_result # Make the safety manager raise a SafetyError error_message = "Operation not allowed in SAFE mode" query_manager.safety_manager.validate_operation.side_effect = SafetyError(error_message) # Execute a query - should raise a SafetyError query = "DROP TABLE users" with pytest.raises(SafetyError) as excinfo: await query_manager.handle_query(query) # Verify the error message assert error_message in str(excinfo.value) # Verify the validator was called with the query query_manager.validator.validate_query.assert_called_once_with(query) # Verify the safety manager was called with the validation result query_manager.safety_manager.validate_operation.assert_called_once_with( ClientType.DATABASE, validation_result, False ) # Verify the db_client was not called query_manager.db_client.execute_query.assert_not_called() @pytest.mark.unit async def test_get_migrations_query(self, query_manager_integration: QueryManager): """Test that get_migrations_query returns a valid query string.""" # Test with default parameters query = query_manager_integration.get_migrations_query() assert isinstance(query, str) assert "supabase_migrations.schema_migrations" in query assert "LIMIT 50" in query # Test with custom parameters custom_query = query_manager_integration.get_migrations_query( limit=10, offset=5, name_pattern="test", include_full_queries=True ) assert isinstance(custom_query, str) assert "supabase_migrations.schema_migrations" in custom_query assert "LIMIT 10" in custom_query assert "OFFSET 5" in custom_query assert "name ILIKE" in custom_query assert "statements" in custom_query # Should include statements column when include_full_queries=True @pytest.mark.unit async def test_init_migration_schema(self): """Test that init_migration_schema initializes the migration schema correctly.""" # Create minimal mocks postgres_client = MagicMock() postgres_client.execute_query = AsyncMock() safety_manager = MagicMock() # Create a real SQLLoader and SQLValidator sql_loader = SQLLoader() sql_validator = SQLValidator() # Create the QueryManager with minimal mocking query_manager = QueryManager( postgres_client=postgres_client, safety_manager=safety_manager, sql_validator=sql_validator, sql_loader=sql_loader, ) # Call the method await query_manager.init_migration_schema() # Verify that the SQL loader was used to get the init migrations query # and that the query was executed assert postgres_client.execute_query.called # Get the arguments that execute_query was called with call_args = postgres_client.execute_query.call_args assert call_args is not None # The first argument should be a QueryValidationResults object args, _ = call_args # Use _ to ignore unused kwargs assert len(args) > 0 validation_result = args[0] assert isinstance(validation_result, QueryValidationResults) # Check that the validation result contains the expected SQL init_query = sql_loader.get_init_migrations_query() assert any(stmt.query and stmt.query in init_query for stmt in validation_result.statements) @pytest.mark.unit async def test_handle_migration(self): """Test that handle_migration correctly handles migrations when needed.""" # Create minimal mocks postgres_client = MagicMock() postgres_client.execute_query = AsyncMock() safety_manager = MagicMock() # Create a real SQLLoader sql_loader = SQLLoader() # Create a mock MigrationManager migration_manager = MagicMock() migration_query = "INSERT INTO _migrations.migrations (name) VALUES ('test_migration')" migration_name = "test_migration" migration_manager.prepare_migration_query.return_value = (migration_query, migration_name) # Create a real SQLValidator sql_validator = SQLValidator() # Create the QueryManager with minimal mocking query_manager = QueryManager( postgres_client=postgres_client, safety_manager=safety_manager, sql_validator=sql_validator, sql_loader=sql_loader, migration_manager=migration_manager, ) # Create a validation result that needs migration validated_statement = ValidatedStatement( category=SQLQueryCategory.DDL, command=SQLQueryCommand.CREATE, risk_level=OperationRiskLevel.MEDIUM, query="CREATE TABLE test (id INT)", needs_migration=True, object_type="TABLE", schema_name="public", ) validation_result = QueryValidationResults( statements=[validated_statement], highest_risk_level=OperationRiskLevel.MEDIUM, has_transaction_control=False, original_query="CREATE TABLE test (id INT)", ) # Call the method await query_manager.handle_migration(validation_result, "CREATE TABLE test (id INT)", "test_migration") # Verify that the migration manager was called to prepare the migration query migration_manager.prepare_migration_query.assert_called_once_with( validation_result, "CREATE TABLE test (id INT)", "test_migration" ) # Verify that execute_query was called at least twice # Once for init_migration_schema and once for the migration query assert postgres_client.execute_query.call_count >= 2