Files
mcp-server/tests/unit/test_submit_tool.py
T
ClaudeandClaude Fable 5 35078f0ae0 refactor(submit): drop spark-submit retry loop; confirm raises on failure
confirm_submit_job no longer retries on SparkSubmitError. A failed
single attempt is recorded as-is:

  - pending.status = 'FAILED' with the error message in pending.error
  - the function re-raises the SparkSubmitError so the caller (and the
    agent) sees the failure immediately
  - the agent can call get_pending_job to see the persisted FAILED
    state, including the captured error

The 'retry' code path was hiding real failures behind transient
recovery: a spark-submit that died because the script is broken looks
the same as one that died because YARN was momentarily unreachable.
Without retries, the agent gets a clean FAILED state and a clear
exception to act on, instead of a delayed, ambiguous result.

Manual recovery is still possible: re-calling confirm_submit_job on a
FAILED pending resets it to PENDING and runs a single fresh attempt
(see test_confirm_resets_failed_then_succeeds). The reset path is
useful for 'I fixed the script, try again' flows; it just is no longer
the default on every failure.

Removed:
  - the for-attempt loop and time.sleep in submit.py
  - import time in submit.py and test_submit_tool.py
  - confirm_max_retries and confirm_retry_delay_seconds from
    common/config.py (env vars, snapshot, reload)
  - 2 retry tests (test_confirm_retries_spark_submit_failure_then_succeeds,
    test_confirm_exhausts_retries_and_sets_failed)
  - test_confirm_marks_failed_on_spark_submit_error updated to assert
    the new single-attempt + raise contract
  - test_confirm_resets_failed_and_retries renamed to
    test_confirm_resets_failed_then_succeeds (no more retry loop to
    exercise)

Full suite 244 passed (3 fewer than before, matching the removed
tests).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-30 14:29:27 +08:00

648 lines
21 KiB
Python

# coding=utf-8
from pathlib import Path
from unittest.mock import patch
import pytest
from common.config import settings
from spark_executor.core import connection_store, pending_store
from spark_executor.core.pending_store import PendingStore
from spark_executor.models import Connection
from spark_executor.tools import submit
from spark_executor.core.spark_submit import SparkSubmitError
def _spark_proc(stderr: str = "tracking URL: http://rm:8088/proxy/application_17400000001/\n"):
return type("P", (), {"returncode": 0, "stderr": stderr})()
@pytest.fixture(autouse=True)
def _fresh(tmp_path: Path, monkeypatch):
monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path))
monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore())
monkeypatch.setattr(pending_store, "DEFAULT_DATA_DIR", str(tmp_path))
monkeypatch.setattr(pending_store, "store", PendingStore())
submit.conn_store = connection_store.store
submit.pending_store = pending_store.store
connection_store.store.save(Connection(
name="prod",
master="yarn",
deploy_mode="cluster",
yarn_rm_url="http://rm:8088",
))
@pytest.fixture
def real_script(tmp_path: Path) -> Path:
"""A real .py file on disk that satisfies prepare_submit_job's path check."""
p = tmp_path / "demo.py"
p.write_text("print('hi from test')\n")
return p
def _last_pending_id() -> str:
pending = submit.pending_store.list_all()
assert pending, "no pending submission was created"
return pending[-1].pending_id
# --- prepare_submit_job ---
def test_prepare_does_not_invoke_spark_submit(monkeypatch, real_script):
called = []
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: called.append(cmd))
out = submit.prepare_submit_job(
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
assert called == [] # never called
assert out["status"] == "PENDING"
assert out["pending_id"].startswith("p_")
def test_prepare_persists_pending_with_snapshot(monkeypatch, real_script):
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
connection_store.store.save(
Connection(
name="prod",
master="yarn",
deploy_mode="cluster",
yarn_rm_url="http://rm:8088",
spark_conf={"spark.sql.shuffle.partitions": "200"},
)
)
submit.prepare_submit_job(
connection="prod",
script_path=str(real_script),
queue="research",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
p = submit.pending_store.get(_last_pending_id())
assert p is not None
assert p.connection == "prod"
assert p.master == "yarn"
assert p.deploy_mode == "cluster"
assert p.yarn_rm_url == "http://rm:8088"
assert p.spark_conf == {"spark.sql.shuffle.partitions": "200"}
assert p.queue == "research"
assert p.script_path == str(real_script)
def test_prepare_persists_app_name(monkeypatch, real_script):
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
submit.prepare_submit_job(
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="my-etl-job",
)
p = submit.pending_store.get(_last_pending_id())
assert p is not None
assert p.app_name == "my-etl-job"
def test_prepare_requires_explicit_app_name(monkeypatch, real_script):
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
with pytest.raises(TypeError, match="app_name"):
submit.prepare_submit_job(connection="prod", script_path=str(real_script))
def test_prepare_accepts_restored_defaults(monkeypatch, real_script):
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
out = submit.prepare_submit_job(
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
assert out["status"] == "PENDING"
p = submit.pending_store.get(_last_pending_id())
assert p.queue == "default"
assert p.executor_memory == "2G"
assert p.executor_cores == 2
assert p.num_executors == 2
def test_prepare_snapshots_extra_args(monkeypatch, real_script):
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
submit.prepare_submit_job(
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
extra_args={"jars": "hdfs:///lib/foo.jar"},
)
pid = _last_pending_id()
p = submit.pending_store.get(pid)
assert p is not None
assert p.extra_args == {"jars": "hdfs:///lib/foo.jar"}
fake_proc = type("P", (), {
"returncode": 0,
"stderr": "tracking URL: http://rm:8088/proxy/application_17400000002/\n",
})()
with patch("spark_executor.tools.submit.run_spark_submit", return_value=fake_proc) as m:
submit.confirm_submit_job(pending_id=pid)
cmd = m.call_args.args[0]
assert ["--jars", "hdfs:///lib/foo.jar"] in [
cmd[i : i + 2] for i in range(len(cmd) - 1)
]
def test_prepare_raises_for_unknown_connection(real_script):
with pytest.raises(KeyError, match="missing"):
submit.prepare_submit_job(
connection="missing",
script_path=str(real_script),
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
# --- script_path validation ---
def test_prepare_rejects_nonexistent_script_path(tmp_path):
missing = str(tmp_path / "does_not_exist.py")
with pytest.raises(ValueError, match="does not exist or is not a file"):
submit.prepare_submit_job(
connection="prod",
script_path=missing,
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
def test_prepare_rejects_directory_as_script_path(tmp_path):
"""A directory is not a file, even if it exists."""
with pytest.raises(ValueError, match="does not exist or is not a file"):
submit.prepare_submit_job(
connection="prod",
script_path=str(tmp_path),
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
def test_prepare_rejects_empty_script_path():
with pytest.raises(ValueError, match="does not exist or is not a file"):
submit.prepare_submit_job(
connection="prod",
script_path="",
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
def test_prepare_error_message_mentions_write_job_file(tmp_path):
"""The agent must be told to call write_job_file first."""
with pytest.raises(ValueError, match="write_job_file"):
submit.prepare_submit_job(
connection="prod",
script_path=str(tmp_path / "x.py"),
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
def test_confirm_rejects_if_script_was_deleted_after_prepare(tmp_path, monkeypatch):
"""Defense in depth: file present at prepare, gone by confirm -> 400."""
script = tmp_path / "demo.py"
script.write_text("print('hi')\n")
submit.prepare_submit_job(
connection="prod",
script_path=str(script),
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
pid = _last_pending_id()
# Simulate the file being deleted (e.g. cleanup job) between prepare and confirm
script.unlink()
with pytest.raises(ValueError, match="does not exist or is not a file"):
submit.confirm_submit_job(pending_id=pid)
def test_prepare_snapshots_connection_at_prepare_time(monkeypatch, real_script):
"""Editing the connection between prepare and confirm must NOT silently retarget."""
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
submit.prepare_submit_job(
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
p = submit.pending_store.get(_last_pending_id())
# Now mutate the connection
connection_store.store.save(
Connection(name="prod", master="spark://attacker:7077", deploy_mode="client")
)
# Snapshot is unchanged
p2 = submit.pending_store.get(p.pending_id)
assert p2.master == "yarn"
assert p2.deploy_mode == "cluster"
# --- confirm_submit_job ---
def test_confirm_invokes_spark_submit_and_marks_submitted(monkeypatch, real_script):
submit.prepare_submit_job(
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
pid = _last_pending_id()
fake_proc = type("P", (), {
"returncode": 0,
"stderr": "tracking URL: http://rm:8088/proxy/application_17400000001/\n",
})()
with patch("spark_executor.tools.submit.run_spark_submit", return_value=fake_proc) as m:
result = submit.confirm_submit_job(pending_id=pid)
cmd = m.call_args.args[0]
assert "yarn" in cmd
assert "cluster" in cmd
assert cmd[-1] == str(real_script)
assert result.application_id == "application_17400000001"
# pending updated
p = submit.pending_store.get(pid)
assert p.status == "SUBMITTED"
assert p.application_id == "application_17400000001"
assert p.job_id is not None
# job carries the connection's yarn_rm_url snapshot
assert submit.job_store.get(p.job_id).yarn_rm_url == "http://rm:8088"
def test_confirm_raises_for_unknown_pending_id():
with pytest.raises(KeyError, match="missing"):
submit.confirm_submit_job(pending_id="missing")
def test_confirm_refuses_non_pending_status(monkeypatch, real_script):
submit.prepare_submit_job(
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
pid = _last_pending_id()
# Mark it CANCELLED first
p = submit.pending_store.get(pid)
p.status = "CANCELLED"
submit.pending_store.save(p)
with pytest.raises(ValueError, match="CANCELLED"):
submit.confirm_submit_job(pending_id=pid)
def test_confirm_marks_failed_on_spark_submit_error(monkeypatch, real_script):
"""confirm_submit_job makes a single spark-submit attempt; on failure
it sets pending.status=FAILED and re-raises the SparkSubmitError."""
submit.prepare_submit_job(
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
pid = _last_pending_id()
call_count = [0]
def _raise(_cmd):
call_count[0] += 1
raise SparkSubmitError("boom")
monkeypatch.setattr(submit, "run_spark_submit", _raise)
with pytest.raises(SparkSubmitError, match="boom"):
submit.confirm_submit_job(pending_id=pid)
assert call_count[0] == 1 # single attempt, no retry
p = submit.pending_store.get(pid)
assert p.status == "FAILED"
assert "boom" in (p.error or "")
# --- confirm idempotency / manual retry ---
def test_confirm_idempotent_when_submitted(monkeypatch, real_script):
submit.prepare_submit_job(
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
pid = _last_pending_id()
monkeypatch.setattr(submit, "run_spark_submit", lambda _cmd: _spark_proc())
first = submit.confirm_submit_job(pending_id=pid)
call_count = [0]
def _counting_run(cmd):
call_count[0] += 1
return _spark_proc()
monkeypatch.setattr(submit, "run_spark_submit", _counting_run)
second = submit.confirm_submit_job(pending_id=pid)
assert call_count[0] == 0
assert second.job_id == first.job_id
assert second.application_id == first.application_id
assert second.tracking_url == first.tracking_url
def test_confirm_resets_failed_then_succeeds(monkeypatch, real_script):
"""A FAILED pending can be re-confirmed after the user fixes the
underlying issue. confirm_submit_job resets FAILED -> PENDING,
then makes a single fresh attempt (no automatic retry loop)."""
submit.prepare_submit_job(
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
pid = _last_pending_id()
monkeypatch.setattr(
submit,
"run_spark_submit",
lambda _cmd: (_ for _ in ()).throw(SparkSubmitError("boom")),
)
with pytest.raises(SparkSubmitError, match="boom"):
submit.confirm_submit_job(pending_id=pid)
p = submit.pending_store.get(pid)
assert p.status == "FAILED"
assert "boom" in (p.error or "")
call_count = [0]
def _counting_run(cmd):
call_count[0] += 1
return _spark_proc()
monkeypatch.setattr(submit, "run_spark_submit", _counting_run)
result = submit.confirm_submit_job(pending_id=pid)
assert result.application_id == "application_17400000001"
assert call_count[0] == 1
p = submit.pending_store.get(pid)
assert p.status == "SUBMITTED"
assert p.error is None
def test_confirm_updates_pending_before_job_store(monkeypatch, real_script):
submit.prepare_submit_job(
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
pid = _last_pending_id()
monkeypatch.setattr(submit, "run_spark_submit", lambda _cmd: _spark_proc())
def _exploding_put(job):
raise RuntimeError("job_store unavailable")
monkeypatch.setattr(submit.job_store, "put", _exploding_put)
with pytest.raises(RuntimeError, match="job_store unavailable"):
submit.confirm_submit_job(pending_id=pid)
p = submit.pending_store.get(pid)
assert p.status == "SUBMITTED"
assert p.application_id == "application_17400000001"
assert p.tracking_url is not None
# --- list_pending_jobs ---
def test_list_pending_jobs_empty():
assert submit.list_pending_jobs() == []
def test_list_pending_jobs_returns_all(monkeypatch, tmp_path):
a = tmp_path / "a.py"
b = tmp_path / "b.py"
a.write_text("a\n"); b.write_text("b\n")
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
submit.prepare_submit_job(
connection="prod",
script_path=str(a),
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
submit.prepare_submit_job(
connection="prod",
script_path=str(b),
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
out = submit.list_pending_jobs()
assert {p["script_path"] for p in out} == {str(a), str(b)}
assert all(p["status"] == "PENDING" for p in out)
# --- get_pending_job ---
def test_get_pending_job_returns_dump(monkeypatch, real_script):
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
submit.prepare_submit_job(
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
pid = _last_pending_id()
out = submit.get_pending_job(pid)
assert out["pending_id"] == pid
assert out["connection"] == "prod"
assert out["status"] == "PENDING"
def test_get_pending_job_unknown_raises():
with pytest.raises(KeyError):
submit.get_pending_job("missing")
# --- cancel_pending_job ---
def test_cancel_pending_job_marks_cancelled(monkeypatch, real_script):
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
submit.prepare_submit_job(
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
pid = _last_pending_id()
out = submit.cancel_pending_job(pid)
assert out == {"pending_id": pid, "status": "CANCELLED"}
assert submit.pending_store.get(pid).status == "CANCELLED"
def test_cancel_pending_job_unknown_raises():
with pytest.raises(KeyError):
submit.cancel_pending_job("missing")
def test_cancel_pending_job_refuses_submitted(monkeypatch, real_script):
submit.prepare_submit_job(
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
pid = _last_pending_id()
p = submit.pending_store.get(pid)
p.status = "SUBMITTED"
submit.pending_store.save(p)
with pytest.raises(ValueError, match="SUBMITTED"):
submit.cancel_pending_job(pid)
# --- update_pending_job ---
def test_update_pending_updates_resource_params(monkeypatch, real_script):
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
submit.prepare_submit_job(
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
pid = _last_pending_id()
out = submit.update_pending_job(
pending_id=pid,
queue="research",
executor_memory="8G",
executor_cores=4,
num_executors=8,
app_name="renamed-app",
extra_args={"jars": "hdfs:///lib/bar.jar"},
)
assert out["pending_id"] == pid
assert out["status"] == "PENDING"
params = out["parameters"]
assert params["queue"] == "research"
assert params["executor_memory"] == "8G"
assert params["executor_cores"] == 4
assert params["num_executors"] == 8
assert params["app_name"] == "renamed-app"
assert params["extra_args"] == {"jars": "hdfs:///lib/bar.jar"}
# unchanged fields are preserved
assert params["connection"] == "prod"
assert params["script_path"] == str(real_script)
def test_update_pending_rejects_non_pending_status(monkeypatch, real_script):
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
submit.prepare_submit_job(
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
pid = _last_pending_id()
p = submit.pending_store.get(pid)
p.status = "SUBMITTED"
submit.pending_store.save(p)
with pytest.raises(ValueError, match="only PENDING submissions can be updated"):
submit.update_pending_job(pending_id=pid, queue="research")
def test_update_pending_rejects_bad_script_path(tmp_path, real_script):
submit.prepare_submit_job(
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
pid = _last_pending_id()
missing = str(tmp_path / "does_not_exist.py")
with pytest.raises(ValueError, match="does not exist or is not a file"):
submit.update_pending_job(pending_id=pid, script_path=missing)
def test_update_pending_rejects_sql_violation_script(tmp_path, real_script):
submit.prepare_submit_job(
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
pid = _last_pending_id()
bad_script = tmp_path / "evil.py"
bad_script.write_text('spark.sql("DROP TABLE users")\n')
with pytest.raises(ValueError, match="DROP"):
submit.update_pending_job(pending_id=pid, script_path=str(bad_script))
def test_update_pending_unknown_id_raises():
with pytest.raises(KeyError, match="missing"):
submit.update_pending_job(pending_id="missing", queue="research")