ClaudeandClaude Fable 5
6cf68439a2
Add 3 new MCP tools for inspecting YARN applications NOT submitted through this service: get_external_job_logs, get_external_job_status, get_external_job_result. Each takes application_id + connection_name and queries YARN directly, bypassing the local JobStore. - spark_executor/tools/external_jobs.py: 3 tool functions - spark_executor/tools/requests.py: 3 new Pydantic body models (ExternalJobLogsRequest, ExternalJobStatusRequest, ExternalJobResultRequest) - spark_executor/server.py: 3 new POST routes with explicit operation_id - tests/unit/test_external_jobs.py: 7 unit tests - tests/integration/test_mcp_routes.py: assert 20 tool routes - README.md: list the 3 new tools To make the LLM pick the right tool and not guess at field values, also: - Add Pydantic field descriptions for 22 fields across 8 request models (SaveConnectionRequest, UpdatePendingJobRequest, GetJobLogsRequest, JobIdRequest, PendingIdRequest, ConnectionNameRequest, plus the new ExternalJob*Request models). - Update 12 route descriptions with cross-references, prerequisite context, and 400 behavior notes. - Refactor _unknown_job_error: an input that looks like a YARN application_id (starts with 'application_') now returns HTTP 400 (ValueError) with a hint message naming the right external tool; other not-found cases still return 404 (KeyError). This catches the common LLM mistake of passing application_id to the internal get_job_* / kill_job tools. - 4 new unit tests for the 400 behavior. Tests: 356 passed (up from 242). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
120 lines
3.9 KiB
Python
120 lines
3.9 KiB
Python
# coding=utf-8
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from spark_executor.core import connection_store
|
|
from spark_executor.core.job_store import JobStore
|
|
from spark_executor.models import Connection, Job
|
|
from spark_executor.tools import kill
|
|
from spark_executor.tools import connections
|
|
|
|
|
|
@pytest.fixture
|
|
def fresh_stores(tmp_path: Path):
|
|
"""Wire up connection + job stores rooted in tmp_path. Per-test isolation
|
|
so file-backed JobStore doesn't leak between cases."""
|
|
store = connection_store.ConnectionStore(data_dir=str(tmp_path))
|
|
connection_store.store = store
|
|
connections.store = store
|
|
kill.conn_store = store
|
|
kill.store = JobStore(data_dir=str(tmp_path))
|
|
|
|
|
|
def _fresh_stores():
|
|
store = connection_store.ConnectionStore()
|
|
connection_store.store = store
|
|
connections.store = store
|
|
kill.conn_store = store
|
|
kill.store = JobStore()
|
|
|
|
|
|
def test_kill_job_calls_yarn_kill():
|
|
_fresh_stores()
|
|
kill.conn_store.save(Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088"))
|
|
kill.store.put(
|
|
Job(
|
|
job_id="abc",
|
|
application_id="application_1",
|
|
script_path="/tmp/j.py",
|
|
queue="default",
|
|
submit_time=datetime(2026, 6, 24),
|
|
connection="prod",
|
|
yarn_rm_url="http://rm:8088",
|
|
)
|
|
)
|
|
with patch("spark_executor.tools.kill.kill_application") as m:
|
|
result = kill.kill_job("abc")
|
|
args = m.call_args.args
|
|
assert args[0] == "application_1"
|
|
assert args[1].yarn_rm_url == "http://rm:8088"
|
|
assert result == {
|
|
"job_id": "abc",
|
|
"application_id": "application_1",
|
|
"status": "KILLED",
|
|
}
|
|
|
|
|
|
def test_kill_job_raises_for_unknown_job():
|
|
_fresh_stores()
|
|
with pytest.raises(KeyError):
|
|
kill.kill_job("missing")
|
|
|
|
|
|
def test_kill_job_raises_400_for_external_application_id(fresh_stores):
|
|
"""Input that looks like a YARN application_id but is not in the local
|
|
JobStore must raise ValueError (-> 400) — kill_job has no external
|
|
equivalent, so the error points to the YARN CLI / UI."""
|
|
with pytest.raises(ValueError, match="YARN CLI"):
|
|
kill.kill_job("application_17400000001_0001")
|
|
|
|
|
|
def test_kill_job_raises_when_connection_missing():
|
|
_fresh_stores()
|
|
kill.store.put(
|
|
Job(
|
|
job_id="abc",
|
|
application_id="application_1",
|
|
script_path="/tmp/j.py",
|
|
queue="default",
|
|
submit_time=datetime(2026, 6, 24),
|
|
connection="missing",
|
|
)
|
|
)
|
|
with pytest.raises(KeyError, match="Connection not found"):
|
|
kill.kill_job("abc")
|
|
|
|
|
|
# --- application_id accepted (regression: "agent passed wrong id" bug) ---
|
|
|
|
def test_kill_job_accepts_application_id(fresh_stores):
|
|
kill.conn_store.save(Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088"))
|
|
kill.store.put(
|
|
Job(
|
|
job_id="a1b2c3d4e5f6",
|
|
application_id="application_17400000001_0001",
|
|
script_path="/tmp/j.py",
|
|
queue="default",
|
|
submit_time=datetime(2026, 6, 24),
|
|
connection="prod",
|
|
yarn_rm_url="http://rm:8088",
|
|
)
|
|
)
|
|
with patch("spark_executor.tools.kill.kill_application") as m:
|
|
result = kill.kill_job("application_17400000001_0001")
|
|
# Underlying YARN call uses the application_id, not the local job_id.
|
|
assert m.call_args.args[0] == "application_17400000001_0001"
|
|
# And the response still surfaces BOTH ids so the agent can confirm.
|
|
assert result["job_id"] == "a1b2c3d4e5f6"
|
|
assert result["application_id"] == "application_17400000001_0001"
|
|
|
|
|
|
def test_kill_job_unknown_error_message_mentions_both_ids(fresh_stores):
|
|
with pytest.raises(KeyError) as ei:
|
|
kill.kill_job("totally-fake")
|
|
msg = str(ei.value)
|
|
assert "job_id" in msg
|
|
assert "application_id" in msg
|