From 7e8f308096da94b5b71436b6e5255d1bf8cc479e Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Thu, 9 Jul 2026 04:59:11 +0000 Subject: [PATCH 1/2] [INFRA] Document Python test base classes in AGENTS.md ### What changes were proposed in this pull request? Add a `## Python Test Base Classes` section to `AGENTS.md` (`CLAUDE.md` is a symlink to it), right after the existing `## Scala Test Base Classes` section and mirroring its structure: a short intro, an ASCII "adds capability on top of the previous" ladder, and a scope/base/notes table. It covers: - the classic ladder `unittest.TestCase` -> `PySparkBaseTestCase` -> `ReusedPySparkTestCase` -> `ReusedSQLTestCase` -> `PandasOnSparkTestCase`, including the fresh-vs-shared `SparkContext` distinction (`PySparkTestCase` vs `ReusedPySparkTestCase`); - a `### Spark Connect test bases` subsection for `PlanOnlyTestFixture`, `ReusedConnectTestCase`, and `ReusedMixedTestCase`, noting the `should_test_connect` auto-skip; - a `### Mixins and helpers` subsection for `SQLTestUtils`, `PySparkErrorTestUtils`, the standalone `assertDataFrameEqual` / `assertSchemaEqual`, and the off-ladder domain bases (`SparkSessionTestCase`, `MLlibTestCase`, `PySparkStreamingTestCase`). ### Why are the changes needed? `AGENTS.md` already guides contributors (and AI agents) to pick the lowest suitable Scala test base class, but has no equivalent for PySpark. The Python test bases under `python/pyspark/testing/` form a similar ladder that is easy to get wrong (e.g. subclassing a heavier base than needed, or missing the Connect-specific bases). Documenting them fills the adjacent gap. ### Does this PR introduce _any_ user-facing change? No. Documentation only. ### How was this patch tested? No code changes. Verified the described class hierarchy and capabilities against the sources under `python/pyspark/testing/` (`utils.py`, `sqlutils.py`, `connectutils.py`, `pandasutils.py`, `mlutils.py`, `mllibutils.py`, `streamingutils.py`), and spot-checked that only em-dashes (U+2014) were added as non-ASCII, consistent with the rest of the file. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 4.8) --- AGENTS.md | 41 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/AGENTS.md b/AGENTS.md index a83ce90bbaaec..4c269a16bc2f8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -47,6 +47,47 @@ When writing a new Scala test suite, pick the lowest base class that provides wh | `SharedSparkSession` | `sql/core` | Already extends `QueryTest` for historical reasons, but still mix in `QueryTest` explicitly, e.g. `class X extends QueryTest with SharedSparkSession`. Default for tests under `sql/core`. | | `TestHiveSingleton` | `sql/hive` | Mixed in alongside `QueryTest`, e.g. `class X extends QueryTest with TestHiveSingleton`. Used by tests under `sql/hive`. | +## Python Test Base Classes + +PySpark tests use the stdlib `unittest` framework: every suite subclasses `unittest.TestCase` (run via `python/run-tests`, see below). As with Scala, pick the lowest base that provides what the test actually needs. The bases live under `python/pyspark/testing/` and each adds capability on top of the previous: + + unittest.TestCase (stdlib) + <- PySparkBaseTestCase (pyspark.testing.utils) + <- ReusedPySparkTestCase (pyspark.testing.utils) + <- ReusedSQLTestCase (pyspark.testing.sqlutils) + <- PandasOnSparkTestCase (pyspark.testing.pandasutils) + +| Test scope | Base | Notes | +|------------|------|-------| +| Plain Python — no Spark | `unittest.TestCase` | Use the stdlib base directly. `PySparkBaseTestCase` is the same thing plus a SIGTERM fault-handler dump enabled when `PYSPARK_TEST_TIMEOUT` is set; subclass it only when you want that. | +| RDD / `SparkContext` — no `SparkSession` | `PySparkTestCase` (fresh) or `ReusedPySparkTestCase` (shared) | Both create a `SparkContext("local[4]")`. `PySparkTestCase` makes a new one per test (isolation); `ReusedPySparkTestCase` shares one per class (faster, the usual choice) and adds `quiet()` and an overridable `conf()` / `master()`. | +| SQL / DataFrame — needs a `SparkSession` | `ReusedSQLTestCase` | The workhorse for classic tests under `python/pyspark/sql`. Adds a shared `cls.spark`, sample `cls.df` / `cls.testData`, and mixes in `SQLTestUtils` + `PySparkErrorTestUtils`. Classic (non-Connect) mode. | +| pandas API on Spark | `PandasOnSparkTestCase` | Extends `ReusedSQLTestCase` with Arrow enabled and pandas-on-Spark assertions (`PandasOnSparkTestUtils`); `ComparisonTestBase` builds on it. | + +### Spark Connect test bases + +Spark Connect suites live in `pyspark.testing.connectutils` and are auto-skipped (via `should_test_connect`) when Connect dependencies are missing: + + PySparkBaseTestCase (pyspark.testing.utils) + <- PlanOnlyTestFixture (pyspark.testing.connectutils) + <- ReusedConnectTestCase (pyspark.testing.connectutils) + <- ReusedMixedTestCase (pyspark.testing.connectutils) + +| Test scope | Base | Notes | +|------------|------|-------| +| Plan / proto construction — no server | `PlanOnlyTestFixture` | Uses a `MockRemoteSession`; builds and inspects plans without a running Connect server. For proto / plan-shape assertions. | +| Connect DataFrame — real session | `ReusedConnectTestCase` | The Connect analog of `ReusedSQLTestCase`; starts a session via `.remote(...)` (honoring `SPARK_CONNECT_TESTING_REMOTE`, default `local[4]`). Mixes in `SQLTestUtils` + `PySparkErrorTestUtils`. | +| Classic + Connect side by side | `ReusedMixedTestCase` | Extends `ReusedConnectTestCase`; exposes both a classic `spark` and a Connect `connect` session for parity checks (`compare_by_show`, `both_conf`). Requires JVM access. | + +### Mixins and helpers + +These are combined with a base above rather than used on their own: + +- `SQLTestUtils` — context managers `sql_conf`, `table`, `temp_view`, `view`, `database`, `function`, `temp_func`, `temp_env`; assumes `self.spark`. Already mixed into `ReusedSQLTestCase` and `ReusedConnectTestCase`. +- `PySparkErrorTestUtils` — `check_error(...)` to assert on a `PySparkException`'s error class and message parameters. The Python counterpart of Scala's `checkError`. +- `assertDataFrameEqual` / `assertSchemaEqual` (public API, from `pyspark.testing`) — standalone assertion functions that work in any test, no particular base class required. +- Domain bases outside the main ladder: `SparkSessionTestCase` (`pyspark.testing.mlutils`, for `ml`), `MLlibTestCase` (`pyspark.testing.mllibutils`), and `PySparkStreamingTestCase` (`pyspark.testing.streamingutils`, DStreams). + ## Build and Test Build and tests can take a long time. If the user explicitly asked to run tests, run them. Otherwise (you are running tests on your own to verify a change), first ask the user if they have more changes to make. From f2fb5bc4c1436770e7e9e516f32e68eca6a4b402 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Thu, 9 Jul 2026 05:06:23 +0000 Subject: [PATCH 2/2] [INFRA] Clarify ReusedMixedTestCase is for classic-vs-Connect comparison Sharpen the ReusedMixedTestCase row: state that its purpose is to directly compare classic vs Connect, exposing a classic `self.spark` and a Connect `self.connect` so a test can run the same operation on each and assert they agree. Generated-by: Claude Code (Opus 4.8) --- AGENTS.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/AGENTS.md b/AGENTS.md index 4c269a16bc2f8..433cff6ccacd9 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -77,7 +77,7 @@ Spark Connect suites live in `pyspark.testing.connectutils` and are auto-skipped |------------|------|-------| | Plan / proto construction — no server | `PlanOnlyTestFixture` | Uses a `MockRemoteSession`; builds and inspects plans without a running Connect server. For proto / plan-shape assertions. | | Connect DataFrame — real session | `ReusedConnectTestCase` | The Connect analog of `ReusedSQLTestCase`; starts a session via `.remote(...)` (honoring `SPARK_CONNECT_TESTING_REMOTE`, default `local[4]`). Mixes in `SQLTestUtils` + `PySparkErrorTestUtils`. | -| Classic + Connect side by side | `ReusedMixedTestCase` | Extends `ReusedConnectTestCase`; exposes both a classic `spark` and a Connect `connect` session for parity checks (`compare_by_show`, `both_conf`). Requires JVM access. | +| Classic + Connect side by side | `ReusedMixedTestCase` | Extends `ReusedConnectTestCase`. For directly comparing classic vs Connect: it exposes a classic `self.spark` and a Connect `self.connect` so a test can run the same operation on each and assert they agree (`compare_by_show`, `both_conf`). Requires JVM access. | ### Mixins and helpers