Skip to content

Remove and rewrite unnecessary aio and Spark cursor tests - #1258

Merged
laughingman7743 merged 4 commits into
masterfrom
test/1245-aio-spark-tests
Oct 18, 2026
Merged

laughingman7743 merged 4 commits into
masterfrom
test/1245-aio-spark-tests

Conversation

@laughingman7743

@laughingman7743 laughingman7743 commented Oct 17, 2026 •

Copy link
Copy Markdown
Member

WHAT

Apply the rows of #1245 to the aio and Spark cursor tests: tests/pyathena/aio/test_connection.py, aio/test_cursor.py, aio/test_result_set.py, aio/test_util.py, aio/test_common.py, aio/{arrow,polars,s3fs,spark}/test_cursor.py, spark/test_async_cursor.py, spark/test_spark_cursor.py, spark/test_common.py, and the merge targets in test_connection.py, test_cursor.py, and tests/pyathena/util.py. It builds on #1253, whose rewritten tests are the merge targets; this diff changes no file under pyathena/.

  • Merge the aio keyword-only, callback, and arraysize tests into the synchronous test_connection.py parametrizations: the 7 aio cursor classes join CURSOR_CLASSES, PAGED_CURSORS, UNCAPPED_CURSORS, and START_CALLBACK_CURSORS, and aio_connect, AioConnection.create, and AioConnection join test_connect_keyword_only. test_cursor_execute_keyword_only and test_cursor_accepts_connection_callbacks become async tests, so they await an aio execute() and close().
  • Merge the aio test_no_result_set_raises and test_cancel_initial copies into aio/test_common.py: test_fetch_without_result_set covers every aio cursor and as_arrow()/as_polars(), and the new test_cancel_without_query covers every aio SQL cursor. test_executemany_failure_closes_result_set_off_event_loop also checks rowcount and the retained query_id.
  • Remove the aio, Arrow, Polars, and S3FS copies of tests that run through shared code (WithResultSet, WithAsyncFetch, the synchronous result sets, the formatter), and fold the execute() is cursor, as_polars/as_arrow, and UNLOAD as_* checks into the remaining query of each class.
  • Rewrite the aio tests that patched private methods to answer the Athena requests offline: answered_requests() (from Remove and rewrite unnecessary top-level cursor, connection, converter, and utility tests #1253) for the execute() options, which checks the StartQueryExecution and ListQueryExecutions parameters, that an execution older than cache_expiration_time is not reused, and that a qmark query with parameters lists nothing; QueryExecutionHistory for the cache search; and a stubbed Glue client for test_glue_request_runs_off_the_event_loop, which checks that the client is built and the request sent off the event loop. cache_search_cursor() moves to tests/pyathena/util.py for the synchronous and aio cache tests, next to the new placeholder_connection().
  • test_execute_cancelled_again_while_start_is_retried checks the attempt count after its default executor has shut down, instead of after asyncio.sleep(0.2).
  • Keep only test_executemany_rowcount[update] and test_executemany_failure[1], and drop the [backoff] cases of the two aio retry-cancellation tests; the synchronous test_retry_api_call_stopped_by_event[backoff] checks the backoff.
  • Spark: run the execute_accepts_none_parameters tests of all three Spark cursors against stubbed calculation requests (offline_spark_cursor() and COMPLETED_CALCULATION_RESPONSES in tests/pyathena/util.py). The session tests construct the cursor instead of calling _wait_for_idle_session/_exists_session, and check the GetSessionStatus and GetSession requests. test_calculate_interrupted_while_starting becomes the final_state parameter of test_execute_interrupted_while_starting, and the AsyncSparkCursor and AioSparkCursor cases of tests of inherited code are dropped. TestAioSparkCursor::test_spark_dataframe runs one small Spark job, print(spark.range(3).count()), instead of the CSV read, groupBy, and saveAsTable that the synchronous test runs.

These rows changed after their cover was re-checked by breaking the behavior in the source:

Row in #1245 Change Reason
aio/test_connection.py::test_connect_keyword_only, test_connection_constructor_keyword_only MERGE into test_connection.py::test_connect_keyword_only The synchronous test called only connect() and Connection; a positional parameter on aio_connect, AioConnection.create, or AioConnection.__init__ passed it. The merged test passes one positional argument, because with two a factory that accepts one still raised TypeError.
aio/test_cursor.py::TestAioCursor::test_executemany_fetch Removed, with its rowcount check moved No other test checked that an aio executemany() of statements without an update count reports rowcount == -1; aio/test_common.py::TestWithAsyncFetch::test_executemany_unknown_rowcount checks it offline for every aio SQL cursor, with no result set left to fetch.
spark/test_common.py::TestSparkBaseCursor::test_async_close_shuts_down_executor_without_terminating_session KEEP The shutdown tests use owned sessions; an AsyncSparkCursor.close() that returned without terminating a supplied session left its executor running and passed them.
aio/spark/test_cursor.py::TestAioSparkCursor::test_context_manager Removed, with a new offline test AioSparkCursor has its own __aexit__, and no test checked that it closes the cursor; the removed test ran print inside async with without checking it either. spark/test_common.py::TestSparkBaseCursor::test_aio_context_exit_terminates_owned_session checks that the owned session is terminated at the end of async with.
spark/test_common.py::TestSparkBaseCursor::test_wait_for_idle_session_waits_until_idle REWRITE, keeping the request count A constructor that returned at the first CREATING status passed a test of only the cursor and the absent termination; the test asserts the three GetSessionStatus requests.

aio/test_result_set.py keeps its ProgrammingError import, which test_fetch_after_failed_fetch uses. For #1244: AioPandasCursor is in the synchronous UNCAPPED_CURSORS, and aio/test_common.py::TestWithAsyncFetch::test_cancel_without_query[AioPandasCursor] exists.

Before After
Test functions 253 164
Collected tests 581 477

Counted over the 13 changed test files at #1253's head (3fa2d7e) and this head. Estimated from the removed and merged queries, the change saves about 80 Athena queries per run (about 30 of them Iceberg statements of the executemany_table fixture), 4 Spark sessions and 9 calculations, including one S3 Parquet write and one Glue table, and the live metadata requests of test_glue_request_runs_off_the_event_loop.

WHY

Closes #1245. Part of #1239: these tests restated the code, patched or asserted private state, waited on sleeps, or repeated a synchronous test through the same code, as described in docs/testing.md.

TEST

Python 3.15.0, at a5b5b34 (on master dc25038, after #1253, #1251, #1254, #1252, and #1256).

  • just format: no changes. just lint: passed.
  • Offline, with --noconftest, placeholder settings, and -W error::RuntimeWarning, the changed files and test_connection.py, test_cursor.py, test_util.py, test_glue.py, tests/pyathena/spark/: 554 passed, 5 skipped; the 6 failed and 109 errors. With tests/pyathena/sqlalchemy/ and tests/pyathena/aio/sqlalchemy/ added for the rebased tests/pyathena/util.py: 1,227 passed, 8 skipped are tests that need AWS or its fixtures.
  • The Test workflow on Ready passed at a5b5b34 (run 39084165463): lint, test, test-sqla, and test-sqla-async. It also passed at ff0ca79, before the rebase onto dc25038 (run 39082697859).
  • After the rebase onto dc25038, the whole tests/pyathena/ offline (--noconftest, placeholder settings): the failures are the same tests that fail offline on master, which need AWS or its fixtures (172 here, 177 on master, none new).
  • Mutation checks, offline with --noconftest and placeholder settings; every source mutation was reverted and git diff -- pyathena/ was empty after each run. Each mutation named in the rows failed the kept, merged, or rewritten test that covers it, except where the table above records a change. Dropping self._rowcount = -1 from the except branch of the aio executemany() changes nothing observable, because the cursor's own _rowcount stays -1 during the loop. A wrong started session ID passes test_close_terminates_session_by_ownership, which compares against cursor.session_id, and fails test_close_does_not_terminate_session_twice and test_wait_for_idle_session_waits_until_idle.
  • Live, once for Remove and rewrite unnecessary top-level cursor, connection, converter, and utility tests #1253 and this PR together, at this PR's c661752 before the rebase (which contains Remove and rewrite unnecessary top-level cursor, connection, converter, and utility tests #1253's 3fa2d7e), with no Test workflow running: uv run --env-file .env pytest -n 4 over the changed test files of both PRs (tests/pyathena/test_{cursor,connection,converter,formatter,options,parser,result_set,util,glue}.py, the changed tests/pyathena/aio/ files and directories, and tests/pyathena/spark/): 965 passed, 1 skipped, in 275 s.

🤖 Generated with Claude Code

assert query_id == "test_query_id"
request_mock.assert_called_once_with(
query="SELECT 1",
"""The options of execute() reach its Athena requests (no AWS)."""

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Self-review round one (implementation behavior, /code-review high and a /simplify review), base 3fa2d7e (#1253's head, the PR base), head 6ae81ae. Reviewed at 0a4056e; repairs in 6ae81ae. FINDINGS, repaired unless noted.

The rewritten options test did not observe cache_expiration_time: with an empty listing, dropping it from the aio _execute() passed. The listing now holds the same query from an hour ago, which cache_expiration_time=100 must not reuse; passing 0 instead fails the test (mutation check, offline).

throttle_metadata_api(aio_cursor.connection.client, monkeypatch)
loop_thread = threading.get_ident()
async def test_glue_request_runs_off_the_event_loop(self, monkeypatch):
"""The Glue client and request of a throttled lookup do not block the loop (no AWS)."""

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Round one: the rewrite dropped the check that the Glue client is built off the event loop. session.client now records its thread and checks that it builds glue; building the client on the loop before asyncio.to_thread fails the test.

}

@pytest.mark.parametrize("cursor_class", SYNC_SPARK_CURSOR_CLASSES)
with pytest.raises(OperationalError, match="session failure reason"):

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Round one: the supplied-session test now matches the reason Athena reports and checks the GetSession request for supplied-session; skipping _exists_session fails it.

Deferred: the format of the Spark session error message (Session: <id> is <STATE>.) has no test after test_wait_for_idle_session_raises_without_reason is removed; the message is not documented, and docs/testing.md does not count undocumented message text as a contract.

with answered_requests(offline_connection.client):
async with _opened_cursor(offline_connection, cursor_class) as cursor:
result = None
with contextlib.suppress(OperationalError):

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Round one: the suppressed OperationalError also covered the Async cursors' future.result(), so a failing future passed. The future is now waited for outside the suppression.

Also from /simplify: AIO_CURSOR_CLASSES is a literal list instead of a module-name test; placeholder_connection() in tests/pyathena/util.py replaces the copy in aio/test_cursor.py and offline_spark_cursor().

AioS3FSCursor,
]

_WHOLE_RESULT_METHODS = {

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Round one: test_fetch_without_result_set skipped as_arrow()/as_polars() through hasattr, so removing either method passed; _WHOLE_RESULT_METHODS names them per class, with the No result set. match of the fetch checks.

SYNC_SPARK_CURSOR_CLASSES = [SparkCursor, AsyncSparkCursor]
# AsyncSparkCursor inherits _calculate and _start_execution from SparkCursor's bases;
# the parameter keeps the test IDs.
CALCULATE_CURSOR_CLASSES = [SparkCursor]

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Round one: the six calculate tests keep a one-class parameter, now CALCULATE_CURSOR_CLASSES, with a comment: it keeps their test IDs, as the test conventions require; test_executemany_failure[1] keeps its ID the same way.

Not changed: aio/test_connection.py keeps its own list of the aio cursors for its two remaining tests; the offline connections are not closed, as in the neighbouring offline tests.

Comment thread tests/pyathena/util.py
}


def placeholder_connection(connection_class=Connection):

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Self-review round two (claims, callers, AWS operation, evidence), base 3fa2d7e, head 6ae81ae. FINDINGS in the PR description only; corrected.

Claims checked:

  • No file under pyathena/ changes (git diff --name-only lists only tests/).
  • 253 → 161 test functions and 581 → 469 collected tests over the 13 changed test files, recounted at the head.
  • 4 Spark sessions: the aio_spark_cursor and async_spark_cursor fixtures start one per test, for TestAioSparkCursor::test_spark_sql, test_session_ownership, test_context_manager (own connection), and TestAsyncSparkCursor::test_session_ownership.
  • 9 calculations: 2 fewer in the rewritten test_spark_dataframe, 2 in test_spark_sql, 2 in each test_session_ownership, 1 in test_context_manager.
  • One S3 Parquet write and one Glue table: the removed saveAsTable.
  • About 80 Athena queries is an estimate from the removed queries and the executemany_table fixture (3 statements per case); no run measured it.
  • The description did not say that the options test checks cache_expiration_time and the Glue test the client's thread; added. It now also names the rows that Remove, merge, and trim unnecessary tests in the pandas tests #1244 needs.

Callers: no source changes. The aio execute() keyword-only contract follows docs/usage.md "Keyword arguments in PyAthena 4.0", which names aio_connect() and the aio cursors.

AWS operation: the live aio and Spark tests are fewer, and TestAioSparkCursor::test_spark_dataframe asserts the calculation properties on a print calculation. That dpu_execution_in_millis, progress, and the result URIs are set for it is checked only by the pending live run.

assert cursor.rowcount == 0

@pytest.mark.parametrize("cursor_class", _ALL_CURSOR_CLASSES)
async def test_executemany_unknown_rowcount(self, cursor_class):

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Independent review (relayed): Codex CLI codex exec -s read-only, codex-cli 0.163.0, model gpt-6-astra, session 01a14a13-763c-7782-9166-8550f135be47. Static review of a detached snapshot at 6ae81ae against 3fa2d7e; no tests, builds, or GitHub access; the prompt left out the PR number, description, commit messages, and earlier findings. The snapshot was unchanged afterwards. Result: FINDINGS (4), all repaired in c661752.

  1. P2: removing TestAioCursor::test_executemany_fetch left no test that an aio executemany() whose statement has no update count, such as a SELECT, reports rowcount == -1; the synchronous test runs separate code. Repaired: TestWithAsyncFetch::test_executemany_unknown_rowcount covers every aio SQL cursor offline; making the unknown branch 0 fails it.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Independent follow-up (relayed): Codex CLI, gpt-6-astra, session 01a14a1d-c392-7833-87bf-4b077a73067a, static review of c661752 against 6ae81ae with the repair diff. All four points are closed. One P3 finding: test_executemany_unknown_rowcount repeats the same faked execute() for six classes that share executemany(). Deferred with reasons: the neighbouring TestWithAsyncFetch tests are parametrized over the same classes, and #1245 keeps that parametrization because each case checks which executemany() the class resolves (its MRO). With one class, a class whose bases put another executemany() first would pass.

with pytest.raises(ProgrammingError, match=r"No result set\."):
await fetch()
# The Arrow and Polars cursors also convert the whole result.
for name in _WHOLE_RESULT_METHODS.get(cursor_class, ()):

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. P3: the No result set. match that round one added to the as_arrow()/as_polars() checks is undocumented message text, which the removed tests did not check. Repaired: those checks assert ProgrammingError only; the fetch checks keep their earlier match.

"""
),
description="test description",
"print(spark.range(3).count())", description="test description"

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. P2: print('hello') runs no Spark job, so dpu_execution_in_millis and progress may be empty for it. Repaired: the calculation runs print(spark.range(3).count()), a Spark job, and expects 3; the live run checks the properties.

connection.client.terminate_session.assert_not_called()

@pytest.mark.parametrize("cursor_class", SPARK_CURSOR_CLASSES)
@pytest.mark.parametrize("cursor_class", SYNC_SPARK_CURSOR_CLASSES)

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. P2: removing test_async_close_shuts_down_executor_without_terminating_session left the executor shutdown of an AsyncSparkCursor on a supplied session without a test; the shutdown tests use owned sessions. Repaired: the test is restored, identical to the base (test_common.py:463); returning from close() without terminating fails it.

cursor.execute = execute
await cursor.executemany("SELECT 1", [None, None, None])

assert cursor.rowcount == -1

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Self-review of the repair c661752, both perspectives; narrow scope by git range-diff 3fa2d7ef..6ae81ae6 3fa2d7ef..c6617522 (one new commit). Base 3fa2d7e, head c661752. CLEAN.

Round one: the new offline rowcount test sets each result set closed, so executemany() sends no close to a thread, and the three counts (1, -1, 1) also check that a later known count does not hide an earlier unknown one. The restored Spark test is identical to the base. Offline: 553 passed; both repair mutations fail their tests.

Round two: the description now lists the moved rowcount check and the kept Spark test, the Spark job, and 163 test functions and 476 collected tests.

@laughingman7743
laughingman7743 marked this pull request as ready for review October 17, 2026 14:21
Base automatically changed from test/1243-top-level-tests to master October 18, 2026 00:06
@laughingman7743
laughingman7743 marked this pull request as draft October 18, 2026 00:18
@laughingman7743
laughingman7743 force-pushed the test/1245-aio-spark-tests branch from c661752 to ff0ca79 Compare October 18, 2026 00:18
else:
connection.client.terminate_session.assert_not_called()

async def test_aio_context_exit_terminates_owned_session(self):

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rebase and repair after #1253 was merged, both self-review perspectives. Base: the merge-base with master, 3a41a84 (#1253's merge). Head: ff0ca79. Scope: git range-diff 3fa2d7ef..c6617522 3a41a84b..ff0ca797. The three earlier commits are unchanged except for one import conflict in tests/pyathena/util.py: #1251 added from types import SimpleNamespace and SQLAlchemy imports next to this PR's MagicMock. The new commit ff0ca79 adds two checks after the maintainer's question. CLEAN.

Round one:

  • TestSparkBaseCursor::test_aio_context_exit_terminates_owned_session checks that async with ends with the owned session's TerminateSession and not before. An empty AioSparkCursor.__aexit__ fails it (offline mutation); before, no test failed for that mutation.
  • test_executemany_unknown_rowcount also asserts that no result set is left. Skipping the final reset of executemany() fails it.
  • The rebase merges Remove and rewrite unnecessary SQLAlchemy dialect tests #1251's tests/pyathena/util.py helpers with this PR's. Offline, the changed files and tests/pyathena/{sqlalchemy,aio/sqlalchemy}: 1,227 passed; the failures and errors are the tests that need AWS.

Round two: the description names the moved Spark row, the result-set check, the rebase base, 164 test functions and 477 collected tests. The live run predates the rebase; the Test workflow on Ready covers the rebased head.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Independent follow-up (relayed): Codex CLI, gpt-6-astra, session 01a14c61-a7b9-7732-b47c-769cbf3ee983. Static review of a detached snapshot at ff0ca79, given the range-diff against the reviewed c661752 and the new commit's diff; no tests, builds, or GitHub access. CLEAN. The 14 test files match the reviewed versions before the new commit, and tests/pyathena/util.py keeps both PRs' helpers and imports. The new __aexit__ test fails if the close is omitted or not awaited, and the result-set assertion fails if the final result set is kept. Neither has an unbounded wait.

laughingman7743 and others added 4 commits October 18, 2026 09:44
Apply the rows of #1245 to the aio connection, cursor, result set, utility,
and DataFrame cursor tests and to the Spark cursor tests. The aio keyword,
callback, and arraysize tests become parameters of the synchronous
connection tests, the aio cursors' fetch-without-result and cancel tests
become offline tests over every aio cursor, and the copies that run
through shared code are removed.

The rewritten tests answer the Athena requests offline: the request
parameters of execute() options, the cache search through
QueryExecutionHistory, the Glue fallback thread, and the Spark
calculations. A cancelled start request is checked after its executor
thread ends instead of after a sleep, and the Spark session tests drive
the cursor constructor instead of its private helpers.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… missed

The aio options test lists an hour-old execution, so it fails if
cache_expiration_time does not reach the cache search. The Glue test also
records the thread that builds the client. The supplied-session test
checks the GetSession request and the reason from Athena. An Async
cursor's future is waited for outside the suppressed OperationalError.
The placeholder connection and the aio cursor list are defined once, and
the Arrow and Polars whole-result methods are named per class.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Check offline that an executemany() with an unknown update count reports
-1, keep the test that closing an AsyncSparkCursor on a supplied session
shuts down its executor, and run a Spark job in the aio calculation test,
so that its DPU time and progress are reported. Leave out the message of
the whole-result errors, which is not documented.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ter executemany

AioSparkCursor has its own __aexit__, which no test checked: an owned
session is now terminated at the end of async with. The aio executemany()
test also checks that no result set is left to fetch, which the removed
test_executemany_fetch tests checked.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@laughingman7743
laughingman7743 force-pushed the test/1245-aio-spark-tests branch from ff0ca79 to a5b5b34 Compare October 18, 2026 00:49
Comment thread tests/pyathena/util.py
from pyathena.model import AthenaCalculationExecutionStatus, AthenaQueryExecution
from pyathena.util import override, retry_api_call
from pyathena.spark.common import SparkBaseCursor
from pyathena.util import RetryConfig, override, retry_api_call

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rebase onto master dc25038 (#1254, #1252, #1256 merged), both self-review perspectives. Base: merge-base dc25038. Head: a5b5b34. Scope: git range-diff 3a41a84b..ff0ca797 dc25038f..a5b5b349; the four commits are unchanged except for one import conflict in tests/pyathena/util.py, where master added override next to this PR's RetryConfig and SparkBaseCursor. CLEAN.

Round one: tests/pyathena/util.py has no duplicate top-level names after the merge. #1256 removes aio pandas tests that cite this PR's UNCAPPED_CURSORS and test_cancel_without_query; both are in this PR. The whole tests/pyathena/ offline fails only the tests that also fail offline on master because they need AWS (172 here, 177 on master, none new). just lint passed.

Round two: the description records the rebase base and these results. The Test workflow on Ready runs on this head.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Remove, merge, and rewrite duplicate tests in the asyncio and Spark tests

1 participant