ClaudeandClaude Fable 5
523e6a9c76
The MCP tool was named `generate_job_file` from Stage 2 but it does
NOT generate PySpark code — the calling LLM writes the code in its own
context, and this tool only persists it to a file under
SPARK_EXECUTOR_JOBS_DIR so `spark-submit` can see it. The misleading
`generate_` prefix sent agents (and humans) looking for a code
generator that doesn't exist.
This commit folds three related polish changes into one (split later
with rebase -i if you want them as separate history):
1. The rename itself:
- `tools/generate.py` → `tools/write_job.py`
- `generate_job_file` → `write_job_file`
- `GenerateJobFileRequest` → `WriteJobFileRequest`
- `/generate_job_file` route → `/write_job_file`
- `operation_id="generate_job_file"` → `operation_id="write_job_file"`
The internal helper `core.job_writer.write_job_file` (which just
writes bytes to disk with no SQL guard) is imported with an
`_write_to_disk` alias to avoid the name collision with the
MCP-exposed function in the same module.
The description for the tool now explicitly states 'this tool
does NOT generate PySpark code. The calling LLM is expected to
have already written the code; this tool only persists it.'
2. Skill for LLM agents operating the service
(`docs/superpowers/skills/spark-executor-mcp-operate/SKILL.md`,
449 lines). Covers the 16 tools, the two-step prepare/confirm
flow, the dual-ID contract (job_id vs application_id), the
PendingSubmission state machine, the Connection profile, the
job-file workflow, the error reference, common pitfalls, and a
full end-to-end word-count example.
3. Default `executor_memory` lowered 4G → 2G
(`_DEFAULTS_TO_CONFIRM` in `server.py`). Mirrors the matching
change in `test_mcp_routes.py` and the 5 unit tests that
reference the default. Aligns with the lighter workloads the
service is sized for in its current container profile.
Also tracked in git for the first time:
- `docs/superpowers/plans/2026-06-24-spark-executor-mcp.md`
(the original Stage 1/2/3 design plan, updated to use the new
tool name throughout).
Test rename:
- `tests/unit/test_generate_tool.py` → `test_write_job_tool.py`
- the new test file picks up an extra assertion that the SQL guard
rejects a `DROP TABLE` statement at write time.
243 tests pass (was 242; +1 new SQL-guard assertion). Zero regressions.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
691 lines
22 KiB
Python
691 lines
22 KiB
Python
# coding=utf-8
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
import time
|
|
|
|
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):
|
|
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()
|
|
def _raise(_cmd):
|
|
raise SparkSubmitError("boom")
|
|
monkeypatch.setattr(submit, "run_spark_submit", _raise)
|
|
settings.confirm_max_retries = 0
|
|
with pytest.raises(SparkSubmitError):
|
|
submit.confirm_submit_job(pending_id=pid)
|
|
p = submit.pending_store.get(pid)
|
|
assert p.status == "FAILED"
|
|
assert "boom" in (p.error or "")
|
|
|
|
|
|
# --- confirm retry / idempotency ---
|
|
|
|
def test_confirm_retries_spark_submit_failure_then_succeeds(monkeypatch, real_script):
|
|
monkeypatch.setattr(time, "sleep", lambda _s: None)
|
|
settings.confirm_max_retries = 2
|
|
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()
|
|
calls = []
|
|
def _maybe_fail(cmd):
|
|
calls.append(cmd)
|
|
if len(calls) < 3:
|
|
raise SparkSubmitError("transient")
|
|
return _spark_proc()
|
|
monkeypatch.setattr(submit, "run_spark_submit", _maybe_fail)
|
|
result = submit.confirm_submit_job(pending_id=pid)
|
|
assert result.application_id == "application_17400000001"
|
|
assert len(calls) == 3
|
|
p = submit.pending_store.get(pid)
|
|
assert p.status == "SUBMITTED"
|
|
|
|
|
|
def test_confirm_exhausts_retries_and_sets_failed(monkeypatch, real_script):
|
|
monkeypatch.setattr(time, "sleep", lambda _s: None)
|
|
settings.confirm_max_retries = 2
|
|
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()
|
|
calls = []
|
|
def _always_fail(cmd):
|
|
calls.append(cmd)
|
|
raise SparkSubmitError("persistent")
|
|
monkeypatch.setattr(submit, "run_spark_submit", _always_fail)
|
|
with pytest.raises(SparkSubmitError, match="persistent"):
|
|
submit.confirm_submit_job(pending_id=pid)
|
|
assert len(calls) == 3
|
|
p = submit.pending_store.get(pid)
|
|
assert p.status == "FAILED"
|
|
assert "persistent" in (p.error or "")
|
|
|
|
|
|
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_and_retries(monkeypatch, real_script):
|
|
monkeypatch.setattr(time, "sleep", lambda _s: 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()
|
|
monkeypatch.setattr(
|
|
submit,
|
|
"run_spark_submit",
|
|
lambda _cmd: (_ for _ in ()).throw(SparkSubmitError("boom")),
|
|
)
|
|
with pytest.raises(SparkSubmitError):
|
|
submit.confirm_submit_job(pending_id=pid)
|
|
p = submit.pending_store.get(pid)
|
|
assert p.status == "FAILED"
|
|
assert "boom" in (p.error or "")
|
|
|
|
monkeypatch.setattr(submit, "run_spark_submit", lambda _cmd: _spark_proc())
|
|
result = submit.confirm_submit_job(pending_id=pid)
|
|
assert result.application_id == "application_17400000001"
|
|
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")
|