# coding=utf-8 import os from pathlib import Path import pytest from common import config from spark_executor.tools import job_file from spark_executor.tools.job_file import ( ScriptFileError, read_job_file, update_job_file, ) @pytest.fixture(autouse=True) def _restore_settings(): snapshot = config.Settings( data_dir=config.settings.data_dir, jobs_dir=config.settings.jobs_dir, log_dir=config.settings.log_dir, yarn_resource_manager_url=config.settings.yarn_resource_manager_url, log_level=config.settings.log_level, ) yield config.settings.data_dir = snapshot.data_dir config.settings.jobs_dir = snapshot.jobs_dir config.settings.log_dir = snapshot.log_dir config.settings.yarn_resource_manager_url = snapshot.yarn_resource_manager_url config.settings.log_level = snapshot.log_level # --- read_job_file --- def test_read_returns_existing_content(tmp_path: Path): p = tmp_path / "demo.py" p.write_text("print('hello')\n") out = read_job_file(str(p)) assert out["content"] == "print('hello')\n" assert out["path"] == str(p) assert out["size"] == len("print('hello')\n") def test_read_works_for_any_existing_file(tmp_path: Path): """read_job_file has no jobs_dir restriction — it's read-only and the agent needs to inspect arbitrary files (e.g. logs/, /etc/hadoop/conf).""" p = tmp_path / "anything.py" p.write_text("x = 1\n") out = read_job_file(str(p)) assert out["content"] == "x = 1\n" def test_read_raises_for_missing_file(tmp_path: Path): with pytest.raises(ScriptFileError, match="does not exist"): read_job_file(str(tmp_path / "nope.py")) def test_read_raises_for_directory(tmp_path: Path): with pytest.raises(ScriptFileError, match="does not exist"): read_job_file(str(tmp_path)) def test_read_raises_for_empty_path(): with pytest.raises(ScriptFileError, match="does not exist"): read_job_file("") # --- update_job_file --- def test_update_overwrites_existing_file_in_jobs_dir(tmp_path: Path): """The file must be under settings.jobs_dir (default './data/jobs').""" config.settings.jobs_dir = str(tmp_path) p = tmp_path / "script.py" p.write_text("v1\n") out = update_job_file(str(p), "v2\n") assert out["bytes_written"] == 3 assert p.read_text() == "v2\n" def test_update_rejects_paths_outside_jobs_dir(tmp_path: Path): """update_job_file is sandboxed to settings.jobs_dir.""" config.settings.jobs_dir = str(tmp_path / "jobs") other = tmp_path / "elsewhere.py" other.write_text("x\n") with pytest.raises(ScriptFileError, match="must be under"): update_job_file(str(other), "y\n") def test_update_rejects_paths_outside_jobs_dir_via_existing_file(tmp_path: Path): """Path validation: even when the file exists, an absolute path OUTSIDE settings.jobs_dir is rejected. (In production, /etc/passwd would exist in the container; here we synthesize the same scenario with a real file outside the configured jobs_dir.)""" config.settings.jobs_dir = str(tmp_path / "jobs") (tmp_path / "jobs").mkdir() outside = tmp_path / "hadoop-core-site.xml" outside.write_text("\n") with pytest.raises(ScriptFileError, match="must be under"): update_job_file(str(outside), "tampered\n") def test_update_rejects_relative_escape_attempt(tmp_path: Path): """Path traversal via '..' must not bypass the jobs_dir check.""" config.settings.jobs_dir = str(tmp_path / "jobs") (tmp_path / "jobs").mkdir() (tmp_path / "etc").mkdir() secret = tmp_path / "etc" / "passwd" secret.write_text("root:x:0:0:...\n") # Try to reach the file via '../etc/passwd' relative to jobs_dir sneaky = str(tmp_path / "jobs" / ".." / "etc" / "passwd") with pytest.raises(ScriptFileError, match="must be under"): update_job_file(sneaky, "tampered\n") def test_update_rejects_missing_file(tmp_path: Path): config.settings.jobs_dir = str(tmp_path) with pytest.raises(ScriptFileError, match="does not exist"): update_job_file(str(tmp_path / "nope.py"), "x\n") def test_update_rejects_oversized_content(tmp_path: Path, monkeypatch): """Cap at 1 MB so the MCP response stays bounded.""" config.settings.jobs_dir = str(tmp_path) p = tmp_path / "big.py" p.write_text("# small\n") # Synthesize 2 MB of content (don't actually write 2 MB to disk) big = "x" * (2 * 1024 * 1024) with pytest.raises(ScriptFileError, match="too large"): update_job_file(str(p), big) def test_update_rejects_1mb_plus_1_byte(tmp_path: Path): """Exactly at the boundary: 1 MB + 1 byte must be rejected.""" config.settings.jobs_dir = str(tmp_path) p = tmp_path / "x.py" p.write_text("# t\n") just_over = "x" * (1024 * 1024 + 1) with pytest.raises(ScriptFileError, match="too large"): update_job_file(str(p), just_over) # --- round-trip: write -> read -> update -> read --- 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, 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 r1 = read_job_file(str(path)) assert r1["content"] == "# v1: initial\n" # 3. Agent edits (LLM or human) edited = "# v2: edited by agent\n" + r1["content"].split("\n", 1)[1] # 4. Agent writes back w = update_job_file(str(path), edited) assert w["bytes_written"] == len(edited) # 5. Read back to verify r2 = read_job_file(str(path)) assert r2["content"] == edited assert "v2: edited" in r2["content"]