Databricks MCP Server

import os import pytest import tempfile from unittest.mock import patch, Mock from dotenv import load_dotenv # Load environment variables load_dotenv() # Import the MCP tool function directly from main import mcp_upload_sample_py_to_databricks, mcp_create_databricks_subfolder def test_mcp_upload_sample_py_to_databricks_real(): """ Test the MCP tool that uploads a hardcoded Python file to Databricks. This test will make actual API calls to Databricks. """ # Check if Databricks credentials are available host = os.getenv("DATABRICKS_HOST") token = os.getenv("DATABRICKS_TOKEN") if not host or not token: pytest.skip("DATABRICKS_HOST and DATABRICKS_TOKEN must be set for real upload test") print(f"Testing with Databricks host: {host}") # Upload the Python file directly to /Users/stephen.hsu/sample_claude_upload.py print("Uploading Python file...") result = mcp_upload_sample_py_to_databricks() # Assert the result structure assert "status" in result # Check if upload was successful if result["status"] == "success": print(f"✅ Upload successful: {result['workspace_path']}") assert "message" in result assert "workspace_path" in result assert result["workspace_path"] == "/Users/stephen.hsu@databricks.com/claude_generated_folder/sample_claude_upload.py" assert "Sample Python file uploaded successfully" in result["message"] else: print(f"❌ Upload failed: {result.get('detail', 'Unknown error')}") pytest.skip(f"Upload failed: {result.get('detail', 'Unknown error')}") def test_mcp_create_databricks_sample_folder(): """ Test the MCP tool that creates a sample folder in Databricks workspace. """ # Check if Databricks credentials are available host = os.getenv("DATABRICKS_HOST") token = os.getenv("DATABRICKS_TOKEN") if not host or not token: pytest.skip("DATABRICKS_HOST and DATABRICKS_TOKEN must be set for real folder creation test") print(f"Testing folder creation with Databricks host: {host}") # Call the folder creation function result = mcp_create_databricks_sample_folder() # Assert the result structure assert "status" in result if result["status"] == "success": print(f"✅ Folder creation successful: {result.get('workspace_path', 'unknown')}") assert "message" in result assert "workspace_path" in result assert result["workspace_path"] == "/Users/stephen.hsu/sample_claude_upload" else: print(f"❌ Folder creation failed: {result.get('detail', 'Unknown error')}") pytest.skip(f"Folder creation failed: {result.get('detail', 'Unknown error')}") def test_mcp_create_databricks_subfolder(): """ Test the MCP tool that creates a subfolder in Databricks workspace. """ # Check if Databricks credentials are available host = os.getenv("DATABRICKS_HOST") token = os.getenv("DATABRICKS_TOKEN") if not host or not token: pytest.skip("DATABRICKS_HOST and DATABRICKS_TOKEN must be set for real folder creation test") print(f"Testing subfolder creation with Databricks host: {host}") # Call the folder creation function result = mcp_create_databricks_subfolder("test_claude_folder") # Assert the result structure assert "status" in result if result["status"] == "success": print(f"✅ Subfolder creation successful: {result.get('workspace_path', 'unknown')}") assert "message" in result assert "workspace_path" in result assert result["workspace_path"] == "/Users/stephen.hsu@databricks.com/test_claude_folder" else: print(f"❌ Subfolder creation failed: {result.get('detail', 'Unknown error')}") pytest.skip(f"Subfolder creation failed: {result.get('detail', 'Unknown error')}") def test_mcp_upload_sample_py_to_databricks_mocked(): """ Test the MCP tool with mocked Databricks manager to ensure the function works correctly. """ with patch('main.get_manager') as mock_get_manager: # Create a mock manager mock_manager = Mock() mock_manager.upload_notebook.return_value = "/Users/stephen.hsu@databricls.com/sample_claude_upload.py" mock_get_manager.return_value = mock_manager # Call the function result = mcp_upload_sample_py_to_databricks() # Verify the result assert result["status"] == "success" assert result["workspace_path"] == "/Users/stephen.hsu@databricls.com/sample_claude_upload.py" assert "Sample Python file uploaded successfully" in result["message"] # Verify the manager was called correctly mock_manager.upload_notebook.assert_called_once() call_args = mock_manager.upload_notebook.call_args assert call_args[0][1] == "/Users/stephen.hsu@databricls.com/sample_claude_upload.py" # workspace_path assert "def hello():" in call_args[0][0] # code content def test_mcp_upload_sample_py_to_databricks_error_handling(): """ Test error handling in the MCP tool. """ with patch('main.get_manager') as mock_get_manager: # Mock an exception mock_manager = Mock() mock_manager.upload_notebook.side_effect = Exception("Test error") mock_get_manager.return_value = mock_manager # Call the function result = mcp_upload_sample_py_to_databricks() # Verify error handling assert result["status"] == "error" assert "Test error" in result["detail"] def test_mcp_create_project_folder_and_upload_assets(): """ Test the comprehensive project folder creation and asset upload tool. """ # Check if Databricks credentials are available host = os.getenv("DATABRICKS_HOST") token = os.getenv("DATABRICKS_TOKEN") if not host or not token: pytest.skip("DATABRICKS_HOST and DATABRICKS_TOKEN must be set for real project test") print(f"Testing project folder creation with Databricks host: {host}") # Test assets assets = { "main.py": """ def main(): print("Hello from project!") if __name__ == "__main__": main() """, "config.json": '{"project": "test", "version": "1.0.0"}', "README.md": "# Test Project\n\nThis is a test project created by Claude." } # Call the project creation function from main import mcp_create_project_folder_and_upload_assets result = mcp_create_project_folder_and_upload_assets("test_claude_project", assets) # Assert the result structure assert "status" in result if result["status"] == "success": print(f"✅ Project creation successful: {result.get('project_folder', 'unknown')}") assert "message" in result assert "project_folder" in result assert "uploaded_files" in result assert result["project_folder"] == "/Users/stephen.hsu@databricks.com/claude_generated_folder/test_claude_project" # Check uploaded files uploaded_files = result["uploaded_files"] assert len(uploaded_files) == 3 for file_info in uploaded_files: if file_info["status"] == "success": print(f" ✅ {file_info['filename']}: {file_info['path']}") else: print(f" ❌ {file_info['filename']}: {file_info.get('detail', 'Unknown error')}") else: print(f"❌ Project creation failed: {result.get('detail', 'Unknown error')}") pytest.skip(f"Project creation failed: {result.get('detail', 'Unknown error')}") def test_claude_generated_folder_handling(): """ Test that the claude_generated_folder is properly handled when it already exists. """ # Check if Databricks credentials are available host = os.getenv("DATABRICKS_HOST") token = os.getenv("DATABRICKS_TOKEN") if not host or not token: pytest.skip("DATABRICKS_HOST and DATABRICKS_TOKEN must be set for real folder test") print(f"Testing claude_generated_folder handling with Databricks host: {host}") # Test uploading to the same folder multiple times to ensure it handles existing folder from main import mcp_upload_sample_py_to_databricks # First upload print("First upload (should create folder):") result1 = mcp_upload_sample_py_to_databricks() assert result1["status"] == "success" # Second upload (should use existing folder) print("Second upload (should use existing folder):") result2 = mcp_upload_sample_py_to_databricks() assert result2["status"] == "success" print("✅ Both uploads successful - folder handling works correctly") def test_mcp_create_and_run_job_for_notebook(): """ Test creating and running a job for the uploaded Python file. """ # Check if Databricks credentials are available host = os.getenv("DATABRICKS_HOST") token = os.getenv("DATABRICKS_TOKEN") if not host or not token: pytest.skip("DATABRICKS_HOST and DATABRICKS_TOKEN must be set for real job test") print(f"Testing job creation and execution with Databricks host: {host}") # Test notebook path (the uploaded Python file) notebook_path = "/Users/stephen.hsu@databricks.com/claude_generated_folder/sample_claude_upload.py" # Call the job creation and execution function from main import mcp_create_and_run_job_for_notebook result = mcp_create_and_run_job_for_notebook(notebook_path, "Test Claude Python Job") # Assert the result structure assert "status" in result if result["status"] == "success": print(f"✅ Job creation and execution successful!") print(f" Job ID: {result.get('job_id', 'unknown')}") print(f" Run ID: {result.get('run_id', 'unknown')}") print(f" Job URL: {result.get('job_url', 'unknown')}") print(f" Run URL: {result.get('run_url', 'unknown')}") print(f" Notebook: {result.get('notebook_path', 'unknown')}") print(f" Message: {result.get('message', 'unknown')}") # Verify all required fields are present assert "job_id" in result assert "run_id" in result assert "job_url" in result assert "run_url" in result assert "notebook_path" in result assert "message" in result # Verify URLs are properly formatted assert host in result["job_url"] assert host in result["run_url"] assert str(result["job_id"]) in result["job_url"] assert str(result["run_id"]) in result["run_url"] print(f"\n📋 To monitor the job:") print(f" Job overview: {result['job_url']}") print(f" Run details: {result['run_url']}") print(f" The job should print 'Hello from Claude!' when it runs successfully.") else: print(f"❌ Job creation failed: {result.get('detail', 'Unknown error')}") pytest.skip(f"Job creation failed: {result.get('detail', 'Unknown error')}") if __name__ == "__main__": # Run the real upload test if executed directly print("Running real Databricks upload test...") test_mcp_upload_sample_py_to_databricks_real()