mcp-server/tests/unit/test_job_file.py

# 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("<config/>\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_accepts_content_larger_than_former_1mb_cap(tmp_path: Path):
    """The 1 MB cap was removed (commit a5e6d1f and following). Content
    larger than the old cap must now be accepted; size is bounded by
    the MCP transport, not the tool."""
    config.settings.jobs_dir = str(tmp_path)
    p = tmp_path / "big.py"
    p.write_text("# small\n")
    big = "x" * (2 * 1024 * 1024)
    out = update_job_file(str(p), big)
    assert out["bytes_written"] == len(big)
    assert p.read_text() == big


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