tao.chen/mcp-server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(mcp): rename generate_job_file to write_job_file to match what it does

Browse Source 
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>
This commit is contained in:
Claude
2026-06-29 19:13:29 +08:00
co-authored by Claude Fable 5
parent 0fef77c0b8
commit 523e6a9c76
18 changed files with 3661 additions and 154 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/integration/test_mcp_routes.py
+22 -22
View File
@@ -52,7 +52,7 @@ def test_sixteen_tool_routes_registered():
"/get_connection",
"/delete_connection",
# LLM-driven PySpark generation (3) — Stage 2
"/generate_job_file",
"/write_job_file",
"/read_job_file",
"/update_job_file",
):
@@ -146,7 +146,7 @@ def test_prepare_submit_job_works_via_body(tmp_path):
"connection": "prod",
"script_path": str(script),
"queue": "research",
"executor_memory": "4G",
"executor_memory": "2G",
"executor_cores": 2,
"num_executors": 2,
"app_name": "test-app",
@@ -177,7 +177,7 @@ def test_prepare_rejects_unconfirmed_defaults(tmp_path):
detail = r.json()["detail"]
assert "Please confirm default values" in detail
assert "queue='default'" in detail
assert "executor_memory='4G'" in detail
assert "executor_memory='2G'" in detail
assert "executor_cores=2" in detail
assert "num_executors=2" in detail
@@ -193,7 +193,7 @@ def test_prepare_accepts_confirmed_defaults(tmp_path):
"connection": "prod",
"script_path": str(script),
"queue": "default",
"executor_memory": "4G",
"executor_memory": "2G",
"executor_cores": 2,
"num_executors": 2,
"app_name": "test-app",
@@ -203,7 +203,7 @@ def test_prepare_accepts_confirmed_defaults(tmp_path):
body = r.json()
assert body["status"] == "PENDING"
assert body["parameters"]["queue"] == "default"
assert body["parameters"]["executor_memory"] == "4G"
assert body["parameters"]["executor_memory"] == "2G"
assert body["parameters"]["executor_cores"] == 2
assert body["parameters"]["num_executors"] == 2
@@ -218,7 +218,7 @@ def test_prepare_rejects_nonexistent_script_path_with_400():
"connection": "prod",
"script_path": "/nope/does_not_exist.py",
"queue": "default",
"executor_memory": "4G",
"executor_memory": "2G",
"executor_cores": 2,
"num_executors": 2,
"app_name": "test-app",
@@ -228,12 +228,12 @@ def test_prepare_rejects_nonexistent_script_path_with_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
assert "write_job_file" in detail
def test_prepare_rejects_sql_policy_violation_with_400(tmp_path):
"""SQL guard: even if the file exists, prepare_submit_job must reject
code containing DROP/DELETE/etc. (defense in depth — generate_job_file
code containing DROP/DELETE/etc. (defense in depth — write_job_file
also enforces this, but a host-mounted file might bypass that)."""
c = TestClient(app)
c.post("/save_connection", json={"name": "prod", "master": "yarn"})
@@ -245,7 +245,7 @@ def test_prepare_rejects_sql_policy_violation_with_400(tmp_path):
"connection": "prod",
"script_path": str(bad_script),
"queue": "default",
"executor_memory": "4G",
"executor_memory": "2G",
"executor_cores": 2,
"num_executors": 2,
"app_name": "test-app",
@@ -257,14 +257,14 @@ def test_prepare_rejects_sql_policy_violation_with_400(tmp_path):
assert "SELECT" in detail and "INSERT" in detail # the policy is explained
def test_generate_rejects_sql_policy_violation_with_400(tmp_path):
"""Same policy at generate_job_file — agent gets immediate feedback
def test_write_rejects_sql_policy_violation_with_400(tmp_path):
"""Same policy at write_job_file — agent gets immediate feedback
before the file is even written to disk."""
from common import config
config.settings.jobs_dir = str(tmp_path / "jobs")
c = TestClient(app)
r = c.post(
"/generate_job_file",
"/write_job_file",
json={"code": 'spark.sql("DELETE FROM events")\n'},
)
assert r.status_code == 400
@@ -331,7 +331,7 @@ def test_list_and_get_pending_job_roundtrip_via_body(tmp_path):
"connection": "prod",
"script_path": str(script),
"queue": "default",
"executor_memory": "4G",
"executor_memory": "2G",
"executor_cores": 2,
"num_executors": 2,
"app_name": "test-app",
@@ -354,15 +354,15 @@ def test_list_connections_works_with_empty_body():
assert r.json() == []
# --- Stage 2: generate_job_file ---
# --- Stage 2: write_job_file ---
def test_generate_job_file_accepts_code_string_in_body(tmp_path, monkeypatch):
def test_write_job_file_accepts_code_string_in_body(tmp_path, monkeypatch):
"""End-to-end: a Pydantic body model lets tools/call pass a code string
that FastAPI would reject if it were a query parameter (length limits)."""
monkeypatch.chdir(tmp_path)
c = TestClient(app)
code = "print('from MCP integration test')\n" * 100 # > FastAPI query limit
r = c.post("/generate_job_file", json={"code": code})
r = c.post("/write_job_file", json={"code": code})
assert r.status_code == 200, r.text
p = r.json()["script_path"]
assert os.path.isfile(p)
@@ -401,7 +401,7 @@ def test_unknown_connection_in_prepare_returns_404():
"connection": "nope",
"script_path": "/tmp/x.py",
"queue": "default",
"executor_memory": "4G",
"executor_memory": "2G",
"executor_cores": 2,
"num_executors": 2,
"app_name": "test-app",
@@ -423,7 +423,7 @@ def test_confirm_non_pending_returns_400(tmp_path):
"connection": "prod",
"script_path": str(script),
"queue": "default",
"executor_memory": "4G",
"executor_memory": "2G",
"executor_cores": 2,
"num_executors": 2,
"app_name": "test-app",
@@ -448,7 +448,7 @@ def test_cancel_already_submitted_returns_400(tmp_path):
"connection": "prod",
"script_path": str(script),
"queue": "default",
"executor_memory": "4G",
"executor_memory": "2G",
"executor_cores": 2,
"num_executors": 2,
"app_name": "test-app",
@@ -477,7 +477,7 @@ def test_update_pending_job_route(tmp_path):
"connection": "prod",
"script_path": str(script),
"queue": "default",
"executor_memory": "4G",
"executor_memory": "2G",
"executor_cores": 2,
"num_executors": 2,
"app_name": "test-app",
@@ -496,7 +496,7 @@ def test_update_pending_job_route(tmp_path):
got = c.post("/get_pending_job", json={"pending_id": pid}).json()
assert got["queue"] == "research"
# untouched fields are preserved
assert got["executor_memory"] == "4G"
assert got["executor_memory"] == "2G"
assert got["app_name"] == "test-app"
@@ -511,7 +511,7 @@ def test_update_pending_job_route_rejects_submitted(tmp_path, monkeypatch):
"connection": "prod",
"script_path": str(script),
"queue": "default",
"executor_memory": "4G",
"executor_memory": "2G",
"executor_cores": 2,
"num_executors": 2,
"app_name": "test-app",
tests/unit/test_job_file.py
+1 -1
View File
@@ -145,7 +145,7 @@ def test_full_edit_cycle_via_tools(tmp_path: Path):
"""Simulate the agent's review-and-edit loop end-to-end."""
config.settings.jobs_dir = str(tmp_path)
path = tmp_path / "script.py"
# 1. Initial state: file doesn't exist (in real flow, generate_job_file
# 1. Initial state: file doesn't exist (in real flow, write_job_file
# would create it; we just create it directly here for the test)
path.write_text("# v1: initial\n")
# 2. Agent reads back
tests/unit/test_logging.py
+1 -1
View File
@@ -51,7 +51,7 @@ def test_prepare_submit_job_emits_debug_and_info(log_capture, tmp_path):
connection="prod",
script_path=str(script),
queue="research",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
tests/unit/test_models.py
+3 -3
View File
@@ -95,7 +95,7 @@ def test_pending_submission_defaults_to_pending_status():
deploy_mode="cluster",
script_path="/tmp/j.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
spark_conf={},
@@ -120,7 +120,7 @@ def test_pending_submission_carries_yarn_rm_url_snapshot():
yarn_rm_url="http://rm-prod:8088",
script_path="/tmp/j.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
spark_conf={},
@@ -139,7 +139,7 @@ def test_pending_submission_can_record_outcome():
deploy_mode="cluster",
script_path="/tmp/j.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
spark_conf={},
tests/unit/test_pending_store.py
+3 -3
View File
@@ -18,7 +18,7 @@ def _pending(pid: str = "p_1", created_at: datetime | None = None) -> PendingSub
deploy_mode="cluster",
script_path="/tmp/j.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
spark_conf={},
@@ -85,7 +85,7 @@ def _legacy_record_text() -> str:
return (
'{"p_legacy": {"pending_id": "p_legacy", "app_name": "legacy", '
'"connection": "prod", "master": "yarn", "deploy_mode": "cluster", '
'"script_path": "/tmp/j.py", "queue": "default", "executor_memory": "4G", '
'"script_path": "/tmp/j.py", "queue": "default", "executor_memory": "2G", '
'"executor_cores": 2, "num_executors": 2, "spark_conf": {}, '
'"created_at": "2026-06-23T00:00:00"}}'
)
@@ -139,7 +139,7 @@ def test_loads_records_with_missing_or_null_app_name(tmp_path: Path):
"deploy_mode": "cluster",
"script_path": "/tmp/j.py",
"queue": "default",
"executor_memory": "4G",
"executor_memory": "2G",
"executor_cores": 2,
"num_executors": 2,
"spark_conf": {},
tests/unit/test_requests.py
+2 -2
View File
@@ -16,7 +16,7 @@ def test_prepare_submit_job_request_restores_defaults_for_queue_and_resources():
app_name="test-app",
)
assert req.queue == "default"
assert req.executor_memory == "4G"
assert req.executor_memory == "2G"
assert req.executor_cores == 2
assert req.num_executors == 2
@@ -26,7 +26,7 @@ def test_prepare_submit_job_request_accepts_extra_args():
connection="prod",
script_path="/tmp/x.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
tests/unit/test_spark_submit.py
+9 -9
View File
@@ -36,7 +36,7 @@ def test_build_command_uses_provided_master_and_deploy_mode():
deploy_mode="cluster",
script_path="/tmp/jobs/job_001.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
)
@@ -44,7 +44,7 @@ def test_build_command_uses_provided_master_and_deploy_mode():
assert "yarn" in cmd
assert "cluster" in cmd
assert "--queue" in cmd and "default" in cmd
assert "--executor-memory" in cmd and "4G" in cmd
assert "--executor-memory" in cmd and "2G" in cmd
assert "--executor-cores" in cmd and "2" in cmd
assert "--num-executors" in cmd and "2" in cmd
assert cmd[-1] == "/tmp/jobs/job_001.py"
@@ -56,7 +56,7 @@ def test_extra_args_none_default():
deploy_mode="cluster",
script_path="/tmp/jobs/job_001.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
)
@@ -70,7 +70,7 @@ def test_extra_args_single_pair():
deploy_mode="cluster",
script_path="/tmp/jobs/job_001.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
extra_args={"jars": "hdfs:///libs/foo.jar"},
@@ -88,7 +88,7 @@ def test_extra_args_multiple_pairs_preserve_order():
deploy_mode="cluster",
script_path="/tmp/jobs/job_001.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
extra_args={
@@ -112,7 +112,7 @@ def test_extra_args_does_not_collide_with_spark_conf():
deploy_mode="cluster",
script_path="/tmp/jobs/job_001.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
spark_conf={"spark.executor.memory": "8G"},
@@ -146,7 +146,7 @@ def test_build_command_appends_spark_conf_entries():
deploy_mode="cluster",
script_path="/tmp/j.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
spark_conf={"spark.sql.shuffle.partitions": "200", "spark.executor.memoryOverhead": "1G"},
@@ -169,7 +169,7 @@ def test_build_command_uses_settings_spark_submit_bin():
deploy_mode="cluster",
script_path="/tmp/j.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
)
@@ -183,7 +183,7 @@ def test_build_command_supports_pyspark_binary():
deploy_mode="cluster",
script_path="/tmp/j.py",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
)
tests/unit/test_submit_tool.py
+33 -33
View File
@@ -56,7 +56,7 @@ def test_prepare_does_not_invoke_spark_submit(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -81,7 +81,7 @@ def test_prepare_persists_pending_with_snapshot(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="research",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -103,7 +103,7 @@ def test_prepare_persists_app_name(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="my-etl-job",
@@ -125,7 +125,7 @@ def test_prepare_accepts_restored_defaults(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -133,7 +133,7 @@ def test_prepare_accepts_restored_defaults(monkeypatch, real_script):
assert out["status"] == "PENDING"
p = submit.pending_store.get(_last_pending_id())
assert p.queue == "default"
assert p.executor_memory == "4G"
assert p.executor_memory == "2G"
assert p.executor_cores == 2
assert p.num_executors == 2
@@ -144,7 +144,7 @@ def test_prepare_snapshots_extra_args(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -173,7 +173,7 @@ def test_prepare_raises_for_unknown_connection(real_script):
connection="missing",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -189,7 +189,7 @@ def test_prepare_rejects_nonexistent_script_path(tmp_path):
connection="prod",
script_path=missing,
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -203,7 +203,7 @@ def test_prepare_rejects_directory_as_script_path(tmp_path):
connection="prod",
script_path=str(tmp_path),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -216,21 +216,21 @@ def test_prepare_rejects_empty_script_path():
connection="prod",
script_path="",
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
)
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"):
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="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -245,7 +245,7 @@ def test_confirm_rejects_if_script_was_deleted_after_prepare(tmp_path, monkeypat
connection="prod",
script_path=str(script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -264,7 +264,7 @@ def test_prepare_snapshots_connection_at_prepare_time(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -287,7 +287,7 @@ def test_confirm_invokes_spark_submit_and_marks_submitted(monkeypatch, real_scri
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -323,7 +323,7 @@ def test_confirm_refuses_non_pending_status(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -342,7 +342,7 @@ def test_confirm_marks_failed_on_spark_submit_error(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -368,7 +368,7 @@ def test_confirm_retries_spark_submit_failure_then_succeeds(monkeypatch, real_sc
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -395,7 +395,7 @@ def test_confirm_exhausts_retries_and_sets_failed(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -419,7 +419,7 @@ def test_confirm_idempotent_when_submitted(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -446,7 +446,7 @@ def test_confirm_resets_failed_and_retries(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -476,7 +476,7 @@ def test_confirm_updates_pending_before_job_store(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -510,7 +510,7 @@ def test_list_pending_jobs_returns_all(monkeypatch, tmp_path):
connection="prod",
script_path=str(a),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -519,7 +519,7 @@ def test_list_pending_jobs_returns_all(monkeypatch, tmp_path):
connection="prod",
script_path=str(b),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -537,7 +537,7 @@ def test_get_pending_job_returns_dump(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -562,7 +562,7 @@ def test_cancel_pending_job_marks_cancelled(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -583,7 +583,7 @@ def test_cancel_pending_job_refuses_submitted(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -604,7 +604,7 @@ def test_update_pending_updates_resource_params(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -639,7 +639,7 @@ def test_update_pending_rejects_non_pending_status(monkeypatch, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -657,7 +657,7 @@ def test_update_pending_rejects_bad_script_path(tmp_path, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
@@ -673,7 +673,7 @@ def test_update_pending_rejects_sql_violation_script(tmp_path, real_script):
connection="prod",
script_path=str(real_script),
queue="default",
executor_memory="4G",
executor_memory="2G",
executor_cores=2,
num_executors=2,
app_name="test-app",
tests/unit/test_generate_tool.py → tests/unit/test_write_job_tool.py
+15 -5
View File
@@ -5,7 +5,7 @@ from pathlib import Path
import pytest
from common import config
from spark_executor.tools import generate
from spark_executor.tools import write_job
@pytest.fixture(autouse=True)
@@ -23,9 +23,9 @@ def _restore_settings():
config.settings.log_level = snapshot.log_level
def test_generate_writes_code_and_returns_path(monkeypatch, tmp_path: Path):
def test_write_writes_code_and_returns_path(monkeypatch, tmp_path: Path):
config.settings.jobs_dir = str(tmp_path / "data" / "jobs")
out = generate.generate_job_file(
out = write_job.write_job_file(
"from pyspark.sql import SparkSession\n"
"spark = SparkSession.builder.getOrCreate()\n"
)
@@ -38,8 +38,18 @@ def test_generate_writes_code_and_returns_path(monkeypatch, tmp_path: Path):
assert "SparkSession.builder.getOrCreate()" in f.read()
def test_generate_uses_settings_jobs_dir(tmp_path: Path):
def test_write_uses_settings_jobs_dir(tmp_path: Path):
config.settings.jobs_dir = str(tmp_path / "custom")
out = generate.generate_job_file("x = 1\n")
out = write_job.write_job_file("x = 1\n")
assert out["script_path"].startswith(str(tmp_path / "custom"))
assert os.path.isfile(out["script_path"])
def test_write_rejects_sql_violation(tmp_path: Path):
"""The SQL guard is the whole reason this isn't a bare file-write
tool: the LLM might have produced a script with a forbidden
statement and we want to fail at write time, not at prepare time."""
from spark_executor.tools.write_job import SqlGuardViolation
config.settings.jobs_dir = str(tmp_path / "data" / "jobs")
with pytest.raises(SqlGuardViolation):
write_job.write_job_file("spark.sql('DROP TABLE users')\n")