Audit of the previous 'salvage application_id from stderr' change
(e4fb97c) revealed the fix was incomplete:
- The pending record now had pending.application_id set on failure.
- But get_job_logs / get_job_status / kill_job all look up the
application_id through JobStore.get_either(application_id), which
scans the on-disk jobs.json for a Job record matching that
application_id. No Job record was being created on the failure
path, so the recovered application_id was unreachable through
the public tool surface — get_job_logs would 404 with 'No Job
found for id=...application_xxx'.
Fix: when confirm_submit_job's failure branch successfully extracts
an application_id from stderr, it now also creates a Job record (same
shape as the success path: 12-char hex job_id, application_id,
script_path, queue, submit_time, connection, yarn_rm_url). The Job's
only job in this case is to act as the lookup index from the
LLM-facing tools into the underlying YARN app.
The job_id is generated even when application_id recovery fails (so
we always set pending.job_id), but the Job record is only put when
application_id is non-None. That keeps the 'spark-submit died locally
before reaching YARN' case clean: no Job record, get_job_logs gets
the standard 'no logs available' error from yarn_client, not a
404 from JobStore.
Tests:
- test_confirm_recovers_application_id_from_failed_spark_submit now
asserts both that pending.application_id is set AND that a Job
record exists and is findable by both job_id AND application_id.
- test_confirm_failed_spark_submit_without_application_id_keeps_it_none
asserts the negative case: no application_id, no Job record.
Fixture fix: _fresh() now rebinds submit.job_store to use tmp_path.
Previously the test wrote to ./data/jobs.json (the real data dir),
which left stale records between test runs and broke any test that
scanned by application_id. This is a latent bug exposed by the new
assertions; fixing it here also makes every other confirm test
hermetic.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
744 lines
25 KiB
Python
744 lines
25 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.job_store import JobStore
|
|
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.job_store is a module-level singleton bound at import time
|
|
# to settings.data_dir (./data/). Rebind to tmp_path so failed-submit
|
|
# tests don't pollute the real data/jobs.json with old records.
|
|
submit.job_store = JobStore(data_dir=str(tmp_path))
|
|
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 "")
|
|
assert p.application_id is None # no process attached, nothing to recover
|
|
|
|
|
|
def test_confirm_recovers_application_id_from_failed_spark_submit(monkeypatch, real_script):
|
|
"""When spark-submit fails AFTER YARN accepted the application (a
|
|
common case: spark-submit loses its connection to the RM, or the
|
|
PySpark script exits non-zero on the cluster side), the stderr
|
|
still contains 'Submitted application <id>'. We must salvage it
|
|
AND create a Job record so get_job_logs / get_job_status /
|
|
kill_job can find the application_id through JobStore."""
|
|
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()
|
|
failed_stderr = (
|
|
"Warning: ...\n"
|
|
"Submitted application application_17400000002_0007\n"
|
|
"tracking URL: http://rm:8088/proxy/application_17400000002_0007/\n"
|
|
"... then spark-submit died for some reason\n"
|
|
)
|
|
failed_proc = type("P", (), {"returncode": 1, "stderr": failed_stderr})()
|
|
def _raise_with_result(_cmd):
|
|
err = SparkSubmitError("lost connection to RM")
|
|
err.result = failed_proc
|
|
raise err
|
|
monkeypatch.setattr(submit, "run_spark_submit", _raise_with_result)
|
|
with pytest.raises(SparkSubmitError, match="lost connection to RM"):
|
|
submit.confirm_submit_job(pending_id=pid)
|
|
p = submit.pending_store.get(pid)
|
|
assert p.status == "FAILED"
|
|
assert "lost connection to RM" in (p.error or "")
|
|
assert p.application_id == "application_17400000002_0007"
|
|
assert p.tracking_url == "http://rm:8088/proxy/application_17400000002_0007/"
|
|
|
|
# The Job record must exist so get_job_logs / get_job_status /
|
|
# kill_job can find the application_id through JobStore. The
|
|
# agent only sees the exception; the persistent log target lives
|
|
# in JobStore.
|
|
assert p.job_id is not None
|
|
job = submit.job_store.get_either(p.job_id)
|
|
assert job is not None
|
|
assert job.application_id == "application_17400000002_0007"
|
|
# And the application_id lookup must also work (this is the path
|
|
# an LLM agent typically takes — they have the YARN id, not the
|
|
# local job_id).
|
|
job_by_app = submit.job_store.get_either("application_17400000002_0007")
|
|
assert job_by_app is not None
|
|
assert job_by_app.job_id == p.job_id
|
|
|
|
|
|
def test_confirm_failed_spark_submit_without_application_id_keeps_it_none(monkeypatch, real_script):
|
|
"""spark-submit can fail BEFORE reaching YARN (e.g. local script
|
|
syntax error, missing spark-submit binary). stderr has no
|
|
application_id — the pending must record FAILED with no log
|
|
target, so get_application_logs will return the clean 'no logs'
|
|
error rather than a confusing 404."""
|
|
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()
|
|
failed_proc = type("P", (), {
|
|
"returncode": 2,
|
|
"stderr": "Error: script_path does not exist: /tmp/whatever.py\n",
|
|
})()
|
|
def _raise_with_result(_cmd):
|
|
err = SparkSubmitError("local failure")
|
|
err.result = failed_proc
|
|
raise err
|
|
monkeypatch.setattr(submit, "run_spark_submit", _raise_with_result)
|
|
with pytest.raises(SparkSubmitError, match="local failure"):
|
|
submit.confirm_submit_job(pending_id=pid)
|
|
p = submit.pending_store.get(pid)
|
|
assert p.status == "FAILED"
|
|
assert p.application_id is None
|
|
assert p.tracking_url is None
|
|
# No application_id → no Job record. get_either for the pending's
|
|
# job_id must return None.
|
|
if p.job_id is not None:
|
|
assert submit.job_store.get_either(p.job_id) is None
|
|
|
|
|
|
# --- 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")
|