Claude
994c2d67ab
Exposes a single new MCP tool: generate_job_file(code) -> {script_path}.
Wiring follows the Stage 1 conventions (Pydantic body model for the
route, loguru DEBUG/INFO logging, summary+description for the startup
log, MCP arg-pattern via the body model so tools/call roundtrips long
code strings without 422-ing on FastAPI query length limits).
Flow:
1. LLM calls generate_job_file(code=...) -> {script_path}
2. LLM calls prepare_submit_job(connection=...,
script_path=...) -> {pending_id}
3. User reviews the file + pending record
4. LLM calls confirm_submit_job(pending_id=...) -> spark-submit runs
The output directory is SPARK_EXECUTOR_JOBS_DIR (default ./data/jobs/,
gitignored, persists across container restarts via the existing
./data volume mount in docker-compose.yml).
3 new tests:
- Unit: env-var override, default fallback in tmp cwd
- Integration: end-to-end body call with a > FastAPI-query-limit code
string (the canary test that would have caught the Stage 1
query-params-422 bug)
Live MCP smoke verified: tools/list shows 14 tools (13 from Stage 1 +
the new one), tools/call generate_job_file returns the absolute path
under ./data/jobs/.
28 lines
921 B
Python
28 lines
921 B
Python
# coding=utf-8
|
|
import os
|
|
from pathlib import Path
|
|
|
|
from spark_executor.tools import generate
|
|
|
|
|
|
def test_generate_writes_code_and_returns_path(monkeypatch, tmp_path: Path):
|
|
monkeypatch.chdir(tmp_path)
|
|
out = generate.generate_job_file(
|
|
"from pyspark.sql import SparkSession\n"
|
|
"spark = SparkSession.builder.getOrCreate()\n"
|
|
)
|
|
assert "script_path" in out
|
|
p = out["script_path"]
|
|
assert os.path.isabs(p)
|
|
assert p.startswith(str(tmp_path / "data" / "jobs"))
|
|
assert p.endswith(".py")
|
|
with open(p) as f:
|
|
assert "SparkSession.builder.getOrCreate()" in f.read()
|
|
|
|
|
|
def test_generate_uses_env_var_when_set(monkeypatch, tmp_path: Path):
|
|
monkeypatch.setenv("SPARK_EXECUTOR_JOBS_DIR", str(tmp_path / "custom"))
|
|
out = generate.generate_job_file("x = 1\n")
|
|
assert out["script_path"].startswith(str(tmp_path / "custom"))
|
|
assert os.path.isfile(out["script_path"])
|