fix: validate script_path exists + tell agent to call generate_job_file
Two problems with the prior prepare_submit_job flow:
1. The agent could pass a script_path that only existed in its own
context (LLM-generated code not yet on disk) or a path on the host
filesystem that's invisible inside the container. The new check
surfaces this as a 400 with a clear remediation hint instead of
letting spark-submit fail later with an opaque FileNotFoundError -> 500.
2. The container-isolation issue: any path the agent gives is interpreted
inside the container. The two ways a file can legitimately exist there
are (a) generate_job_file(code=...) just wrote it to
SPARK_EXECUTOR_JOBS_DIR (the default ./data/jobs/ is the only
gitignored dir that survives restarts), or (b) a host dir was
mounted via -v. The error message spells both out so an agent can
self-correct.
Implementation:
- submit.py: new _check_script_path() that raises ValueError (-> 400)
when the path is missing, empty, or a directory. Called in both
prepare_submit_job and confirm_submit_job (defense in depth).
- prepare_submit_job checks the connection FIRST (KeyError -> 404)
before the script (ValueError -> 400), so an agent with both problems
sees the more fundamental 'unknown connection' error first.
- server.py / requests.py: route description and Pydantic field
description spell out the generate_job_file pattern so an LLM
reading the tool schema learns the right next step.
8 new tests; 8 existing tests adjusted to create real files (they used
synthetic /tmp/*.py paths that don't exist).
This commit is contained in:
@@ -73,7 +73,13 @@ def health_check():
|
|||||||
"Snapshot the named Connection's master / deploy_mode / spark_conf / "
|
"Snapshot the named Connection's master / deploy_mode / spark_conf / "
|
||||||
"yarn_rm_url into a PendingSubmission record and persist it. "
|
"yarn_rm_url into a PendingSubmission record and persist it. "
|
||||||
"Does NOT invoke spark-submit. Returns pending_id for use with "
|
"Does NOT invoke spark-submit. Returns pending_id for use with "
|
||||||
"confirm_submit_job (the user-second-confirmation step)."
|
"confirm_submit_job (the user-second-confirmation step).\n\n"
|
||||||
|
"REQUIRED PATTERN for LLM-generated code: call generate_job_file(code=...) "
|
||||||
|
"first, then pass the returned script_path here. Direct submission with a "
|
||||||
|
"synthetic path (one that only exists in the agent's context) will be "
|
||||||
|
"rejected with HTTP 400 — the script must exist inside the container's "
|
||||||
|
"filesystem. For pre-existing files, mount the host directory into the "
|
||||||
|
"container and pass the in-container path."
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
def _prepare_submit_job(req: PrepareSubmitJobRequest):
|
def _prepare_submit_job(req: PrepareSubmitJobRequest):
|
||||||
|
|||||||
@@ -26,7 +26,18 @@ class SaveConnectionRequest(BaseModel):
|
|||||||
|
|
||||||
class PrepareSubmitJobRequest(BaseModel):
|
class PrepareSubmitJobRequest(BaseModel):
|
||||||
connection: str
|
connection: str
|
||||||
script_path: str
|
script_path: str = Field(
|
||||||
|
...,
|
||||||
|
description=(
|
||||||
|
"Absolute path to the PySpark script inside the container's "
|
||||||
|
"filesystem. Must point at an existing regular file. For "
|
||||||
|
"LLM-generated code, call generate_job_file(code=...) first "
|
||||||
|
"and pass the returned script_path here. For pre-existing "
|
||||||
|
"files on the host, mount them via a docker volume and pass "
|
||||||
|
"the in-container path. Returns 400 with a remediation hint "
|
||||||
|
"if the path is missing or not a file."
|
||||||
|
),
|
||||||
|
)
|
||||||
queue: str = "default"
|
queue: str = "default"
|
||||||
executor_memory: str = "4G"
|
executor_memory: str = "4G"
|
||||||
executor_cores: int = 2
|
executor_cores: int = 2
|
||||||
|
|||||||
@@ -3,6 +3,7 @@
|
|||||||
@Time :2026/6/24
|
@Time :2026/6/24
|
||||||
@Author :tao.chen
|
@Author :tao.chen
|
||||||
"""
|
"""
|
||||||
|
import os
|
||||||
import secrets
|
import secrets
|
||||||
import uuid
|
import uuid
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
@@ -20,6 +21,27 @@ from spark_executor.core.spark_submit import (
|
|||||||
from spark_executor.models import Job, PendingSubmission, SubmitResult
|
from spark_executor.models import Job, PendingSubmission, SubmitResult
|
||||||
|
|
||||||
|
|
||||||
|
def _script_path_error(path: str) -> ValueError:
|
||||||
|
"""Instructive error for an invalid script_path. The MCP client receives
|
||||||
|
this via the ValueError -> 400 exception handler in server.py."""
|
||||||
|
return ValueError(
|
||||||
|
f"script_path does not exist or is not a file: {path!r}. "
|
||||||
|
f"Two ways to fix this:\n"
|
||||||
|
f" 1. (Recommended) Call generate_job_file(code=...) first to write "
|
||||||
|
f"the PySpark source to disk, then pass the returned script_path.\n"
|
||||||
|
f" 2. If the file already exists on the host, mount it into the "
|
||||||
|
f"container (e.g. -v /host/path:/app/scripts:ro in docker run) and "
|
||||||
|
f"pass the in-container path here."
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _check_script_path(script_path: str) -> None:
|
||||||
|
"""Verify script_path points at an existing regular file. Raises ValueError
|
||||||
|
(-> 400 via the FastAPI exception handler) with an instructive message."""
|
||||||
|
if not script_path or not os.path.isfile(script_path):
|
||||||
|
raise _script_path_error(script_path)
|
||||||
|
|
||||||
|
|
||||||
def _new_pending_id() -> str:
|
def _new_pending_id() -> str:
|
||||||
return "p_" + secrets.token_hex(6)
|
return "p_" + secrets.token_hex(6)
|
||||||
|
|
||||||
@@ -39,9 +61,23 @@ def prepare_submit_job(
|
|||||||
f"queue={queue} executor_memory={executor_memory} executor_cores={executor_cores} "
|
f"queue={queue} executor_memory={executor_memory} executor_cores={executor_cores} "
|
||||||
f"num_executors={num_executors}"
|
f"num_executors={num_executors}"
|
||||||
)
|
)
|
||||||
|
# Order of checks matters for the error the agent sees:
|
||||||
|
# 1. Unknown connection -> 404 (KeyError -> 404 in server.py)
|
||||||
|
# 2. Missing script file -> 400 (ValueError -> 400)
|
||||||
|
# Connection is checked first because it is a more fundamental problem
|
||||||
|
# (the agent is asking about a cluster that doesn't exist), and the
|
||||||
|
# agent shouldn't have to fix the script path only to learn the
|
||||||
|
# connection name is wrong.
|
||||||
conn = conn_store.get(connection)
|
conn = conn_store.get(connection)
|
||||||
if conn is None:
|
if conn is None:
|
||||||
raise KeyError(f"Unknown connection: {connection}")
|
raise KeyError(f"Unknown connection: {connection}")
|
||||||
|
# Fail fast: a non-existent path is the most common agent mistake (it
|
||||||
|
# generated the code in its own context but forgot to call
|
||||||
|
# generate_job_file first, or its path refers to the host filesystem
|
||||||
|
# which is invisible inside the container). Better to surface this with
|
||||||
|
# a 400 + clear remediation than to let spark-submit fail later with
|
||||||
|
# an opaque FileNotFoundError -> 500.
|
||||||
|
_check_script_path(script_path)
|
||||||
|
|
||||||
pending_id = _new_pending_id()
|
pending_id = _new_pending_id()
|
||||||
pending = PendingSubmission(
|
pending = PendingSubmission(
|
||||||
@@ -85,6 +121,10 @@ def confirm_submit_job(*, pending_id: str) -> SubmitResult:
|
|||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"pending_id {pending_id} is in status {pending.status!r}, not PENDING"
|
f"pending_id {pending_id} is in status {pending.status!r}, not PENDING"
|
||||||
)
|
)
|
||||||
|
# Defense in depth: re-verify the script still exists. A user could
|
||||||
|
# delete the file between prepare and confirm (or an external cleanup
|
||||||
|
# job could remove it). 400 via the ValueError -> 400 handler.
|
||||||
|
_check_script_path(pending.script_path)
|
||||||
|
|
||||||
cmd = build_spark_submit_command(
|
cmd = build_spark_submit_command(
|
||||||
master=pending.master,
|
master=pending.master,
|
||||||
|
|||||||
@@ -76,12 +76,15 @@ def test_save_connection_accepts_dict_spark_conf_in_body():
|
|||||||
assert r.json() == {"name": "prod", "status": "SAVED"}
|
assert r.json() == {"name": "prod", "status": "SAVED"}
|
||||||
|
|
||||||
|
|
||||||
def test_prepare_submit_job_works_via_body():
|
def test_prepare_submit_job_works_via_body(tmp_path):
|
||||||
c = TestClient(app)
|
c = TestClient(app)
|
||||||
c.post("/save_connection", json={"name": "prod", "master": "yarn"})
|
c.post("/save_connection", json={"name": "prod", "master": "yarn"})
|
||||||
|
# The script must exist (Stage 2 cleanup) — write a real file first.
|
||||||
|
script = tmp_path / "demo.py"
|
||||||
|
script.write_text("print('hi')\n")
|
||||||
r = c.post(
|
r = c.post(
|
||||||
"/prepare_submit_job",
|
"/prepare_submit_job",
|
||||||
json={"connection": "prod", "script_path": "/tmp/demo.py", "queue": "research"},
|
json={"connection": "prod", "script_path": str(script), "queue": "research"},
|
||||||
)
|
)
|
||||||
assert r.status_code == 200, r.text
|
assert r.status_code == 200, r.text
|
||||||
body = r.json()
|
body = r.json()
|
||||||
@@ -91,12 +94,29 @@ def test_prepare_submit_job_works_via_body():
|
|||||||
assert body["parameters"]["master"] == "yarn" # snapshotted from connection
|
assert body["parameters"]["master"] == "yarn" # snapshotted from connection
|
||||||
|
|
||||||
|
|
||||||
def test_list_and_get_pending_job_roundtrip_via_body():
|
def test_prepare_rejects_nonexistent_script_path_with_400():
|
||||||
|
"""MCP clients should see a clean 400 with a remediation hint, not a 500."""
|
||||||
c = TestClient(app)
|
c = TestClient(app)
|
||||||
c.post("/save_connection", json={"name": "prod", "master": "yarn"})
|
c.post("/save_connection", json={"name": "prod", "master": "yarn"})
|
||||||
|
r = c.post(
|
||||||
|
"/prepare_submit_job",
|
||||||
|
json={"connection": "prod", "script_path": "/nope/does_not_exist.py"},
|
||||||
|
)
|
||||||
|
assert r.status_code == 400
|
||||||
|
detail = r.json()["detail"]
|
||||||
|
# The message must guide the agent to the right next step
|
||||||
|
assert "does not exist" in detail
|
||||||
|
assert "generate_job_file" in detail
|
||||||
|
|
||||||
|
|
||||||
|
def test_list_and_get_pending_job_roundtrip_via_body(tmp_path):
|
||||||
|
c = TestClient(app)
|
||||||
|
c.post("/save_connection", json={"name": "prod", "master": "yarn"})
|
||||||
|
script = tmp_path / "a.py"
|
||||||
|
script.write_text("print('hi')\n")
|
||||||
prep = c.post(
|
prep = c.post(
|
||||||
"/prepare_submit_job",
|
"/prepare_submit_job",
|
||||||
json={"connection": "prod", "script_path": "/tmp/a.py"},
|
json={"connection": "prod", "script_path": str(script)},
|
||||||
).json()
|
).json()
|
||||||
pid = prep["pending_id"]
|
pid = prep["pending_id"]
|
||||||
|
|
||||||
@@ -104,7 +124,7 @@ def test_list_and_get_pending_job_roundtrip_via_body():
|
|||||||
assert any(p["pending_id"] == pid for p in listed)
|
assert any(p["pending_id"] == pid for p in listed)
|
||||||
|
|
||||||
got = c.post("/get_pending_job", json={"pending_id": pid}).json()
|
got = c.post("/get_pending_job", json={"pending_id": pid}).json()
|
||||||
assert got["script_path"] == "/tmp/a.py"
|
assert got["script_path"] == str(script)
|
||||||
assert got["status"] == "PENDING"
|
assert got["status"] == "PENDING"
|
||||||
|
|
||||||
|
|
||||||
@@ -164,13 +184,15 @@ def test_unknown_connection_in_prepare_returns_404():
|
|||||||
assert "nope" in r.json()["detail"]
|
assert "nope" in r.json()["detail"]
|
||||||
|
|
||||||
|
|
||||||
def test_confirm_non_pending_returns_400():
|
def test_confirm_non_pending_returns_400(tmp_path):
|
||||||
"""A CANCELLED pending should refuse confirm; FastAPI should surface ValueError as 400."""
|
"""A CANCELLED pending should refuse confirm; FastAPI should surface ValueError as 400."""
|
||||||
c = TestClient(app)
|
c = TestClient(app)
|
||||||
c.post("/save_connection", json={"name": "prod", "master": "yarn"})
|
c.post("/save_connection", json={"name": "prod", "master": "yarn"})
|
||||||
|
script = tmp_path / "x.py"
|
||||||
|
script.write_text("print('hi')\n")
|
||||||
prep = c.post(
|
prep = c.post(
|
||||||
"/prepare_submit_job",
|
"/prepare_submit_job",
|
||||||
json={"connection": "prod", "script_path": "/tmp/x.py"},
|
json={"connection": "prod", "script_path": str(script)},
|
||||||
).json()
|
).json()
|
||||||
pid = prep["pending_id"]
|
pid = prep["pending_id"]
|
||||||
c.post("/cancel_pending_job", json={"pending_id": pid})
|
c.post("/cancel_pending_job", json={"pending_id": pid})
|
||||||
@@ -179,13 +201,15 @@ def test_confirm_non_pending_returns_400():
|
|||||||
assert "CANCELLED" in r.json()["detail"]
|
assert "CANCELLED" in r.json()["detail"]
|
||||||
|
|
||||||
|
|
||||||
def test_cancel_already_submitted_returns_400():
|
def test_cancel_already_submitted_returns_400(tmp_path):
|
||||||
"""Cancelling a SUBMITTED pending should refuse with ValueError -> 400."""
|
"""Cancelling a SUBMITTED pending should refuse with ValueError -> 400."""
|
||||||
c = TestClient(app)
|
c = TestClient(app)
|
||||||
c.post("/save_connection", json={"name": "prod", "master": "yarn"})
|
c.post("/save_connection", json={"name": "prod", "master": "yarn"})
|
||||||
|
script = tmp_path / "x.py"
|
||||||
|
script.write_text("print('hi')\n")
|
||||||
prep = c.post(
|
prep = c.post(
|
||||||
"/prepare_submit_job",
|
"/prepare_submit_job",
|
||||||
json={"connection": "prod", "script_path": "/tmp/x.py"},
|
json={"connection": "prod", "script_path": str(script)},
|
||||||
).json()
|
).json()
|
||||||
pid = prep["pending_id"]
|
pid = prep["pending_id"]
|
||||||
# Simulate a SUBMITTED state by mutating the pending directly
|
# Simulate a SUBMITTED state by mutating the pending directly
|
||||||
|
|||||||
@@ -42,14 +42,16 @@ def test_save_connection_emits_info_log(log_capture):
|
|||||||
assert "connection saved" in text
|
assert "connection saved" in text
|
||||||
|
|
||||||
|
|
||||||
def test_prepare_submit_job_emits_debug_and_info(log_capture):
|
def test_prepare_submit_job_emits_debug_and_info(log_capture, tmp_path):
|
||||||
connections.save_connection(name="prod", master="yarn", deploy_mode="cluster")
|
connections.save_connection(name="prod", master="yarn", deploy_mode="cluster")
|
||||||
|
script = tmp_path / "demo.py"
|
||||||
|
script.write_text("print('hi')\n")
|
||||||
log_capture.truncate(0); log_capture.seek(0)
|
log_capture.truncate(0); log_capture.seek(0)
|
||||||
submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py", queue="research")
|
submit.prepare_submit_job(connection="prod", script_path=str(script), queue="research")
|
||||||
text = log_capture.getvalue()
|
text = log_capture.getvalue()
|
||||||
assert "DEBUG|prepare_submit_job enter" in text
|
assert "DEBUG|prepare_submit_job enter" in text
|
||||||
assert "INFO|prepare_submit_job ok" in text
|
assert "INFO|prepare_submit_job ok" in text
|
||||||
assert "script_path=/tmp/j.py" in text
|
assert f"script_path={script}" in text
|
||||||
assert "queue=research" in text
|
assert "queue=research" in text
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -26,6 +26,14 @@ def _fresh(tmp_path: Path, monkeypatch):
|
|||||||
))
|
))
|
||||||
|
|
||||||
|
|
||||||
|
@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:
|
def _last_pending_id() -> str:
|
||||||
pending = submit.pending_store.list_all()
|
pending = submit.pending_store.list_all()
|
||||||
assert pending, "no pending submission was created"
|
assert pending, "no pending submission was created"
|
||||||
@@ -34,16 +42,16 @@ def _last_pending_id() -> str:
|
|||||||
|
|
||||||
# --- prepare_submit_job ---
|
# --- prepare_submit_job ---
|
||||||
|
|
||||||
def test_prepare_does_not_invoke_spark_submit(monkeypatch):
|
def test_prepare_does_not_invoke_spark_submit(monkeypatch, real_script):
|
||||||
called = []
|
called = []
|
||||||
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: called.append(cmd))
|
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: called.append(cmd))
|
||||||
out = submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py")
|
out = submit.prepare_submit_job(connection="prod", script_path=str(real_script))
|
||||||
assert called == [] # never called
|
assert called == [] # never called
|
||||||
assert out["status"] == "PENDING"
|
assert out["status"] == "PENDING"
|
||||||
assert out["pending_id"].startswith("p_")
|
assert out["pending_id"].startswith("p_")
|
||||||
|
|
||||||
|
|
||||||
def test_prepare_persists_pending_with_snapshot(monkeypatch):
|
def test_prepare_persists_pending_with_snapshot(monkeypatch, real_script):
|
||||||
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
|
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
|
||||||
connection_store.store.save(
|
connection_store.store.save(
|
||||||
Connection(
|
Connection(
|
||||||
@@ -54,7 +62,7 @@ def test_prepare_persists_pending_with_snapshot(monkeypatch):
|
|||||||
spark_conf={"spark.sql.shuffle.partitions": "200"},
|
spark_conf={"spark.sql.shuffle.partitions": "200"},
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py", queue="research")
|
submit.prepare_submit_job(connection="prod", script_path=str(real_script), queue="research")
|
||||||
p = submit.pending_store.get(_last_pending_id())
|
p = submit.pending_store.get(_last_pending_id())
|
||||||
assert p is not None
|
assert p is not None
|
||||||
assert p.connection == "prod"
|
assert p.connection == "prod"
|
||||||
@@ -63,18 +71,57 @@ def test_prepare_persists_pending_with_snapshot(monkeypatch):
|
|||||||
assert p.yarn_rm_url == "http://rm:8088"
|
assert p.yarn_rm_url == "http://rm:8088"
|
||||||
assert p.spark_conf == {"spark.sql.shuffle.partitions": "200"}
|
assert p.spark_conf == {"spark.sql.shuffle.partitions": "200"}
|
||||||
assert p.queue == "research"
|
assert p.queue == "research"
|
||||||
assert p.script_path == "/tmp/j.py"
|
assert p.script_path == str(real_script)
|
||||||
|
|
||||||
|
|
||||||
def test_prepare_raises_for_unknown_connection():
|
def test_prepare_raises_for_unknown_connection(real_script):
|
||||||
with pytest.raises(KeyError, match="missing"):
|
with pytest.raises(KeyError, match="missing"):
|
||||||
submit.prepare_submit_job(connection="missing", script_path="/tmp/j.py")
|
submit.prepare_submit_job(connection="missing", script_path=str(real_script))
|
||||||
|
|
||||||
|
|
||||||
def test_prepare_snapshots_connection_at_prepare_time(monkeypatch):
|
# --- 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)
|
||||||
|
|
||||||
|
|
||||||
|
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))
|
||||||
|
|
||||||
|
|
||||||
|
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="")
|
||||||
|
|
||||||
|
|
||||||
|
def test_prepare_error_message_mentions_generate_job_file(tmp_path):
|
||||||
|
"""The agent must be told to call generate_job_file first."""
|
||||||
|
with pytest.raises(ValueError, match="generate_job_file"):
|
||||||
|
submit.prepare_submit_job(
|
||||||
|
connection="prod", script_path=str(tmp_path / "x.py")
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
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))
|
||||||
|
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."""
|
"""Editing the connection between prepare and confirm must NOT silently retarget."""
|
||||||
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
|
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
|
||||||
submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py")
|
submit.prepare_submit_job(connection="prod", script_path=str(real_script))
|
||||||
p = submit.pending_store.get(_last_pending_id())
|
p = submit.pending_store.get(_last_pending_id())
|
||||||
# Now mutate the connection
|
# Now mutate the connection
|
||||||
connection_store.store.save(
|
connection_store.store.save(
|
||||||
@@ -88,8 +135,8 @@ def test_prepare_snapshots_connection_at_prepare_time(monkeypatch):
|
|||||||
|
|
||||||
# --- confirm_submit_job ---
|
# --- confirm_submit_job ---
|
||||||
|
|
||||||
def test_confirm_invokes_spark_submit_and_marks_submitted(monkeypatch):
|
def test_confirm_invokes_spark_submit_and_marks_submitted(monkeypatch, real_script):
|
||||||
submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py")
|
submit.prepare_submit_job(connection="prod", script_path=str(real_script))
|
||||||
pid = _last_pending_id()
|
pid = _last_pending_id()
|
||||||
fake_proc = type("P", (), {
|
fake_proc = type("P", (), {
|
||||||
"returncode": 0,
|
"returncode": 0,
|
||||||
@@ -100,7 +147,7 @@ def test_confirm_invokes_spark_submit_and_marks_submitted(monkeypatch):
|
|||||||
cmd = m.call_args.args[0]
|
cmd = m.call_args.args[0]
|
||||||
assert "yarn" in cmd
|
assert "yarn" in cmd
|
||||||
assert "cluster" in cmd
|
assert "cluster" in cmd
|
||||||
assert cmd[-1] == "/tmp/j.py"
|
assert cmd[-1] == str(real_script)
|
||||||
assert result.application_id == "application_17400000001"
|
assert result.application_id == "application_17400000001"
|
||||||
# pending updated
|
# pending updated
|
||||||
p = submit.pending_store.get(pid)
|
p = submit.pending_store.get(pid)
|
||||||
@@ -116,8 +163,8 @@ def test_confirm_raises_for_unknown_pending_id():
|
|||||||
submit.confirm_submit_job(pending_id="missing")
|
submit.confirm_submit_job(pending_id="missing")
|
||||||
|
|
||||||
|
|
||||||
def test_confirm_refuses_non_pending_status(monkeypatch):
|
def test_confirm_refuses_non_pending_status(monkeypatch, real_script):
|
||||||
submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py")
|
submit.prepare_submit_job(connection="prod", script_path=str(real_script))
|
||||||
pid = _last_pending_id()
|
pid = _last_pending_id()
|
||||||
# Mark it CANCELLED first
|
# Mark it CANCELLED first
|
||||||
p = submit.pending_store.get(pid)
|
p = submit.pending_store.get(pid)
|
||||||
@@ -127,9 +174,9 @@ def test_confirm_refuses_non_pending_status(monkeypatch):
|
|||||||
submit.confirm_submit_job(pending_id=pid)
|
submit.confirm_submit_job(pending_id=pid)
|
||||||
|
|
||||||
|
|
||||||
def test_confirm_marks_failed_on_spark_submit_error(monkeypatch):
|
def test_confirm_marks_failed_on_spark_submit_error(monkeypatch, real_script):
|
||||||
from spark_executor.core.spark_submit import SparkSubmitError
|
from spark_executor.core.spark_submit import SparkSubmitError
|
||||||
submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py")
|
submit.prepare_submit_job(connection="prod", script_path=str(real_script))
|
||||||
pid = _last_pending_id()
|
pid = _last_pending_id()
|
||||||
def _raise(_cmd):
|
def _raise(_cmd):
|
||||||
raise SparkSubmitError("boom")
|
raise SparkSubmitError("boom")
|
||||||
@@ -147,20 +194,23 @@ def test_list_pending_jobs_empty():
|
|||||||
assert submit.list_pending_jobs() == []
|
assert submit.list_pending_jobs() == []
|
||||||
|
|
||||||
|
|
||||||
def test_list_pending_jobs_returns_all(monkeypatch):
|
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)
|
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
|
||||||
submit.prepare_submit_job(connection="prod", script_path="/tmp/a.py")
|
submit.prepare_submit_job(connection="prod", script_path=str(a))
|
||||||
submit.prepare_submit_job(connection="prod", script_path="/tmp/b.py")
|
submit.prepare_submit_job(connection="prod", script_path=str(b))
|
||||||
out = submit.list_pending_jobs()
|
out = submit.list_pending_jobs()
|
||||||
assert {p["script_path"] for p in out} == {"/tmp/a.py", "/tmp/b.py"}
|
assert {p["script_path"] for p in out} == {str(a), str(b)}
|
||||||
assert all(p["status"] == "PENDING" for p in out)
|
assert all(p["status"] == "PENDING" for p in out)
|
||||||
|
|
||||||
|
|
||||||
# --- get_pending_job ---
|
# --- get_pending_job ---
|
||||||
|
|
||||||
def test_get_pending_job_returns_dump(monkeypatch):
|
def test_get_pending_job_returns_dump(monkeypatch, real_script):
|
||||||
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
|
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
|
||||||
submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py")
|
submit.prepare_submit_job(connection="prod", script_path=str(real_script))
|
||||||
pid = _last_pending_id()
|
pid = _last_pending_id()
|
||||||
out = submit.get_pending_job(pid)
|
out = submit.get_pending_job(pid)
|
||||||
assert out["pending_id"] == pid
|
assert out["pending_id"] == pid
|
||||||
@@ -175,9 +225,9 @@ def test_get_pending_job_unknown_raises():
|
|||||||
|
|
||||||
# --- cancel_pending_job ---
|
# --- cancel_pending_job ---
|
||||||
|
|
||||||
def test_cancel_pending_job_marks_cancelled(monkeypatch):
|
def test_cancel_pending_job_marks_cancelled(monkeypatch, real_script):
|
||||||
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
|
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
|
||||||
submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py")
|
submit.prepare_submit_job(connection="prod", script_path=str(real_script))
|
||||||
pid = _last_pending_id()
|
pid = _last_pending_id()
|
||||||
out = submit.cancel_pending_job(pid)
|
out = submit.cancel_pending_job(pid)
|
||||||
assert out == {"pending_id": pid, "status": "CANCELLED"}
|
assert out == {"pending_id": pid, "status": "CANCELLED"}
|
||||||
@@ -189,8 +239,8 @@ def test_cancel_pending_job_unknown_raises():
|
|||||||
submit.cancel_pending_job("missing")
|
submit.cancel_pending_job("missing")
|
||||||
|
|
||||||
|
|
||||||
def test_cancel_pending_job_refuses_submitted(monkeypatch):
|
def test_cancel_pending_job_refuses_submitted(monkeypatch, real_script):
|
||||||
submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py")
|
submit.prepare_submit_job(connection="prod", script_path=str(real_script))
|
||||||
pid = _last_pending_id()
|
pid = _last_pending_id()
|
||||||
p = submit.pending_store.get(pid)
|
p = submit.pending_store.get(pid)
|
||||||
p.status = "SUBMITTED"
|
p.status = "SUBMITTED"
|
||||||
|
|||||||
Reference in New Issue
Block a user