Files
mcp-server/tests/unit/test_status_tool.py
T
ClaudeandClaude Fable 5 0fef77c0b8 fix(job_store): persist Jobs to disk and accept either ID in job tools
Two user-reported bugs, same root cause: the in-memory JobStore + the
'job_id must be the 12-char hex' tool contract.

Bug 1: 'Unknown job_id' reported frequently
  JobStore was a process-local dict (spark_executor/core/job_store.py).
  Under gunicorn workers > 1, a job created by confirm_submit_job
  landing on worker A was invisible to worker B, so a follow-up
  get_job_status / get_job_result / get_job_logs / kill_job landing on
  a different worker returned 'Unknown job_id'. Same multi-worker
  problem that bit the MCP session layer; only the affected data was
  different.

Bug 2: 'get_job_logs frequently confuses job_id and application_id'
  confirm_submit_job returns BOTH identifiers in SubmitResult, but
  get_job_logs (and friends) only accepted the local 12-char job_id
  and never said so in their description. When the agent passed the
  YARN application_id, the error message itself was misleading:
  'Unknown job_id: application_17400000001_0001' — the agent had
  passed an id, just the wrong kind.

This change fixes both at the root:

  * JobStore is now JSON-backed at data/jobs.json (atomic tempfile +
    os.replace), with cross-process safety via fcntl.flock on a sibling
    .lock file. Stage 3's SQLite migration is still planned; the file
    format is intentionally simple so it is a straight
    'for j in read_all(): db.insert(j)'.

  * New JobStore.get_either(uid) looks up by job_id first, then
    application_id. All four job-lifecycle tools (get_job_status,
    get_job_result, get_job_logs, kill_job) call get_either instead
    of get(job_id), so the agent can pass either identifier and get
    the same answer.

  * The 'neither matched' KeyError now spells out both id forms and
    what they look like, so the agent isn't left guessing.

  * server.py tool descriptions for the four job tools explicitly
    state 'job_id accepts BOTH identifiers' so this is visible to the
    LLM at tool-selection time, not only at error time.

Tests:
  * test_job_store.py: tmp_path isolation, persistence across
    instances, human-readable JSON, corrupt-file resilience,
    get_either (by job_id, by application_id, collision preference,
    unknown), put idempotency.
  * test_{logs,status,kill,result}_tool.py: per-test tmp_path fixture,
    'accepts application_id' regression for each tool, and an
    explicit assertion that the unknown-id error message mentions
    BOTH id forms. test_result_raises_keyerror_for_unknown_job's
    match pattern updated for the new message.

242 tests pass (was 226; +16 new). Zero regressions.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-29 18:57:31 +08:00

150 lines
4.6 KiB
Python

# coding=utf-8
from datetime import datetime
from pathlib import Path
from unittest.mock import patch
import pytest
from spark_executor.core.job_store import JobStore
from spark_executor.core import connection_store
from spark_executor.models import Connection, Job
from spark_executor.tools import connections, status
from spark_executor.core.yarn_client import YarnClientConfig
@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
status.conn_store = store
status.store = JobStore(data_dir=str(tmp_path))
def _fresh_stores():
"""Reset job store and connection store singletons for a single test."""
store = connection_store.ConnectionStore()
connection_store.store = store
connections.store = store
status.conn_store = store
status.store = JobStore()
def test_get_job_status_returns_state():
_fresh_stores()
status.conn_store.save(Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088"))
status.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.status.get_application_status",
return_value=("RUNNING", "State : RUNNING\n"),
) as m:
out = status.get_job_status("abc")
assert out.application_id == "application_1"
assert out.state == "RUNNING"
assert "RUNNING" in out.raw
# resolved YarnClientConfig is forwarded to the REST client
args = m.call_args.args
assert args[0] == "application_1"
assert args[1].yarn_rm_url == "http://rm:8088"
def test_get_job_status_raises_for_unknown_job():
_fresh_stores()
with pytest.raises(KeyError):
status.get_job_status("missing")
def test_get_job_status_raises_when_connection_missing():
_fresh_stores()
status.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"):
status.get_job_status("abc")
def test_get_job_status_passes_auth_config():
_fresh_stores()
status.conn_store.save(
Connection(
name="secure",
master="yarn",
yarn_rm_url="http://rm:8088",
auth_type="basic",
auth_user="hdfs",
auth_password="secret",
)
)
status.store.put(
Job(
job_id="abc",
application_id="application_1",
script_path="/tmp/j.py",
queue="default",
submit_time=datetime(2026, 6, 24),
connection="secure",
yarn_rm_url="http://rm:8088",
)
)
with patch(
"spark_executor.tools.status.get_application_status",
return_value=("RUNNING", "State : RUNNING\n"),
) as m:
status.get_job_status("abc")
args = m.call_args.args
assert args[0] == "application_1"
assert isinstance(args[1], YarnClientConfig)
assert args[1].auth_type == "basic"
assert args[1].auth_user == "hdfs"
# --- application_id accepted (regression: "agent passed wrong id" bug) ---
def test_get_job_status_accepts_application_id(fresh_stores):
status.conn_store.save(Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088"))
status.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.status.get_application_status",
return_value=("RUNNING", "{}"),
):
out = status.get_job_status("application_17400000001_0001")
assert out.state == "RUNNING"
assert out.application_id == "application_17400000001_0001"
def test_get_job_status_unknown_error_message_mentions_both_ids(fresh_stores):
with pytest.raises(KeyError) as ei:
status.get_job_status("totally-fake")
msg = str(ei.value)
assert "job_id" in msg
assert "application_id" in msg