Databricks MCP Server

import os import pytest from fastapi.testclient import TestClient from unittest.mock import patch, Mock from main import app client = TestClient(app) # File endpoints def test_create_folder(tmp_path): folder = tmp_path / "test_folder" resp = client.post("/file/create_folder", json={"path": str(folder)}) assert resp.status_code == 200 assert resp.json()["status"] == "success" assert folder.exists() def test_create_py_file(tmp_path): file_path = tmp_path / "test.py" resp = client.post("/file/create_py_file", json={"path": str(file_path), "content": "print('hi')"}) assert resp.status_code == 200 assert resp.json()["status"] == "success" assert file_path.exists() assert file_path.read_text() == "print('hi')" def test_edit_file_content(tmp_path): file_path = tmp_path / "edit.py" file_path.write_text("old") resp = client.post("/file/edit_file", json={"path": str(file_path), "content": "new"}) assert resp.status_code == 200 assert resp.json()["status"] == "success" assert file_path.read_text() == "new" # Databricks endpoints (mocked) @patch('main.submit_code') def test_submit_code(mock_submit_code): mock_submit_code.return_value = {"result": "ok"} resp = client.post("/databricks/submit_code", json={"code": "print('hi')", "cluster_id": "cid"}) assert resp.status_code == 200 assert resp.json()["status"] == "success" @patch('main.create_job') def test_create_job(mock_create_job): mock_create_job.return_value = {"job_id": "jid"} resp = client.post("/databricks/create_job", json={"job_config": {"name": "job"}}) assert resp.status_code == 200 assert resp.json()["status"] == "success" @patch('main.run_job') def test_run_job(mock_run_job): mock_run_job.return_value = {"run_id": "rid"} resp = client.post("/databricks/run_job", json={"job_id": "jid"}) assert resp.status_code == 200 assert resp.json()["status"] == "success" @patch('main.create_dlt_pipeline') def test_create_dlt_pipeline(mock_create_dlt): mock_create_dlt.return_value = {"pipeline_id": "pid"} resp = client.post("/databricks/create_dlt_pipeline", json={"pipeline_config": {"name": "pipeline"}}) assert resp.status_code == 200 assert resp.json()["status"] == "success" @patch('main.get_job_error') def test_get_job_error(mock_get_error): mock_get_error.return_value = "error message" resp = client.post("/databricks/get_job_error", json={"run_id": "rid"}) assert resp.status_code == 200 assert resp.json()["status"] == "success" assert resp.json()["error"] == "error message" @patch('main.get_manager') def test_check_job_status(mock_get_manager): mock_manager = Mock() mock_status = { "job_id": "jid", "run_id": "rid", "state": "RUNNING", "result_state": "UNKNOWN", "state_message": "", "start_time": 1234567890, "end_time": 0 } mock_manager.check_job_status.return_value = mock_status mock_get_manager.return_value = mock_manager resp = client.post("/databricks/check_job_status", json={"job_id": "jid", "run_id": "rid"}) assert resp.status_code == 200 assert resp.json()["status"] == "success" assert resp.json()["job_status"] == mock_status