Something went wrong. Try again.
My version of https://github.com/vercel-labs/fx
Something went wrong. Try again.
Python
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362from __future__ import annotations
import contextlibimport ioimport jsonimport osimport pathlibimport shleximport signalimport sysimport tempfileimport threadingimport timeimport unittestfrom unittest import mock
from scripts.pgso.model import PgsoErrorfrom scripts.pgso.runner import cancellation_guard, run_checked
class PgsoRunnerTests(unittest.TestCase): def test_cancellation_guard_converts_sigterm_and_restores_handlers(self) -> None: installed: dict[object, object] = {}
def install(signum, handler): installed[signum] = handler
with ( mock.patch("scripts.pgso.runner.signal.getsignal", return_value="original"), mock.patch("scripts.pgso.runner.signal.signal", side_effect=install), ): with self.assertRaisesRegex(PgsoError, "cancelled by SIGTERM"): with cancellation_guard(): handler = installed[signal.SIGTERM] assert callable(handler) handler(signal.SIGTERM, None)
self.assertEqual("original", installed[signal.SIGINT]) self.assertEqual("original", installed[signal.SIGTERM])
def test_child_output_is_visible_before_the_command_finishes(self) -> None: class LiveOutput(io.StringIO): def __init__(self) -> None: super().__init__() self.marker_seen = threading.Event()
def write(self, value: str) -> int: written = super().write(value) if "early output" in value: self.marker_seen.set() return written
with tempfile.TemporaryDirectory(prefix="fx-pgso-runner-") as tmp: root = pathlib.Path(tmp) stdout = LiveOutput() errors: list[Exception] = []
def invoke() -> None: try: run_checked( [ sys.executable, "-c", ( "import time; " "print('early output', flush=True); " "time.sleep(1)" ), ], cwd=root, env=os.environ.copy(), timeout_s=5, log_path=root / "live-before-exit.json", ) except Exception as error: errors.append(error)
with contextlib.redirect_stdout(stdout): worker = threading.Thread(target=invoke, daemon=True) worker.start() observed_while_running = stdout.marker_seen.wait(timeout=0.5) command_still_running = worker.is_alive() worker.join(timeout=3)
self.assertTrue(observed_while_running) self.assertTrue(command_still_running) self.assertFalse(worker.is_alive()) self.assertEqual([], errors)
def test_silent_command_emits_periodic_heartbeats(self) -> None: with tempfile.TemporaryDirectory(prefix="fx-pgso-runner-") as tmp: root = pathlib.Path(tmp) stdout = io.StringIO()
with ( contextlib.redirect_stdout(stdout), mock.patch( "scripts.pgso.runner.COMMAND_HEARTBEAT_SECONDS", 0.05, create=True, ), ): run_checked( [sys.executable, "-c", "import time; time.sleep(0.2)"], cwd=root, env=os.environ.copy(), timeout_s=5, log_path=root / "heartbeat.json", )
output = stdout.getvalue() self.assertRegex( output, r"\[pgso\] command still running after \d+\.\d+s", ) self.assertIn("heartbeat.json", output)
def test_success_streams_command_lifecycle_and_child_output(self) -> None: with tempfile.TemporaryDirectory(prefix="fx-pgso-runner-") as tmp: root = pathlib.Path(tmp) stdout = io.StringIO() stderr = io.StringIO()
with contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr): run_checked( [ sys.executable, "-c", ( "import sys; " "print('live stdout', flush=True); " "print('live stderr', file=sys.stderr, flush=True)" ), ], cwd=root, env=os.environ.copy(), timeout_s=5, log_path=root / "live.json", )
output = stdout.getvalue() self.assertIn("[pgso] command started", output) self.assertIn("live stdout\n", output) self.assertRegex(output, r"\[pgso\] command passed in \d+\.\d{3}s") self.assertLess(output.index("command started"), output.index("live stdout")) self.assertLess(output.index("live stdout"), output.index("command passed")) self.assertIn("live stderr\n", stderr.getvalue())
def test_failure_streams_terminal_status_after_child_output(self) -> None: with tempfile.TemporaryDirectory(prefix="fx-pgso-runner-") as tmp: root = pathlib.Path(tmp) stdout = io.StringIO()
with ( contextlib.redirect_stdout(stdout), self.assertRaisesRegex(PgsoError, "command failed with exit code 7"), ): run_checked( [ sys.executable, "-c", "print('failure detail', flush=True); raise SystemExit(7)", ], cwd=root, env=os.environ.copy(), timeout_s=5, log_path=root / "failure.json", )
output = stdout.getvalue() self.assertIn("failure detail\n", output) self.assertRegex(output, r"\[pgso\] command failed in \d+\.\d{3}s") self.assertLess(output.index("failure detail"), output.index("command failed"))
def test_success_captures_output_and_writes_a_bounded_log(self) -> None: with tempfile.TemporaryDirectory(prefix="fx-pgso-runner-") as tmp: log_path = pathlib.Path(tmp) / "stage.json" env = os.environ.copy() env["PGSO_TEST_SECRET"] = "must-not-be-logged"
result = run_checked( [ sys.executable, "-c", "import sys; print('stdout'); print('stderr', file=sys.stderr)", ], cwd=pathlib.Path(tmp), env=env, timeout_s=5, log_path=log_path, )
self.assertEqual(0, result.returncode) self.assertEqual("stdout\n", result.stdout) self.assertEqual("stderr\n", result.stderr) self.assertGreaterEqual(result.elapsed_seconds, 0) log = json.loads(log_path.read_text()) self.assertEqual(pathlib.Path(sys.executable).name, log["command_class"]) self.assertEqual("passed", log["status"]) self.assertNotIn("environment", log) self.assertNotIn("must-not-be-logged", log_path.read_text())
def test_noisy_output_is_truncated_in_memory_and_in_the_json_log(self) -> None: with tempfile.TemporaryDirectory(prefix="fx-pgso-runner-") as tmp: root = pathlib.Path(tmp) log_path = root / "bounded.json" with ( contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()), mock.patch( "scripts.pgso.runner.MAX_CAPTURED_OUTPUT_CHARS", 128, ), ): result = run_checked( [ sys.executable, "-c", "import sys; print('x' * 4096); print('y' * 4096, file=sys.stderr)", ], cwd=root, env=os.environ.copy(), timeout_s=5, log_path=log_path, )
log = json.loads(log_path.read_text()) self.assertTrue(result.stdout_truncated) self.assertTrue(result.stderr_truncated) self.assertTrue(log["stdout_truncated"]) self.assertTrue(log["stderr_truncated"]) self.assertGreater(log["stdout_total_chars"], 4096) self.assertGreater(log["stderr_total_chars"], 4096) self.assertLess(len(log["stdout"]), 256) self.assertLess(len(log["stderr"]), 256)
def test_interruption_terminates_the_process_group_and_logs_cancellation(self) -> None: with tempfile.TemporaryDirectory(prefix="fx-pgso-runner-") as tmp: root = pathlib.Path(tmp) process = mock.Mock() process.stdout = io.StringIO("partial stdout\n") process.stderr = io.StringIO("") process.poll.return_value = None process.wait.side_effect = KeyboardInterrupt() process.returncode = -2 log_path = root / "cancelled.json"
with ( mock.patch("scripts.pgso.runner.subprocess.Popen", return_value=process), mock.patch("scripts.pgso.runner._terminate_process_group") as terminate, self.assertRaises(KeyboardInterrupt), ): run_checked( ["fixture"], cwd=root, env={}, timeout_s=5, log_path=log_path, )
terminate.assert_called_once_with(process) log = json.loads(log_path.read_text()) self.assertEqual("cancelled", log["status"]) self.assertEqual("partial stdout\n", log["stdout"])
def test_nonzero_exit_raises_after_logging_output(self) -> None: with tempfile.TemporaryDirectory(prefix="fx-pgso-runner-") as tmp: log_path = pathlib.Path(tmp) / "stage.json"
with self.assertRaisesRegex( PgsoError, "command failed with exit code 7", ): run_checked( [ sys.executable, "-c", "import sys; print('before failure'); sys.exit(7)", ], cwd=pathlib.Path(tmp), env=os.environ.copy(), timeout_s=5, log_path=log_path, )
log = json.loads(log_path.read_text()) self.assertEqual(7, log["returncode"]) self.assertEqual("before failure\n", log["stdout"]) self.assertEqual("failed", log["status"])
def test_empty_stderr_policy_fails_closed(self) -> None: with tempfile.TemporaryDirectory(prefix="fx-pgso-runner-") as tmp: log_path = pathlib.Path(tmp) / "stage.json"
with self.assertRaisesRegex(PgsoError, "wrote unexpected stderr"): run_checked( [ sys.executable, "-c", "import sys; print('optimizer warning', file=sys.stderr)", ], cwd=pathlib.Path(tmp), env=os.environ.copy(), timeout_s=5, log_path=log_path, require_empty_stderr=True, )
log = json.loads(log_path.read_text()) self.assertEqual("unexpected_stderr", log["status"]) self.assertEqual("optimizer warning\n", log["stderr"])
def test_missing_executable_fails_closed_and_writes_a_log(self) -> None: with tempfile.TemporaryDirectory(prefix="fx-pgso-runner-") as tmp: log_path = pathlib.Path(tmp) / "stage.json"
with self.assertRaisesRegex(PgsoError, "executable not found"): run_checked( [str(pathlib.Path(tmp) / "missing-command")], cwd=pathlib.Path(tmp), env=os.environ.copy(), timeout_s=5, log_path=log_path, )
log = json.loads(log_path.read_text()) self.assertEqual("spawn_failed", log["status"]) self.assertIsNone(log["returncode"])
def test_timeout_terminates_the_whole_process_group(self) -> None: with tempfile.TemporaryDirectory(prefix="fx-pgso-runner-") as tmp: root = pathlib.Path(tmp) marker = root / "child-survived" log_path = root / "stage.json" script = root / "timeout-fixture.sh" script.write_text( "#!/bin/sh\n" f"(sleep 1; printf survived > {shlex.quote(str(marker))}) &\n" "printf 'started\\n'\n" "sleep 30\n", encoding="utf-8", )
with self.assertRaisesRegex(PgsoError, "timed out after 0.5 seconds"): run_checked( ["/bin/sh", str(script)], cwd=root, env=os.environ.copy(), timeout_s=0.5, log_path=log_path, )
time.sleep(1.2) self.assertFalse(marker.exists()) log = json.loads(log_path.read_text()) self.assertEqual("timed_out", log["status"]) self.assertIn("started\n", log["stdout"])
if __name__ == "__main__": unittest.main()