test_complete_serverless_workflow.py•3.56 kB
#!/usr/bin/env python3
"""
Complete test of the serverless Databricks MCP workflow:
1. Create folder
2. Upload Python file
3. Create and run serverless job
"""
import os
import sys
import tempfile
from dotenv import load_dotenv
# Load environment variables
load_dotenv()
def test_complete_serverless_workflow():
"""Test the complete serverless workflow"""
print("🧪 Testing Complete Serverless Databricks MCP Workflow")
print("=" * 60)
# Import MCP functions
from main import (
mcp_create_folder,
mcp_upload_py_file_to_databricks,
mcp_create_and_run_job_for_notebook
)
# Test parameters
username = "stephen.hsu@databricks.com"
folder_path = f"/Users/{username}/claude_generated_folder"
# Create a test Python file
test_code = '''
# Test Python file for serverless execution
import datetime
print("🚀 Starting serverless job execution...")
print(f"Timestamp: {datetime.datetime.now()}")
# Simulate some work
import time
time.sleep(2)
print("✅ Serverless job completed successfully!")
print(f"Final timestamp: {datetime.datetime.now()}")
'''
try:
# Step 1: Create folder
print("\n📁 Step 1: Creating folder...")
folder_result = mcp_create_folder(folder_path)
print(f"Folder creation result: {folder_result}")
if folder_result.get('status') != 'success':
print("❌ Folder creation failed")
return
# Step 2: Upload Python file
print("\n📤 Step 2: Uploading Python file...")
filename = "serverless_test.py"
upload_result = mcp_upload_py_file_to_databricks(
filename=filename,
content=test_code,
username=username
)
print(f"Upload result: {upload_result}")
if upload_result.get('status') != 'success':
print("❌ File upload failed")
return
notebook_path = upload_result.get('notebook_path')
print(f"✅ File uploaded to: {notebook_path}")
# Step 3: Create and run serverless job
print("\n⚡ Step 3: Creating and running serverless job...")
job_result = mcp_create_and_run_job_for_notebook(
notebook_path=notebook_path,
job_name="Complete Serverless Test Job",
username=username
)
print(f"Job result: {job_result}")
if job_result.get('status') == 'success':
print("\n🎉 SUCCESS: Complete serverless workflow executed!")
print(f"Job ID: {job_result.get('job_id')}")
print(f"Run ID: {job_result.get('run_id')}")
print(f"Job URL: {job_result.get('job_url')}")
print(f"Run URL: {job_result.get('run_url')}")
print(f"Message: {job_result.get('message')}")
print("\n📊 Summary:")
print("✅ Folder created successfully")
print("✅ Python file uploaded to Databricks")
print("✅ Serverless job created and started")
print("✅ No cluster permission errors encountered")
print("\n🔗 Monitor your job in Databricks UI using the provided URLs")
else:
print(f"❌ Job creation failed: {job_result.get('detail')}")
except Exception as e:
print(f"❌ Exception occurred: {str(e)}")
import traceback
traceback.print_exc()
if __name__ == "__main__":
test_complete_serverless_workflow()