tests/test_executor.py

"""Bounded subprocess tests for generated SpreadsheetBench scripts."""

from __future__ import annotations

import os
from pathlib import Path
import subprocess
import time

import openpyxl

from spreadsheetbench_env.executor import ExecutorLimits, ScriptExecutor


def _input_workbook(path: Path) -> None:
    workbook = openpyxl.Workbook()
    workbook.active["A1"] = 4
    workbook.save(path)
    workbook.close()


def test_executor_injects_private_paths_and_scrubs_parent_credentials(
    tmp_path: Path,
    monkeypatch,
) -> None:
    source = tmp_path / "source.xlsx"
    _input_workbook(source)
    monkeypatch.setenv("SPREADSHEETBENCH_TEST_SECRET", "must-not-leak")
    executor = ScriptExecutor(tmp_path / "session")
    code = """
import os
import openpyxl
INPUT_PATH = '/tmp/attacker-input.xlsx'
OUTPUT_PATH = '/tmp/attacker-output.xlsx'
wb = openpyxl.load_workbook(INPUT_PATH)
wb.active['A2'] = os.environ.get('SPREADSHEETBENCH_TEST_SECRET', 'scrubbed')
wb.active['A3'] = os.environ.get('OPENBLAS_NUM_THREADS')
wb.save(OUTPUT_PATH)
"""

    result = executor.run(code, source, case_name="case-1")

    assert result.ok is True
    assert result.output_path is not None
    assert result.output_path.parent == executor.root / "case-1"
    output = openpyxl.load_workbook(result.output_path, data_only=True)
    assert output.active["A1"].value == 4
    assert output.active["A2"].value == "scrubbed"
    assert output.active["A3"].value == "1"
    output.close()
    assert not Path("/tmp/attacker-output.xlsx").exists()
    executor.close()
    assert not executor.root.exists()


def test_executor_uses_a_distinct_directory_for_each_case(tmp_path: Path) -> None:
    source = tmp_path / "source.xlsx"
    _input_workbook(source)
    executor = ScriptExecutor(tmp_path / "session")
    code = "import shutil\nshutil.copyfile(INPUT_PATH, OUTPUT_PATH)\n"

    first = executor.run(code, source, case_name="case-1")
    second = executor.run(code, source, case_name="case-2")

    assert first.ok and second.ok
    assert first.output_path != second.output_path
    assert first.output_path is not None and second.output_path is not None
    assert first.output_path.parent != second.output_path.parent
    executor.close()


def test_executor_does_not_use_preexec_fn_in_threaded_server(
    tmp_path: Path,
    monkeypatch,
) -> None:
    source = tmp_path / "source.xlsx"
    _input_workbook(source)
    executor = ScriptExecutor(tmp_path / "session")
    real_popen = subprocess.Popen
    received_preexec: list[object] = []

    def recording_popen(*args, **kwargs):
        received_preexec.append(kwargs.get("preexec_fn"))
        return real_popen(*args, **kwargs)

    monkeypatch.setattr(subprocess, "Popen", recording_popen)

    result = executor.run(
        "import shutil\nshutil.copyfile(INPUT_PATH, OUTPUT_PATH)\n",
        source,
        case_name="no-preexec",
    )

    assert result.ok is True
    assert received_preexec == [None]
    executor.close()


def test_timeout_kills_the_generated_process_group(tmp_path: Path) -> None:
    source = tmp_path / "source.xlsx"
    _input_workbook(source)
    executor = ScriptExecutor(
        tmp_path / "session",
        limits=ExecutorLimits(timeout_s=0.5, cpu_seconds=5),
    )
    code = """
import subprocess
import sys
import time
child = subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(60)'])
print(child.pid, flush=True)
time.sleep(60)
"""

    result = executor.run(code, source, case_name="timeout")

    assert result.ok is False
    assert result.timed_out is True
    assert "timeout" in result.error.lower()
    child_pid = int(result.stdout.strip().splitlines()[0])
    deadline = time.monotonic() + 3
    while time.monotonic() < deadline:
        try:
            os.kill(child_pid, 0)
        except ProcessLookupError:
            break
        time.sleep(0.05)
    else:
        raise AssertionError(f"generated child process {child_pid} survived timeout")
    executor.close()


def test_timeout_does_not_wait_for_detached_child_output_pipes(tmp_path: Path) -> None:
    source = tmp_path / "source.xlsx"
    _input_workbook(source)
    executor = ScriptExecutor(
        tmp_path / "session",
        limits=ExecutorLimits(timeout_s=0.2, cpu_seconds=5),
    )
    code = """
import subprocess
import sys
import time
child = subprocess.Popen(
    [sys.executable, '-c', 'import time; time.sleep(10)'],
    start_new_session=True,
)
print(child.pid, flush=True)
time.sleep(10)
"""

    started = time.monotonic()
    result = executor.run(code, source, case_name="detached-child")
    elapsed = time.monotonic() - started

    assert result.ok is False
    assert result.timed_out is True
    assert elapsed < 2.0
    child_pid = int(result.stdout.strip().splitlines()[0])
    try:
        os.kill(child_pid, 0)
    except ProcessLookupError:
        pass
    else:
        os.kill(child_pid, 9)
    executor.close()


def test_executor_rejects_unsafe_case_names(tmp_path: Path) -> None:
    source = tmp_path / "source.xlsx"
    _input_workbook(source)
    executor = ScriptExecutor(tmp_path / "session")

    try:
        executor.run("", source, case_name="../escape")
    except ValueError as error:
        assert "case name" in str(error)
    else:
        raise AssertionError("unsafe case name was accepted")
    finally:
        executor.close()


def test_concurrent_executors_lease_distinct_runner_uids_and_release_them(
    tmp_path: Path,
) -> None:
    limits = ExecutorLimits(run_uid_base=20_000, run_gid=65_534)
    first = ScriptExecutor(tmp_path / "first", limits=limits)
    second = ScriptExecutor(tmp_path / "second", limits=limits)

    assert first.leased_uid == 20_000
    assert second.leased_uid == 20_001

    first.close()
    replacement = ScriptExecutor(tmp_path / "replacement", limits=limits)
    assert replacement.leased_uid == 20_000

    second.close()
    replacement.close()