write_job_file now validates that the incoming code string is encodable
as UTF-8 BEFORE doing any file I/O. A lone surrogate (U+D800..U+DFFF) —
typically the result of a broken decode somewhere upstream in the
agent's pipeline — cannot be represented in UTF-8, and would otherwise
either silently produce mojibake on disk or crash with an opaque
UnicodeEncodeError from open().
Behavior:
- try: code.encode('utf-8')
- on UnicodeEncodeError: log a warning and raise ValueError with a
message that points the agent at 're-generate the code as plain
Unicode'. ValueError -> HTTP 400 via the existing server.py
exception handler, so the agent gets a clean error and no file is
written.
- normal CJK / emoji / accented Latin still work — the encode check
is essentially free.
Tests:
- test_write_accepts_valid_utf8_including_cjk_and_emoji: positive
case for the common LLM output patterns.
- test_write_rejects_lone_surrogate: negative case, also asserts
that tmp_path is empty (no half-written .py left behind).
Full suite 246 passed (+2 from 244).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
118 lines
4.1 KiB
Python
118 lines
4.1 KiB
Python
# coding=utf-8
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from common import config
|
|
from spark_executor.core.job_writer import (
|
|
ENV_JOBS_DIR,
|
|
resolve_jobs_dir,
|
|
write_job_file,
|
|
)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _restore_settings():
|
|
"""job_writer uses settings.jobs_dir which is captured at module import.
|
|
Tests that set env vars then call settings.reload() — this fixture saves
|
|
and restores the original settings so one test doesn't leak into another."""
|
|
snapshot = config.Settings(
|
|
data_dir=config.settings.data_dir,
|
|
jobs_dir=config.settings.jobs_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.yarn_resource_manager_url = snapshot.yarn_resource_manager_url
|
|
config.settings.log_level = snapshot.log_level
|
|
|
|
|
|
# --- resolve_jobs_dir (priority: arg > settings.jobs_dir) ---
|
|
|
|
def test_resolve_explicit_arg_wins(tmp_path: Path, monkeypatch):
|
|
config.settings.jobs_dir = "/from/settings"
|
|
assert resolve_jobs_dir(str(tmp_path)) == str(tmp_path)
|
|
|
|
|
|
def test_resolve_uses_settings_jobs_dir_when_no_arg(monkeypatch, tmp_path: Path):
|
|
config.settings.jobs_dir = str(tmp_path)
|
|
assert resolve_jobs_dir() == str(tmp_path)
|
|
|
|
|
|
def test_resolve_default_when_settings_unset(monkeypatch):
|
|
config.settings.jobs_dir = "./data/jobs"
|
|
assert resolve_jobs_dir() == "./data/jobs"
|
|
|
|
|
|
# --- write_job_file ---
|
|
|
|
def test_write_creates_file_with_code_and_returns_abs_path(tmp_path: Path):
|
|
out = write_job_file("print('hi')\n", jobs_dir=str(tmp_path))
|
|
assert os.path.isfile(out)
|
|
assert os.path.isabs(out)
|
|
assert out.startswith(str(tmp_path))
|
|
assert out.endswith(".py")
|
|
with open(out) as f:
|
|
assert f.read() == "print('hi')\n"
|
|
|
|
|
|
def test_write_creates_jobs_dir_if_missing(tmp_path: Path):
|
|
target = tmp_path / "newdir"
|
|
assert not target.exists()
|
|
out = write_job_file("x = 1\n", jobs_dir=str(target))
|
|
assert target.is_dir()
|
|
assert os.path.isfile(out)
|
|
|
|
|
|
def test_write_uses_settings_jobs_dir_when_no_arg(tmp_path: Path):
|
|
config.settings.jobs_dir = str(tmp_path)
|
|
out = write_job_file("settings-driven\n")
|
|
assert out.startswith(str(tmp_path))
|
|
with open(out) as f:
|
|
assert f.read() == "settings-driven\n"
|
|
|
|
|
|
def test_write_uses_default_jobs_dir(tmp_path: Path):
|
|
"""When settings.jobs_dir is the default './data/jobs', writes resolve
|
|
relative to cwd. chdir to tmp so the test doesn't pollute the project."""
|
|
config.settings.jobs_dir = "./data/jobs"
|
|
old_cwd = os.getcwd()
|
|
try:
|
|
os.chdir(tmp_path)
|
|
out = write_job_file("default\n")
|
|
assert os.path.isabs(out)
|
|
expected_dir = tmp_path / "data" / "jobs"
|
|
assert expected_dir.is_dir()
|
|
assert str(out).startswith(str(expected_dir))
|
|
finally:
|
|
os.chdir(old_cwd)
|
|
|
|
|
|
def test_write_returns_unique_paths_for_concurrent_calls(tmp_path: Path):
|
|
"""Two writes in the same second must still get distinct filenames (via random suffix)."""
|
|
out1 = write_job_file("a", jobs_dir=str(tmp_path))
|
|
out2 = write_job_file("b", jobs_dir=str(tmp_path))
|
|
assert out1 != out2
|
|
|
|
|
|
def test_write_accepts_valid_utf8_including_cjk_and_emoji(tmp_path: Path):
|
|
"""Common LLM output — CJK + emoji — must encode as UTF-8 cleanly."""
|
|
code = "# 中文注释 ✨\nspark.read.parquet('/data/中文.parquet')\n"
|
|
out = write_job_file(code, jobs_dir=str(tmp_path))
|
|
with open(out, encoding="utf-8") as f:
|
|
assert f.read() == code
|
|
|
|
|
|
def test_write_rejects_lone_surrogate(tmp_path: Path):
|
|
"""A lone surrogate (U+D800..U+DFFF) is not a valid Unicode scalar value
|
|
and cannot be encoded as UTF-8. write_job_file must reject it BEFORE
|
|
writing — never leave a half-written .py on disk."""
|
|
bad = "print('ok')\n\udcff\noops"
|
|
with pytest.raises(ValueError, match="not valid UTF-8"):
|
|
write_job_file(bad, jobs_dir=str(tmp_path))
|
|
# No file was written.
|
|
assert list(tmp_path.iterdir()) == []
|