mcp-server/tests/unit/test_submit_tool.py

# 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 "")


# --- list_pending_jobs ---

def test_list_pending_jobs_empty():
    assert submit.list_pending_jobs() == []


def test_list_pending_jobs_returns_all(monkeypatch):
    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="/tmp/b.py")
    out = submit.list_pending_jobs()
    assert {p["script_path"] for p in out} == {"/tmp/a.py", "/tmp/b.py"}
    assert all(p["status"] == "PENDING" for p in out)


# --- get_pending_job ---

def test_get_pending_job_returns_dump(monkeypatch):
    monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
    submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py")
    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")