"""Tests for Livy MCP tools."""
from unittest.mock import Mock
import pytest
from ms_fabric_mcp_server.client.exceptions import (
FabricLivyError,
FabricLivyTimeoutError,
)
from ms_fabric_mcp_server.tools.livy_tools import register_livy_tools
from tests.fabric.tools.utils import capture_tools
@pytest.mark.unit
class TestLivyTools:
def test_livy_tools_smoke(self):
tools, mcp = capture_tools()
livy_service = Mock()
livy_service.create_session.return_value = {"id": "1", "state": "idle"}
livy_service.list_sessions.return_value = {"sessions": [{"id": "1"}]}
livy_service.get_session_status.return_value = {"id": "1", "state": "idle"}
livy_service.close_session.return_value = {"msg": "closed"}
livy_service.run_statement.return_value = {"id": "1", "state": "available"}
livy_service.get_statement_status.return_value = {"id": "1", "state": "available"}
livy_service.cancel_statement.return_value = {"msg": "canceled"}
livy_service.get_session_log.return_value = {"log": ["line"]}
register_livy_tools(mcp, livy_service)
assert tools["livy_create_session"](
workspace_id="ws-1", lakehouse_id="lh-1"
)["id"] == "1"
assert tools["livy_list_sessions"](
workspace_id="ws-1", lakehouse_id="lh-1"
)["sessions"][0]["id"] == "1"
assert tools["livy_get_session_status"](
workspace_id="ws-1", lakehouse_id="lh-1", session_id="1"
)["state"] == "idle"
assert tools["livy_close_session"](
workspace_id="ws-1", lakehouse_id="lh-1", session_id="1"
)["msg"] == "closed"
assert tools["livy_run_statement"](
workspace_id="ws-1", lakehouse_id="lh-1", session_id="1", code="1+1"
)["state"] == "available"
assert tools["livy_get_statement_status"](
workspace_id="ws-1", lakehouse_id="lh-1", session_id="1", statement_id="1"
)["state"] == "available"
assert tools["livy_cancel_statement"](
workspace_id="ws-1", lakehouse_id="lh-1", session_id="1", statement_id="1"
)["msg"] == "canceled"
assert tools["livy_get_session_log"](
workspace_id="ws-1", lakehouse_id="lh-1", session_id="1"
)["log"] == ["line"]
def test_livy_run_statement_passes_code_verbatim(self):
tools, mcp = capture_tools()
livy_service = Mock()
livy_service.run_statement.return_value = {"id": "1", "state": "available"}
register_livy_tools(mcp, livy_service)
code = "x = 1\\nx + 1"
tools["livy_run_statement"](
workspace_id="ws-1",
lakehouse_id="lh-1",
session_id="1",
code=code,
)
assert livy_service.run_statement.call_args.kwargs["code"] == code
def test_run_statement_timeout_returns_error(self):
"""livy_run_statement returns structured error on timeout, not empty dict."""
tools, mcp = capture_tools()
livy_service = Mock()
livy_service.run_statement.side_effect = FabricLivyTimeoutError(
"statement execution", 600
)
register_livy_tools(mcp, livy_service)
result = tools["livy_run_statement"](
workspace_id="ws-1", lakehouse_id="lh-1", session_id="1", code="long_code()"
)
assert result["status"] == "error"
assert result["error_code"] == "LIVY_TIMEOUT"
assert "timed out" in result["message"].lower()
def test_run_statement_unexpected_error_returns_error(self):
"""livy_run_statement returns structured error on unexpected exceptions."""
tools, mcp = capture_tools()
livy_service = Mock()
livy_service.run_statement.side_effect = ConnectionError("pipe broken")
register_livy_tools(mcp, livy_service)
result = tools["livy_run_statement"](
workspace_id="ws-1", lakehouse_id="lh-1", session_id="1", code="x"
)
assert result["status"] == "error"
assert "pipe broken" in result["message"]
def test_run_statement_spark_error_enriched(self):
"""livy_run_statement surfaces output.status=error at top level."""
tools, mcp = capture_tools()
livy_service = Mock()
livy_service.run_statement.return_value = {
"id": "5",
"state": "available",
"output": {
"status": "error",
"ename": "KeyError",
"evalue": "'column_x' not in index",
},
}
register_livy_tools(mcp, livy_service)
result = tools["livy_run_statement"](
workspace_id="ws-1", lakehouse_id="lh-1", session_id="1", code="df['x']"
)
assert result["status"] == "error"
assert result["error_summary"] == "KeyError: 'column_x' not in index"
# Original fields still present
assert result["state"] == "available"
assert result["output"]["ename"] == "KeyError"
def test_get_statement_status_spark_error_enriched(self):
"""livy_get_statement_status surfaces output.status=error at top level."""
tools, mcp = capture_tools()
livy_service = Mock()
livy_service.get_statement_status.return_value = {
"id": "3",
"state": "available",
"output": {
"status": "error",
"ename": "NameError",
"evalue": "name 'df' is not defined",
},
}
register_livy_tools(mcp, livy_service)
result = tools["livy_get_statement_status"](
workspace_id="ws-1", lakehouse_id="lh-1", session_id="1", statement_id="3"
)
assert result["status"] == "error"
assert result["error_summary"] == "NameError: name 'df' is not defined"
assert result["state"] == "available"
def test_get_statement_status_ok_no_enrichment(self):
"""livy_get_statement_status does NOT add status field when output is ok."""
tools, mcp = capture_tools()
livy_service = Mock()
livy_service.get_statement_status.return_value = {
"id": "3",
"state": "available",
"output": {"status": "ok", "data": {"text/plain": "42"}},
}
register_livy_tools(mcp, livy_service)
result = tools["livy_get_statement_status"](
workspace_id="ws-1", lakehouse_id="lh-1", session_id="1", statement_id="3"
)
assert "status" not in result or result.get("status") != "error"
assert "error_summary" not in result
