Keep a PySpark Test Flood Out of Your Agent's Context Window

Published 7 min read

This started with a broken Spark fixture on an ordinary day. Every test that needed Spark failed on the same setup error, and my coding agent had to read all of it before it could fix a thing. Reproduced and measured, that flood was 19 failing tests and 797 lines of output, almost all of it one Py4J stack repeated.

That's the PySpark tax. A failing Spark test drags a long trace across the Python frames and the JVM underneath, so one failure is already a screenful, and a shared mistake multiplies it across every Spark-backed test.

You would scroll past it. An agent can't: the flood fills its context window and crowds out the room it needs to work. And it really is Spark, not pytest. In that run, the only tests that passed were the ones that never start Spark, the config, JSON, and schema checks. So the fix isn't quieting pytest; it's keeping Spark's flood off the agent.

One broken fixture, 19 failing tests Same run, two views of the output raw pytest output 797 lines the same Py4J stack, 19 times over wrapper bounded digest ~70 lines Result: FAILED ExitCode: 1 Collected: 30 Passed: 11 Errors: 19 --- FAILING TESTS (19 of 19) --- test_cleaning_utils.py::TestCleanStr test_datetime_transforms.py::Test... ... 17 more runnable node ids ... --- SIGNATURE 1 of 1 (19 tests) --- py4j.protocol.Py4JJavaError: ... Example: test_cleaning_utils.py::... <traceback excerpt, head+tail trimmed> ... 15 lines trimmed, see log ... one cause = one block. full log on disk. One broken fixture, 19 failing tests Same run, two views of the output raw pytest output 797 lines the same Py4J stack, 19 times over wrapper bounded digest ~70 lines Result: FAILED ExitCode: 1 Collected: 30 Passed: 11 Errors: 19 --- FAILING TESTS (19 of 19) --- test_cleaning_utils.py::TestCleanStr test_datetime_transforms.py::Test... ... 17 more runnable node ids ... --- SIGNATURE 1 of 1 (19 tests) --- py4j.protocol.Py4JJavaError: ... Example: test_cleaning_utils.py::... <traceback excerpt, head+tail trimmed> ... 15 lines trimmed, see log ... one cause = one block. full log on disk.

Here's the wrapper I built for my own agent. It runs pytest for you, writes the complete stdout and stderr to a log file, and prints a compact digest instead of the raw flood.

The digest is built from the JUnit XML that pytest writes, not by scraping stdout, so heavy interleaved Spark and JVM logging can't corrupt it. It carries four things:

  • the exit code (PASSED, FAILED, ERROR, and so on)
  • the counts: collected, passed, failed, errored
  • the runnable failing node ids (more on those below)
  • the failures themselves, deduplicated by a normalized signature, so many tests failing with one cause collapse into a single block with one trimmed excerpt

A "node id" is just pytest's address for a single test: the file, the class, and the test name, joined by ::, like tests/test_cleaning_utils.py::TestCleanStr::test_clean_str_logic. That matters here because the digest prints them as runnable pytest targets, so one works as the next command without editing.

On that same run, 797 raw lines became about 70:

=== PYSPARK TEST SUMMARY ===
Result: FAILED
ExitCode: 1
Collected: 30  Passed: 11  Failed: 0  Errors: 19  Skipped: 0
TimeSec: 2.66
Log:   .pyspark-test-logs/pytest-<stamp>.log
JUnit: .pyspark-test-logs/pytest-<stamp>.xml
--- FAILING TESTS (showing 19 of 19) ---
tests/test_cleaning_utils.py::TestCleanStr::test_clean_str_logic  -  failed on setup with "py4j.protocol.Py4JJavaError: ..."
tests/test_datetime_transforms.py::TestCombineDateTimeBk::test_combine_logic  -  failed on setup with "py4j.protocol.Py4JJavaError: ..."
... (17 more runnable node ids) ...
--- SIGNATURE 1 of 1  (19 tests) ---
failed on setup with "py4j.protocol.Py4JJavaError: ..."
Example: tests/test_cleaning_utils.py::TestCleanStr::test_clean_str_logic
<traceback excerpt, head+tail trimmed; "... (15 lines trimmed, see log) ..." >

To see what those 70 lines are standing in for, here is one of the nineteen raw blocks the digest replaced:

___________________ ERROR at setup of TestCleanStr.test_clean_str_logic ___________________

    @pytest.fixture(scope="session")
    def spark():
>       session = SparkSession.builder.master("local[oops]").appName("tests").getOrCreate()

tests/conftest.py:14:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
.venv/lib/site-packages/pyspark/sql/session.py:497: in getOrCreate
    sc = SparkContext.getOrCreate(sparkConf)
.venv/lib/site-packages/pyspark/context.py:515: in getOrCreate
    SparkContext(conf=conf or SparkConf())
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

py4j.protocol.Py4JJavaError: An error occurred while calling None.org.apache.spark.api.java.JavaSparkContext.
: org.apache.spark.SparkException: Could not parse Master URL: 'local[oops]'
    at org.apache.spark.SparkContext$.org$apache$spark$SparkContext$$createTaskScheduler(SparkContext.scala:3194)
    at org.apache.spark.SparkContext.<init>(SparkContext.scala:577)
    at org.apache.spark.api.java.JavaSparkContext.<init>(JavaSparkContext.scala:58)
    at java.base/jdk.internal.reflect.NativeConstructorAccessorImpl.newInstance0(Native Method)
    ... 38 more JVM frames ...

All those failures share one cause, so they collapse into a single SIGNATURE 1 of 1 block instead of a wall of tracebacks. The full detail still lives in the log file for the moment you actually need it.

That bound is by construction. The digest is one signature block plus a capped list of node ids, so it stays around 70 lines whether 19 tests share the cause or 1900 do. I measured it at 19; the shape is what holds it flat above that.

You install the skill into your project, and your agent runs the wrapper instead of calling pytest raw. The digest shows every failure at once, grouped by cause, so the agent sees the whole picture and can fix a shared cause across the class in one pass, not test by test. Since these are native pytest node ids, the agent can copy them directly from the digest to rerun the specific tests it just modified.

The wrapper doesn't fix anything. It runs pytest and prints the digest; the agent reads it and makes the fix. As the skill file tells the agent: the wrapper reports, it does not change your code.

It installs under a neutral .agents directory by default and works with whatever agentic environment you use (Claude Code, Codex, Omnigent, Cursor, etc.), with a one-line override per tool if you want auto-discovery.

  • It runs pytest, nothing more. It does not deploy or run Databricks bundles, it is not a CI runner, and it is not a pytest plugin. It also adds no dependencies of its own (Python 3.9+, standard library), but it runs your suite, so you still need pytest and PySpark installed and a Spark that can start on the machine.
  • Signature grouping needs the JUnit XML. The wrapper always asks pytest to write it, but if a suite disables that report or pytest dies before writing it, the digest falls back to a bounded tail of the log instead of grouped signatures. It never crashes on missing XML.
  • The signature strips numbers. Two assertions that differ only by an expected value group together. That is the point when you are collapsing a flood, but it means the digest points you to the shared cause, not to every distinct value. The log keeps the full per-test detail.

It's packaged as a skill: a SKILL.md the agent reads plus the wrapper script, installed into your project with one command.

databricks bundle init https://github.com/vmariiechko/databricks-bundle-template \
  --template-dir assets/pyspark-test-runner

You'll be prompted for a target_dir (default .agents). The asset lives in my databricks-bundle-template repo.

One honest pointer: this wrapper is deliberately narrow, it only knows about PySpark test runs. If you want to compress everything your agent reads, tool outputs, logs, and files, not just tests, there is a general-purpose tool for that in Headroom. This is the test-shaped version of the same idea: less noise reaching the model, but built around pytest's structure instead of generic compression.

Copy it, point your agent at it, and the next time a fixture takes out half your suite, your agent reads 70 lines instead of 797.