133 lines
4.9 KiB
Python
133 lines
4.9 KiB
Python
# coding=utf-8
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
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
|
|
|
|
|
|
@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"))
|
|
|
|
|
|
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):
|
|
called = []
|
|
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: called.append(cmd))
|
|
out = submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py")
|
|
assert called == [] # never called
|
|
assert out["status"] == "PENDING"
|
|
assert out["pending_id"].startswith("p_")
|
|
|
|
|
|
def test_prepare_persists_pending_with_snapshot(monkeypatch):
|
|
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
|
|
connection_store.store.save(
|
|
Connection(
|
|
name="prod",
|
|
master="yarn",
|
|
deploy_mode="cluster",
|
|
spark_conf={"spark.sql.shuffle.partitions": "200"},
|
|
)
|
|
)
|
|
submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py", queue="research")
|
|
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.spark_conf == {"spark.sql.shuffle.partitions": "200"}
|
|
assert p.queue == "research"
|
|
assert p.script_path == "/tmp/j.py"
|
|
|
|
|
|
def test_prepare_raises_for_unknown_connection():
|
|
with pytest.raises(KeyError, match="missing"):
|
|
submit.prepare_submit_job(connection="missing", script_path="/tmp/j.py")
|
|
|
|
|
|
def test_prepare_snapshots_connection_at_prepare_time(monkeypatch):
|
|
"""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="/tmp/j.py")
|
|
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):
|
|
submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py")
|
|
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] == "/tmp/j.py"
|
|
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
|
|
|
|
|
|
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):
|
|
submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py")
|
|
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):
|
|
from spark_executor.core.spark_submit import SparkSubmitError
|
|
submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py")
|
|
pid = _last_pending_id()
|
|
def _raise(_cmd):
|
|
raise SparkSubmitError("boom")
|
|
monkeypatch.setattr(submit, "run_spark_submit", _raise)
|
|
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 "")
|