Skip to content

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

Closed
#1258
Closed
Remove, merge, and rewrite duplicate tests in the asyncio and Spark tests#1245
#1258
Labels
enhancement

Description

@laughingman7743
laughingman7743
opened on Oct 17, 2026

Summary

Part of #1239. Files: aio/test_connection.py, aio/test_cursor.py, aio/test_result_set.py, aio/test_util.py, aio/arrow/test_cursor.py, aio/polars/test_cursor.py, aio/s3fs/test_cursor.py, aio/spark/test_cursor.py, spark/test_async_cursor.py, spark/test_spark_cursor.py, spark/test_common.py (aio/test_common.py is all KEEP). At master 59169c5 with the cross-check corrections: 224 test functions: 102 KEEP, 65 REMOVE, 29 MERGE, 28 REWRITE. One REMOVE (aio/test_util.py::test_async_retry_api_call_stop_when) is deleted by #1238 and is not listed, so the tables hold 64 REMOVE. Saved per run: about 85 Athena queries (about 30 of them Iceberg DDL/DML from the executemany_table fixture, and 2 UNLOADs), 4 Spark sessions and about 8 calculations (with one S3 Parquet write and one Glue table), and 1 live Glue request.

Paths are relative to tests/pyathena/. [original] marks a test that existed at 21cbf05 (v3.15.0).

Order

  • After Break result sets after failed or cancelled fetches and drop the cross-thread guards #1238. It deletes aio/test_util.py::test_async_retry_api_call_stop_when and the stop_when parameter, and the KEEP tests aio/test_result_set.py::TestAthenaAioResultSet::test_fetch_counts_rows_taken_before_error, test_fetch_with_close_during_request, test_fetch_with_close_during_retries. It adds aio/test_result_set.py::TestAthenaAioResultSet::test_fetch_after_failed_fetch, test_fetch_after_cancelled_fetch, aio/test_common.py::TestWithAsyncFetch::test_cancelled_fetch_detaches_result_set and four aio/test_util.py::test_run_in_thread_or_close* tests. Rows in aio/test_result_set.py, aio/test_util.py and the new aio/test_common.py test are applied on top of it.
  • With or after the top-level sub-issue, whose MERGE/REWRITE targets are the covers here:
    • aio/test_connection.py::test_empty_work_group_disables_environment_fallback → test_connection.py::test_environment_fallback (its MERGE keeps the Connection work_group="" case).
    • The aio keyword-only and callback tests merge into the rewritten test_connection.py::test_connect_keyword_only, test_cursor_constructor_keyword_only, test_cursor_execute_keyword_only and test_cursor_accepts_connection_callbacks.
    • aio/test_cursor.py::TestAioCursor::test_binary_parameter → test_cursor.py::TestCursor::test_binary_parameter (REWRITE to one query).
    • The aio arraysize tests merge into test_connection.py::TestConnection::test_cursor_arraysize*. The top-level sub-issue keeps TestConnection::test_cursor_arraysize_over_page_size for this merge, and the merged test must keep the AioCursor/AioDictCursor cases, because AioCursor has its own arraysize setter (pyathena/aio/cursor.py:96) and this sub-issue removes TestAioCursor::test_invalid_arraysize.
  • The pandas sub-issue depends on rows here: its aio/pandas/test_cursor.py::TestAioPandasCursor::test_invalid_arraysize removal needs the arraysize MERGE in section 1 (which adds AioPandasCursor to the sync UNCAPPED_CURSORS), and its test_cancel_initial removal needs the new aio/test_common.py::TestWithAsyncFetch::test_cancel_without_query from section 2.
  • Within this sub-issue: the test_cancel_initial copies are removed only after test_cancel_without_query exists; TestAioCursor::test_executemany and test_executemany_fetch land with the test_executemany_rowcount REWRITE that keeps [update]; spark/test_common.py::TestSparkBaseCursor::test_wait_for_idle_session_raises_on_failure_state lands with the REWRITE of test_init_terminates_new_session_in_failure_state, which keeps the three states, the exception type and the termination.

Findings

1. aio/test_connection.py

test change criteria reason covered by
test_empty_work_group_disables_environment_fallback REMOVE 2 aio_connect/AioConnection.create forward kwargs to Connection.__init__, where the environment fallback lives. test_connection.py::test_environment_fallback (Connection, work_group="")
test_connect_keyword_only, test_connection_constructor_keyword_only REMOVE 2, 6 aio_connect(**kwargs), create(cls, **kwargs) and AioConnection.__init__(self, **kwargs) are keyword-only by their signatures and forward to Connection. test_connection.py::test_connect_keyword_only (rewritten)
test_cursor_constructor_keyword_only, test_cursor_execute_keyword_only, test_cursor_unknown_constructor_keyword MERGE 2 Line-for-line copies of the sync tests that need no await; add the 7 aio classes to the sync CURSOR_CLASSES. into test_connection.py::test_cursor_constructor_keyword_only, test_cursor_execute_keyword_only (rewritten), test_cursor_unknown_constructor_keyword
test_cursor_accepts_connection_callbacks MERGE 1, 2 Copy of the sync test that asserts the private _on_start_query_execution/_on_poll; merge into the sync parametrization and close the cursor explicitly. into test_connection.py::test_cursor_accepts_connection_callbacks (rewritten)
TestAioConnection::test_cursor_arguments_override_cursor_kwargs REMOVE 2 Connection.cursor() is inherited unchanged; the sync test checks the same merge with more keys. test_connection.py::TestConnection::test_cursor_arguments_override_cursor_kwargs
TestAioConnection::test_cursor_arraysize, test_cursor_arraysize_over_page_size, test_cursor_arraysize_uncapped, test_cursor_arraysize_not_positive MERGE 2 Same bodies as the sync tests; add AioCursor/AioDictCursor to the sync PAGED_CURSORS and AioArrowCursor, AioPandasCursor, AioPolarsCursor, AioS3FSCursor to the sync UNCAPPED_CURSORS. into test_connection.py::TestConnection::test_cursor_arraysize, test_cursor_arraysize_over_page_size, test_cursor_arraysize_uncapped, test_cursor_arraysize_not_positive
TestAioConnection::test_internal_cursor_leaves_out_default_cursor_options REMOVE 1, 2 The private Connection._internal_cursor() is inherited; only the class names differ from the sync copy. test_connection.py::TestConnection::test_internal_cursor_leaves_out_default_cursor_options

2. aio/test_cursor.py

test change criteria reason covered by
TestAioCursor::test_binary_parameter (7 cases) REMOVE 2, 4 Binary formatting is the shared formatter, and the sync test has the same 7 cases (7 queries). test_cursor.py::TestCursor::test_binary_parameter (rewritten to one query)
TestAioCursor::test_async_iterator REMOVE 2, 4 __anext__ is WithAsyncFetch's and is tested offline for AioCursor. aio/test_common.py::TestWithAsyncFetch::test_async_iteration[AioCursor], TestAioCursor::test_fetchone
TestAioCursor::test_execute_returns_self MERGE 4 One extra query for assert result is cursor. into TestAioCursor::test_fetchone
TestAioCursor::test_execute_internal_options_passthrough REWRITE 1 Patches the private _build_start_query_execution_request/_find_previous_query_id; drive execute() with a mocked client and assert the StartQueryExecution request (WorkGroup, OutputLocation, ResultReuseConfiguration) and the ListQueryExecutions WorkGroup/MaxResults. —
TestAioCursor::test_execute_qmark_parameters_skip_cache REWRITE 1 The #941 regression for aio code, asserted through patched private methods; call execute("SELECT ?", [...], paramstyle="qmark", cache_size=10) with a mocked client and assert no list_query_executions call. —
TestAioCursor::test_execute_kill_on_interrupt_timeout REMOVE 2 A wait_for timeout is a task cancellation, whose path and task.cancelled() the cover asserts; the TimeoutError is asyncio's. TestAioCursor::test_execute_kill_on_interrupt
TestAioCursor::test_execute_timeout_while_starting REMOVE 2, 3 The same start-phase cancellation through wait_for, with sleep(0.1) against a 0.05 s timeout. TestAioCursor::test_execute_cancelled_while_starting
TestAioCursor::test_execute_cancelled_again_while_start_is_retried REWRITE 3, 5 Proves "no retry" by asyncio.sleep(0.2) and calls == 1; wait for the start worker to finish (an event set when retry_api_call returns) before asserting calls == 1. —
TestAioCursor::test_cache_size_different_schema, test_cache_size_different_catalog, test_cache_search_prefers_latest_execution REMOVE 1, 2 They exercise _match_previous_query, shared code (pyathena/common.py:1122, called from pyathena/aio/common.py:357); identical sync copies exist. test_cursor.py::TestCursor::test_cache_size_different_schema, test_cache_size_different_catalog, test_cache_search_prefers_latest_execution
TestAioCursor::test_cache_search_stops_at_expired_execution, test_cache_search_reads_pages_up_to_cache_size REWRITE 1 The page loop is aio code, but the tests patch the private _list_query_executions and match log text; stub list_query_executions/batch_get_query_execution and assert the number and MaxResults of the ListQueryExecutions requests. —
TestAioCursor::test_no_result_set_raises MERGE 2, 4 The offline test covers the other 4 aio cursors; add AioCursor and AioDictCursor (own fetch overrides) to its parametrization. into aio/test_common.py::TestWithAsyncFetch::test_fetch_without_result_set
TestAioCursor::test_description_initial REMOVE 2, 6 The description property is the shared WithResultSet's. test_cursor.py::TestCursor::test_description_initial [original]
TestAioCursor::test_query_id REMOVE 2, 4 Checks the UUID shape of Athena's ID; query_id after execute is asserted by test_fetchone. TestAioCursor::test_fetchone, test_cursor.py::TestCursor::test_query_id [original]
TestAioCursor::test_query_execution_initial REMOVE 2, 6 24 shared WithResultSet properties that return None before execute. test_cursor.py::TestCursor::test_query_execution_initial [original]
TestAioCursor::test_cancel_initial MERGE 2, 4 cancel() without a query is the shared WithAsyncFetch.cancel; write one offline test over _ALL_CURSOR_CLASSES that raises ProgrammingError. into the new aio/test_common.py::TestWithAsyncFetch::test_cancel_without_query
TestAioCursor::test_executemany REMOVE 2, 4 Rowcount summing is asserted by test_executemany_rowcount[update], and the quoting of jim o'rourke is the shared formatter (6 queries). TestAioCursor::test_executemany_rowcount[update], test_cursor.py::TestCursor::test_executemany [original]
TestAioCursor::test_executemany_rowcount[insert], [delete] REWRITE 2, 4 Keep only [update]; the per-statement counts come from the shared _process_update_count (about 20 queries). test_cursor.py::TestCursor::test_iceberg_table [original], arrow/test_cursor.py::TestArrowCursor::test_iceberg_table [original]
TestAioCursor::test_executemany_failure[0] REWRITE 2, 4 Both cases reach the same except branch; keep [1], which also shows earlier executions are not rolled back (about 7 queries). TestAioCursor::test_executemany_failure[1], test_cursor.py::TestCursor::test_executemany_failure[0]
TestAioCursor::test_executemany_parameter_iteration_failure MERGE 2, 4 The offline test drives the same failure; have its fake execute set query_id and assert rowcount == -1 and the retained query_id. into aio/test_common.py::TestWithAsyncFetch::test_executemany_failure_closes_result_set_off_event_loop
TestAioCursor::test_executemany_fetch REMOVE 2, 4 After executemany the result set is None, and fetch without one raises; both are asserted offline. aio/test_common.py::TestWithAsyncFetch::test_executemany_closes_previous_result_set_off_event_loop, TestAioCursor::test_executemany_rowcount[update]
TestAioCursor::test_context_manager REMOVE 2, 4 Same as test_aio_connect (default cursor class under async with). TestAioCursor::test_aio_connect
TestAioCursor::test_open_close REMOVE 2, 6 Every fixture opens and closes a connection. TestAioCursor::test_aio_connect
TestAioCursor::test_arraysize_default REMOVE 6 Asserts a class constant. —
TestAioCursor::test_invalid_arraysize REMOVE 2 The constructor goes through the same AioCursor.arraysize setter. test_connection.py::TestConnection::test_cursor_arraysize_over_page_size[AioCursor], test_cursor_arraysize_not_positive[AioCursor] (after the section 1 merge)
TestAioCursor::test_glue_request_runs_off_the_event_loop REWRITE 1, 4 Asserts the internal call order ["request", "client"] against live Glue; keep only "the Glue request does not run on the event loop thread", offline with a stubbed Glue client. —
TestAioCursor::test_list_databases, test_get_table_metadata, test_list_table_metadata MERGE 2, 4 test_throttled_metadata_reads_glue makes the same three calls; assert "default"/ENV.schema in the databases, name == "one_row", and one_row in the listing on its expected. into TestAioCursor::test_throttled_metadata_reads_glue
TestAioDictCursor::test_fetchmany REMOVE 2, 4 Dict rows across pages are tested offline, aio paging by TestAioCursor::test_fetchmany. aio/test_result_set.py::TestAthenaAioResultSet::test_fetchmany_across_pages[AthenaAioDictResultSet], TestAioCursor::test_fetchmany
TestAioDictCursor::test_fetchall REMOVE 2, 4 The dict shape is test_fetchone's, the 10,000-row paging TestAioCursor::test_fetchall's (2 queries). TestAioDictCursor::test_fetchone, TestAioCursor::test_fetchall
TestAioDictCursor::test_dict_type_custom_result_set MERGE 4 Use the custom subclass as the OrderedDict cursor in test_dict_type; one query covers both. into TestAioDictCursor::test_dict_type

3. aio/test_result_set.py

test change criteria reason covered by
TestAthenaAioResultSet::test_get_query_results_after_close REMOVE 1, 7 Calls the private _async_get_query_results and matches message text; fetch after close is tested through the public API. TestAthenaAioResultSet::test_fetch_after_close

4. aio/test_util.py

#1238 deletes test_async_retry_api_call_stop_when; it is not listed.

test change criteria reason covered by
test_async_retry_api_call_cancelled[backoff], test_create_in_thread_cancelled_stops_retries[backoff] REWRITE 2, 3 Drop the backoff case, which uses asyncio.sleep(0.1) to "let the backoff begin" and passes if it cancels during the request; keep [request]. test_util.py::test_retry_api_call_stopped_by_event[backoff]
test_async_retry_api_call_not_stoppable REMOVE 2 The same scenario (cancel during the first attempt, stoppable=False, 3 attempts) at the cursor level. aio/test_cursor.py::TestAioCursor::test_cancel_cancelled_while_retrying
test_async_retry_api_call_retries REMOVE 2, 6 Retrying until success without a cancellation is retry_api_call's basic behavior. test_util.py::test_retry_metadata_errors (cases with 2 calls), aio/test_cursor.py::TestAioCursor::test_execute_cancelled_while_start_is_retried, test_async_retry_api_call_not_stoppable_in_stopped_scope

5. aio/arrow/test_cursor.py

test change criteria reason covered by
TestAioArrowCursor::test_binary_null_vs_empty REMOVE 2, 4 Parsing is the sync AthenaArrowResultSet's; the query is identical. arrow/test_cursor.py::TestArrowCursor::test_binary_null_vs_empty
TestAioArrowCursor::test_fetchmany REMOVE 2, 4 The wrapper forwards fetchmany(size) to a thread (tested offline); slicing is the sync result set's. aio/test_common.py::TestWithAsyncFetch::test_fetch_runs_result_set_fetch, arrow/test_cursor.py::TestArrowCursor::test_fetchmany [original]
TestAioArrowCursor::test_fetchall REMOVE 2, 4 One to_thread call; test_fetchone proves the aio execute builds a working result set (2 queries, one 10,000-row read). TestAioArrowCursor::test_fetchone, arrow/test_cursor.py::TestArrowCursor::test_fetchall
TestAioArrowCursor::test_execute_returns_self MERGE 4 Fold assert await cursor.execute(...) is cursor into test_fetchone. into TestAioArrowCursor::test_fetchone
TestAioArrowCursor::test_as_polars MERGE 4 Same one_row query as test_as_arrow; call both on one result. into TestAioArrowCursor::test_as_arrow
TestAioArrowCursor::test_as_arrow_unload MERGE 4 Same UNLOAD query as test_fetchone_unload; call as_arrow() there. into TestAioArrowCursor::test_fetchone_unload
TestAioArrowCursor::test_no_result_set_raises MERGE 2, 4 The fetch part duplicates the offline test; move the as_arrow() assertion there. into aio/test_common.py::TestWithAsyncFetch::test_fetch_without_result_set
TestAioArrowCursor::test_cancel_initial MERGE 2, 4 Shared WithAsyncFetch.cancel. into the new aio/test_common.py::TestWithAsyncFetch::test_cancel_without_query
TestAioArrowCursor::test_context_manager REMOVE 2, 4 __aexit__ is shared and tested offline; every fixture uses async with. aio/test_common.py::TestWithAsyncFetch::test_context_exit_closes_off_event_loop
TestAioArrowCursor::test_arraysize_default REMOVE 6 Asserts the result-set class constant. —
TestAioArrowCursor::test_invalid_arraysize REMOVE 2 The same setter as the constructor path. test_connection.py::TestConnection::test_cursor_arraysize_not_positive[AioArrowCursor] (after the section 1 merge)
TestAioArrowCursor::test_description REMOVE 2, 4 The description comes from the shared metadata handling of the sync result set. aio/test_cursor.py::TestAioCursor::test_description, test_cursor.py::TestCursor::test_description
TestAioArrowCursor::test_description_initial REMOVE 2, 6 Shared WithResultSet.description. test_cursor.py::TestCursor::test_description_initial [original]
TestAioArrowCursor::test_executemany_fetch REMOVE 2, 4 The shared WithAsyncFetch.executemany resets the result set, and fetch without one raises; both are tested offline. aio/test_common.py::TestWithAsyncFetch::test_executemany_closes_previous_result_set_off_event_loop, test_fetch_without_result_set

6. aio/polars/test_cursor.py

test change criteria reason covered by
TestAioPolarsCursor::test_fetchmany REMOVE 2, 4 The wrapper forwards fetchmany(size) to a thread; slicing is the sync result set's. aio/test_common.py::TestWithAsyncFetch::test_fetch_runs_result_set_fetch, polars/test_cursor.py::TestPolarsCursor::test_fetchmany
TestAioPolarsCursor::test_fetchall REMOVE 2, 4 One to_thread call (2 queries, one 10,000-row read). TestAioPolarsCursor::test_fetchone, polars/test_cursor.py::TestPolarsCursor::test_fetchall
TestAioPolarsCursor::test_execute_returns_self MERGE 4 Fold the is cursor assertion into test_fetchone. into TestAioPolarsCursor::test_fetchone
TestAioPolarsCursor::test_as_arrow MERGE 4 Same one_row query as test_as_polars; call both on one result. into TestAioPolarsCursor::test_as_polars
TestAioPolarsCursor::test_as_polars_unload MERGE 4 Same UNLOAD query as test_fetchone_unload; call as_polars() there. into TestAioPolarsCursor::test_fetchone_unload
TestAioPolarsCursor::test_no_result_set_raises MERGE 2, 4 Move the as_polars() assertion into the offline test. into aio/test_common.py::TestWithAsyncFetch::test_fetch_without_result_set
TestAioPolarsCursor::test_cancel_initial MERGE 2, 4 Shared WithAsyncFetch.cancel. into the new aio/test_common.py::TestWithAsyncFetch::test_cancel_without_query
TestAioPolarsCursor::test_context_manager REMOVE 2, 4 Shared __aexit__, tested offline. aio/test_common.py::TestWithAsyncFetch::test_context_exit_closes_off_event_loop
TestAioPolarsCursor::test_arraysize_default REMOVE 6 Asserts the result-set class constant. —
TestAioPolarsCursor::test_invalid_arraysize REMOVE 2 The same setter as the constructor path. test_connection.py::TestConnection::test_cursor_arraysize_not_positive[AioPolarsCursor] (after the section 1 merge)
TestAioPolarsCursor::test_description REMOVE 2, 4 Shared metadata handling of the sync result set. aio/test_cursor.py::TestAioCursor::test_description, test_cursor.py::TestCursor::test_description
TestAioPolarsCursor::test_description_initial REMOVE 2, 6 Shared WithResultSet.description. test_cursor.py::TestCursor::test_description_initial [original]
TestAioPolarsCursor::test_executemany_fetch REMOVE 2, 4 Shared WithAsyncFetch.executemany, tested offline. aio/test_common.py::TestWithAsyncFetch::test_executemany_closes_previous_result_set_off_event_loop, test_fetch_without_result_set

7. aio/s3fs/test_cursor.py

test change criteria reason covered by
TestAioS3FSCursor::test_null_vs_empty_string (2 cases) REMOVE 2, 4 Reader behavior is the sync result set's; the cursor option reaching the result set is asserted offline (2 queries). s3fs/test_cursor.py::TestS3FSCursor::test_null_vs_empty_string, TestAioS3FSCursor::test_read_options
TestAioS3FSCursor::test_fetchmany REMOVE 2, 4 The wrapper forwards fetchmany(size) to a thread; slicing is the sync result set's. aio/test_common.py::TestWithAsyncFetch::test_fetch_runs_result_set_fetch, s3fs/test_cursor.py::TestS3FSCursor::test_fetchmany
TestAioS3FSCursor::test_fetchall REMOVE 2, 4 One to_thread call (2 queries, one 10,000-row read). TestAioS3FSCursor::test_fetchone, s3fs/test_cursor.py::TestS3FSCursor::test_fetchall
TestAioS3FSCursor::test_execute_returns_self MERGE 4 Fold the is cursor assertion into test_fetchone. into TestAioS3FSCursor::test_fetchone
TestAioS3FSCursor::test_cancel_initial MERGE 2, 4 Shared WithAsyncFetch.cancel. into the new aio/test_common.py::TestWithAsyncFetch::test_cancel_without_query
TestAioS3FSCursor::test_context_manager REMOVE 2, 4 Shared __aexit__, tested offline. aio/test_common.py::TestWithAsyncFetch::test_context_exit_closes_off_event_loop
TestAioS3FSCursor::test_async_iterator REMOVE 2, 4 Shared __anext__. aio/test_common.py::TestWithAsyncFetch::test_async_iteration[AioS3FSCursor]
TestAioS3FSCursor::test_arraysize_default REMOVE 6 Asserts the result-set class constant. —
TestAioS3FSCursor::test_invalid_arraysize REMOVE 2 The same setter as the constructor path. test_connection.py::TestConnection::test_cursor_arraysize_not_positive[AioS3FSCursor] (after the section 1 merge)
TestAioS3FSCursor::test_description REMOVE 2, 4 Shared metadata handling of the sync result set. aio/test_cursor.py::TestAioCursor::test_description, test_cursor.py::TestCursor::test_description
TestAioS3FSCursor::test_description_initial REMOVE 2, 6 Shared WithResultSet.description. test_cursor.py::TestCursor::test_description_initial [original]
TestAioS3FSCursor::test_executemany_fetch REMOVE 2, 4 Shared WithAsyncFetch.executemany, tested offline. aio/test_common.py::TestWithAsyncFetch::test_executemany_closes_previous_result_set_off_event_loop, test_fetch_without_result_set

8. aio/spark/test_cursor.py

test change criteria reason covered by
TestAioSparkCursor::test_spark_dataframe REWRITE 2, 4 Run one short calculation with description= and assert the cursor properties and get_std_out(); the CSV read, groupBy and saveAsTable exercise Spark, and they write the same ENV.schema.spark_group_by table and path as the sync test. spark/test_spark_cursor.py::TestSparkCursor::test_spark_dataframe [original]
TestAioSparkCursor::test_spark_sql REMOVE 2, 4 Engine behavior already in the sync test; its other job, dropping the table that test_spark_dataframe writes, ends with that rewrite (1 session, 2 calculations). spark/test_spark_cursor.py::TestSparkCursor::test_spark_sql [original]
TestAioSparkCursor::test_session_ownership REMOVE 2, 4 The ownership logic is the shared SparkBaseCursor.__init__, and aio close() per ownership case is tested offline (1 session, 2 calculations). spark/test_common.py::TestSparkBaseCursor::test_close_terminates_session_by_ownership[AioSparkCursor-*], spark/test_spark_cursor.py::TestSparkCursor::test_session_ownership
TestAioSparkCursor::test_execute_timeout_while_starting REMOVE 2, 3 A wait_for timeout is the same task cancellation, with sleep(0.1) against a 0.05 s timeout. TestAioSparkCursor::test_execute_cancelled_while_starting
TestAioSparkCursor::test_execute_accepts_none_parameters (3 cases) REWRITE 1 Asserts the exact kwargs of the private _calculate; assert one StartCalculationExecution on a mocked client and that calculation_id and the state are set. —
TestAioSparkCursor::test_executemany REWRITE 4 Starts a Spark session to see NotSupportedError from a method that always raises; build the cursor with __new__ as the neighbouring offline tests do and assert the same error. —
TestAioSparkCursor::test_context_manager REMOVE 2, 4 __aexit__ is await self.close(), and close is tested offline (1 session, 1 calculation). spark/test_common.py::TestSparkBaseCursor::test_close_terminates_session_by_ownership[AioSparkCursor-*]

9. spark/test_async_cursor.py

test change criteria reason covered by
TestAsyncSparkCursor::test_session_ownership REMOVE 2, 4 AsyncSparkCursor.close() delegates the termination to SparkBaseCursor.close(); the matrix is tested offline (1 session, 2 calculations). spark/test_common.py::TestSparkBaseCursor::test_close_terminates_session_by_ownership[AsyncSparkCursor-*], spark/test_spark_cursor.py::TestSparkCursor::test_session_ownership
TestAsyncSparkCursor::test_close_retries_termination_after_failure REMOVE 1, 2 Patches the private _terminate_session; the same contract is tested with a mocked client. spark/test_common.py::TestSparkBaseCursor::test_close_retries_failed_termination[AsyncSparkCursor]
TestAsyncSparkCursor::test_execute_accepts_none_parameters (3 cases) REWRITE 1 Asserts the exact args of the private _calculate; assert the returned (query_id, future) and the submitted calculation instead. —

10. spark/test_spark_cursor.py

test change criteria reason covered by
TestSparkCursor::test_execute_accepts_none_parameters (3 cases) REWRITE 1 Asserts the exact args of the private _calculate; assert that the calculation ran (calculation_id and the state set). —

11. spark/test_common.py

test change criteria reason covered by
TestSparkBaseCursor::test_wait_for_idle_session_raises_on_failure_state (3 cases) REMOVE 1, 2 Calls the private _wait_for_idle_session with a patched _get_session_status; the same 3 states are asserted through construction. TestSparkBaseCursor::test_init_terminates_new_session_in_failure_state (rewritten)
TestSparkBaseCursor::test_wait_for_idle_session_raises_without_reason REMOVE 1, 7 Checks only the message format without a reason, with an anchored regex. —
TestSparkBaseCursor::test_wait_for_idle_session_waits_until_idle REWRITE 1 Asserts get_status.call_count == 3 and sleep.call_count == 2; drive _init_cursor with the statuses CREATING, BUSY, IDLE and assert that the cursor is created and no termination is sent. —
TestSparkBaseCursor::test_exists_session_raises_on_failure_state REWRITE 1 Calls the private _exists_session and asserts the exact get_session args; construct with session_id="supplied-session" and a TERMINATED status and assert OperationalError and no termination. —
TestSparkBaseCursor::test_calculate_reuses_generated_token_on_retry, test_calculate_generates_token_per_call, test_calculate_keeps_caller_token, test_calculate_interrupted_while_starting_failure, test_calculate_second_interrupt_while_starting, test_calculate_interrupted_before_request_is_sent REWRITE 2 Drop the AsyncSparkCursor case; it inherits SparkBaseCursor._calculate and _start_execution unchanged, and the tests call _calculate on the caller thread (about 20 duplicate cases). the SparkCursor cases of the same tests
TestSparkBaseCursor::test_calculate_interrupted_while_starting (2 classes × 2 states) MERGE 1, 2 Calls the private _calculate; parametrize the public execute() version over final_state. into TestSparkBaseCursor::test_execute_interrupted_while_starting
TestSparkBaseCursor::test_calculate_interrupt_without_kill_on_interrupt (2 cases) REMOVE 1, 2 Asserts that threading.Thread is not called; the kill_on_interrupt=False branch is the shared BaseCursor._start_execution, tested with the same patch for the query cursor. test_cursor.py::TestCursor::test_execute_interrupted_while_starting_without_kill_on_interrupt
TestSparkBaseCursor::test_init_starts_session (3 cases) REMOVE 2, 6 Asserts that the StartSession ID is stored, which every default-init ownership case depends on. TestSparkBaseCursor::test_close_terminates_session_by_ownership[*-None-None]
TestSparkBaseCursor::test_init_terminates_new_session_that_does_not_become_idle, test_init_does_not_terminate_supplied_session, test_init_does_not_start_session_when_s3_client_fails REWRITE 2 Drop the AioSparkCursor case; it has no __init__ and the same bases as SparkCursor. Keep SparkCursor and AsyncSparkCursor. the SparkCursor cases of the same tests
TestSparkBaseCursor::test_init_terminates_new_session_in_failure_state (3 × 3) REWRITE 1, 2 Drop the AioSparkCursor case, loosen the anchored message regex and drop get_session_status.assert_called_once_with; assert the exception type and the termination for each of the 3 states. —
TestSparkBaseCursor::test_init_keeps_original_error_when_cleanup_fails (3 cases) REWRITE 1, 2 Drop the AioSparkCursor case and the caplog text assertion; assert exc_info.value is error. —
TestSparkBaseCursor::test_async_close_shuts_down_executor_without_terminating_session REMOVE 2 "Supplied session not terminated" is a case of the ownership matrix, and the executor shutdown is asserted by the close tests. TestSparkBaseCursor::test_close_terminates_session_by_ownership[AsyncSparkCursor-supplied-session-None], spark/test_async_cursor.py::TestAsyncSparkCursor::test_close_no_wait_shuts_down_executor

Notes

  • Borderline KEEPs:
    • aio/test_common.py (all 10): parametrized over all aio SQL cursors although the behavior is inherited from WithAsyncFetch; the cases guard each class's MRO and its own execute().
    • aio/test_cursor.py::TestAioCursor::test_execute_with_options, test_arraysize: the only checks that the aio execute() passes options and arraysize to its result set (1 query each).
    • aio/test_cursor.py::TestAioCursor::test_executemany_cancellation (about 5 queries): the only check of query_id and cancel() after a cancelled executemany.
    • aio/test_cursor.py::TestAioCursor::test_cancel_cancelled_while_retrying: calls the private _cancel; the only check of the unstoppable StopQueryExecution retries.
    • aio/test_util.py::test_create_in_thread_cancelled_after_construction: patches pyathena.aio.util.asyncio.wait to reach the race; the close-in-new-thread path has no other test.
    • aio/test_result_set.py::TestAthenaAioResultSet::test_fetch_with_close_during_retries[backoff] uses a 0.1 s threading.Timer; Break result sets after failed or cancelled fetches and drop the cross-thread guards #1238 deletes it.
    • aio/{arrow,polars,s3fs}/test_cursor.py::*::test_read_options, aio/test_connection.py::test_cursor_execute_supported_keywords: patch private _execute/_poll and assert constructor kwargs; the override merge is each cursor's own code and has no offline observable effect.
    • aio/spark/test_cursor.py::TestAioSparkCursor::test_cancel (1 session, 2 long calculations): the only live check of aio cancel().
    • spark/test_spark_cursor.py::TestSparkCursor::test_cleanup_in_stopped_scope, test_spark_on_poll_invoked_each_iteration: patch private methods; the only checks of the unstoppable Spark cleanup (Retries continue after a result set is closed or an asyncio API call is cancelled #1202) and of the Spark on_poll loop.
    • spark/test_common.py::TestSparkBaseCursor::test_calculate_interrupted_before_request_is_sent: patches pyathena.common.threading.Thread; the only test of the "abandon before the helper starts" branch of BaseCursor._start_execution.
    • spark/test_async_cursor.py::TestAsyncSparkCursor::test_close_wait_shuts_down_executor, test_close_no_wait_shuts_down_executor: the shutdown.assert_called_once_with(wait=...) and _terminate_session.assert_called_once_with() assertions can be dropped; the futures and submit after close carry the contract.
    • The [original] spark/test_spark_cursor.py::TestSparkCursor::test_spark_dataframe/test_spark_sql and their aio copies write and drop the same ENV.schema.spark_group_by table and S3 path; the section 8 rewrite removes the overlap.
  • Unused after the change:
    • aio/test_cursor.py: imports re, BINARY, Binary, ProgrammingError, AthenaResultSet, _aio_connect; call and patch once the cache-search tests stub the client.
    • aio/test_connection.py: RecordingAioCursor; the empty TestAioConnection; imports pyathena, Parameter, signature, DefaultTypeConverter, ProgrammingError; PAGED_CURSORS/UNCAPPED_CURSORS (inline them into CURSOR_CLASSES).
    • aio/test_result_set.py: ProgrammingError import.
    • aio/test_util.py: ClientError import; the stop_in parameter and "backoff" branch of _cancel_once_started.
    • aio/arrow/test_cursor.py, aio/polars/test_cursor.py, aio/s3fs/test_cursor.py: imports of the sync result-set class, ProgrammingError, ENV, _aio_connect.
    • aio/spark/test_cursor.py: imports _aio_connect, AthenaSessionStatus, wait_for_spark_session_state; ENV if the rewritten test_spark_dataframe no longer reads the CSV.
    • spark/test_async_cursor.py: imports AthenaSessionStatus, wait_for_spark_session_state.
    • spark/test_common.py: helpers _cursor() and _session_status(); SYNC_SPARK_CURSOR_CLASSES (replace with SparkCursor); logging once the caplog assertion is dropped.
  • Each row's cover is re-checked before the row is applied. A row whose cover does not hold becomes MERGE or KEEP.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancement

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions