From c22be6c16700084a11bed1c382953ee2a8a14612 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 21:23:29 +0900 Subject: [PATCH 1/5] Remove and rewrite unnecessary top-level tests Apply the rows of #1243 to the cursor, connection, converter, formatter, options, parser, result set, utility, and Glue tests. Remove tests that restate the code or repeat another test's contract, rewrite tests that asserted private state to assert results through public calls, and merge the Glue client comparison into the cursor's Glue fallback test. Rows whose cover did not fail when the behavior was broken become merges or stay: the array(varchar) of numbers only, the error without a response, the retry jitter, the retry policy copy, and the typed-to-untyped conversion fallback. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/test_connection.py | 152 +++--- tests/pyathena/test_converter.py | 190 +------- tests/pyathena/test_cursor.py | 761 ++++-------------------------- tests/pyathena/test_formatter.py | 6 +- tests/pyathena/test_glue.py | 67 +-- tests/pyathena/test_options.py | 65 --- tests/pyathena/test_parser.py | 53 --- tests/pyathena/test_result_set.py | 25 +- tests/pyathena/test_util.py | 53 +-- tests/pyathena/util.py | 91 +++- 10 files changed, 292 insertions(+), 1171 deletions(-) diff --git a/tests/pyathena/test_connection.py b/tests/pyathena/test_connection.py index 6df0b1686..577fd8597 100644 --- a/tests/pyathena/test_connection.py +++ b/tests/pyathena/test_connection.py @@ -8,7 +8,7 @@ import os import threading from concurrent.futures import ThreadPoolExecutor -from inspect import Parameter, signature +from inspect import signature from typing import Any from unittest.mock import MagicMock, patch @@ -19,11 +19,10 @@ from pyathena.arrow.async_cursor import AsyncArrowCursor from pyathena.arrow.cursor import ArrowCursor from pyathena.async_cursor import AsyncCursor, AsyncDictCursor -from pyathena.common import BaseCursor from pyathena.connection import Connection from pyathena.converter import DefaultTypeConverter from pyathena.cursor import Cursor, DictCursor -from pyathena.error import ProgrammingError +from pyathena.error import OperationalError, ProgrammingError from pyathena.filesystem.s3 import S3FileSystem from pyathena.pandas.async_cursor import AsyncPandasCursor from pyathena.pandas.cursor import PandasCursor @@ -36,6 +35,7 @@ from pyathena.spark.common import SparkBaseCursor from pyathena.spark.cursor import SparkCursor from pyathena.util import RetryConfig +from tests.pyathena.util import EVENT_TIMEOUT # Cursors whose arraysize is the GetQueryResults page size, capped at 1000. PAGED_CURSORS = [Cursor, DictCursor, AsyncCursor, AsyncDictCursor] @@ -51,11 +51,12 @@ AsyncS3FSCursor, ] CURSOR_CLASSES = PAGED_CURSORS + UNCAPPED_CURSORS + [SparkCursor, AsyncSparkCursor] +# Cursors that run on_start_query_execution. +START_CALLBACK_CURSORS = [Cursor, DictCursor, ArrowCursor, PandasCursor, PolarsCursor, S3FSCursor] # Cursors that pass execute() keyword arguments to a DataFrame reader. READER_CURSOR_CLASSES = [PandasCursor, PolarsCursor, AsyncPandasCursor, AsyncPolarsCursor] -@pytest.mark.parametrize("factory", [Connection, pyathena.connect]) @pytest.mark.parametrize( ("setting", "kwargs", "expected"), [ @@ -69,10 +70,10 @@ ("work_group", {"work_group": "explicit-workgroup"}, "explicit-workgroup"), ], ) -def test_environment_fallback(factory, setting, kwargs, expected, monkeypatch): +def test_environment_fallback(setting, kwargs, expected, monkeypatch): monkeypatch.setenv("AWS_ATHENA_S3_STAGING_DIR", "s3://environment/path/") monkeypatch.setenv("AWS_ATHENA_WORK_GROUP", "environment-workgroup") - conn = factory( + conn = Connection( region_name="us-east-1", aws_access_key_id="access_key", aws_secret_access_key="secret_key", @@ -84,60 +85,17 @@ def test_environment_fallback(factory, setting, kwargs, expected, monkeypatch): conn.close() -@pytest.mark.parametrize("factory", [Connection, pyathena.connect]) -def test_empty_settings_disable_both_environment_fallbacks(factory, monkeypatch): +def test_empty_settings_disable_both_environment_fallbacks(monkeypatch): monkeypatch.setenv("AWS_ATHENA_S3_STAGING_DIR", "s3://environment/path/") monkeypatch.setenv("AWS_ATHENA_WORK_GROUP", "environment-workgroup") with pytest.raises(ProgrammingError, match="Required argument"): - factory(region_name="us-east-1", s3_staging_dir="", work_group="") + Connection(region_name="us-east-1", s3_staging_dir="", work_group="") @pytest.mark.parametrize("factory", [pyathena.connect, Connection]) def test_connect_keyword_only(factory): - assert all( - parameter.kind in (Parameter.KEYWORD_ONLY, Parameter.VAR_KEYWORD) - for parameter in signature(factory).parameters.values() - ) - with patch("pyathena.connection.Session") as session: - with pytest.raises(TypeError, match="positional"): - factory("s3://bucket/path/", "us-east-1") - session.assert_not_called() - - -@pytest.mark.parametrize( - "cursor_class", [*CURSOR_CLASSES, BaseCursor, SparkBaseCursor, WithResultSet] -) -def test_cursor_constructor_keyword_only(cursor_class): - constructor = cursor_class.__init__ - assert [ - parameter.name - for parameter in signature(constructor).parameters.values() - if parameter.kind in (Parameter.POSITIONAL_ONLY, Parameter.POSITIONAL_OR_KEYWORD) - ] == ["self"] - with pytest.raises(TypeError, match="positional"): - constructor(object(), "positional-setting") - - -@pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) -def test_cursor_execute_keyword_only(cursor_class): - execute_signature = signature(cursor_class.execute) - assert [ - parameter.name - for parameter in execute_signature.parameters.values() - if parameter.kind in (Parameter.POSITIONAL_ONLY, Parameter.POSITIONAL_OR_KEYWORD) - ] == ["self", "operation", "parameters"] - parameters = {"value": 1} - cursor = object.__new__(cursor_class) - bound = execute_signature.bind(cursor, "SELECT %(value)s", parameters) - assert bound.arguments["parameters"] is parameters - if cursor_class in (SparkCursor, AsyncSparkCursor): - assert "work_group" not in execute_signature.parameters - else: - bound = execute_signature.bind(cursor, "SELECT %(value)s", parameters, work_group="group") - assert bound.arguments["work_group"] == "group" - execute_signature.bind(cursor, operation="SELECT 1") - with pytest.raises(TypeError, match="positional"): - cursor.execute("SELECT %(value)s", parameters, "group") + with pytest.raises(TypeError): + factory("s3://bucket/path/", "us-east-1") class RecordingCursor(Cursor): @@ -196,6 +154,21 @@ def _open_cursor(conn: Connection[Any], cursor_class: type[Any], **kwargs: Any) return conn.cursor(cursor_class, **kwargs) +@pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) +def test_cursor_constructor_keyword_only(cursor_class): + # Binds the arguments as a call does before the constructor runs. A call + # through Connection.cursor() passes every setting by keyword, so a positional + # setting would fail there as a duplicate keyword argument too. + with pytest.raises(TypeError): + signature(cursor_class).bind("positional-setting") + + +@pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) +def test_cursor_execute_keyword_only(offline_connection, cursor_class): + with _open_cursor(offline_connection, cursor_class) as cursor, pytest.raises(TypeError): + cursor.execute("SELECT %(value)s", {"value": 1}, "group") + + @pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) def test_cursor_unknown_constructor_keyword(offline_connection, cursor_class): with pytest.raises(TypeError, match="unexpected keyword argument 'work_gruop'"): @@ -205,15 +178,62 @@ def test_cursor_unknown_constructor_keyword(offline_connection, cursor_class): offline_connection.cursor(cursor_class) +# Athena's answers to a query or calculation that fails at its first status request. +FAILED_EXECUTION_RESPONSES = { + "StartQueryExecution": {"QueryExecutionId": "query-id"}, + "GetQueryExecution": { + "QueryExecution": { + "QueryExecutionId": "query-id", + "Query": "SELECT 1", + "StatementType": "DML", + "Status": {"State": "FAILED", "StateChangeReason": "failed"}, + "Statistics": {}, + "ResultConfiguration": {}, + "QueryExecutionContext": {}, + "ResultReuseConfiguration": {"ResultReuseByAgeConfiguration": {"Enabled": False}}, + "EngineVersion": {}, + } + }, + "StartCalculationExecution": {"CalculationExecutionId": "calculation-id", "State": "CREATING"}, + "GetCalculationExecutionStatus": { + "Status": {"State": "FAILED"}, + "Statistics": {"DpuExecutionInMillis": 0}, + }, + "GetCalculationExecution": { + "CalculationExecutionId": "calculation-id", + "SessionId": "session", + "Status": {"State": "FAILED", "StateChangeReason": "failed"}, + "Statistics": {"DpuExecutionInMillis": 0}, + "Result": {}, + }, +} + + @pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) def test_cursor_accepts_connection_callbacks(offline_connection, cursor_class): - on_start_query_execution, on_poll = MagicMock(), MagicMock() - offline_connection.on_start_query_execution = on_start_query_execution - offline_connection.on_poll = on_poll + started, polled = [], [] + offline_connection.on_start_query_execution = started.append + offline_connection.on_poll = polled.append - with _open_cursor(offline_connection, cursor_class) as cursor: - assert cursor._on_start_query_execution is on_start_query_execution - assert cursor._on_poll is on_poll + with ( + patch.object( + offline_connection.client, + "_make_api_call", + side_effect=lambda operation, _: FAILED_EXECUTION_RESPONSES[operation], + ), + _open_cursor(offline_connection, cursor_class) as cursor, + ): + try: + result = cursor.execute("SELECT 1") + except OperationalError: + pass + else: + # The Async cursors run the execution in a future. + result[1].result(timeout=EVENT_TIMEOUT) + + assert [execution.state for execution in polled] == ["FAILED"] + # The Async and Spark cursors do not run on_start_query_execution. + assert started == (["query-id"] if cursor_class in START_CALLBACK_CURSORS else []) @pytest.mark.parametrize( @@ -336,10 +356,10 @@ def test_s3_client_built_once_across_threads(self): session_client = conn.session.client def contended_client(*args, **kwargs): - created.append((args, conn._s3_client_lock.locked())) + created.append(args) return session_client(*args, **kwargs) - conn._session.client = contended_client + conn.session.client = contended_client barrier = threading.Barrier(8) def get_client(_): @@ -349,7 +369,7 @@ def get_client(_): with ThreadPoolExecutor(max_workers=8) as executor: clients = list(executor.map(get_client, range(8))) - assert created == [(("s3",), True)] + assert created == [("s3",)] assert all(client is clients[0] for client in clients) assert clients[0].meta.service_model.service_name == "s3" @@ -401,13 +421,15 @@ def test_s3_config_merged_over_config(self): def test_close_closes_built_clients(self): conn = _connection() - with patch.object(conn.client, "close") as athena_close: + with ( + patch.object(conn.session, "client") as build, + patch.object(conn.client, "close") as athena_close, + ): conn.close() athena_close.assert_called_once_with() # Clients that were not used are not built to be closed. - assert conn._glue._client is None - assert conn._s3_client is None + build.assert_not_called() with ( patch.object(conn.client, "close") as athena_close, diff --git a/tests/pyathena/test_converter.py b/tests/pyathena/test_converter.py index bac5a8e0c..ac334a80a 100644 --- a/tests/pyathena/test_converter.py +++ b/tests/pyathena/test_converter.py @@ -123,18 +123,15 @@ def test_to_struct_athena_nested_formats(input_value, expected): @pytest.mark.parametrize( - "input_value", + ("input_value", "expected"), [ - "{formula=x=y+1, status=active}", - '{json={"key": "value"}, name=test}', - '{message=He said "hello", name=John}', + ("{formula=x=y+1, status=active}", {"formula": "x=y+1", "status": "active"}), + ('{json={"key": "value"}, name=test}', {"json": {"key": "value"}, "name": "test"}), + ('{message=He said "hello", name=John}', {"message": 'He said "hello"', "name": "John"}), ], ) -def test_to_struct_athena_complex_cases(input_value): - result = _to_struct(input_value) - # Values too complex to parse fall back to the original string - # instead of being silently dropped - assert result == input_value or isinstance(result, dict) +def test_to_struct_athena_complex_cases(input_value, expected): + assert _to_struct(input_value) == expected @pytest.mark.parametrize( @@ -332,17 +329,6 @@ def test_to_array_invalid_formats(input_value): assert _to_array(input_value) is None -@pytest.mark.parametrize( - ("value", "type_hint", "expected"), - [ - ('[""]', "array(json)", [""]), - ('{"k": ""}', "map(varchar,json)", {"k": ""}), - ], -) -def test_nested_json_empty_string(value, type_hint, expected): - assert DefaultTypeConverter().convert(type_hint.split("(")[0], value, type_hint) == expected - - class TestDefaultTypeConverter: @pytest.mark.parametrize( ("value", "expected"), @@ -351,77 +337,11 @@ class TestDefaultTypeConverter: def test_boolean_conversion(self, value, expected): assert DefaultTypeConverter().convert("boolean", value) is expected - @pytest.mark.parametrize( - ("input_value", "expected"), - [ - ('{"name": "Alice", "age": 25}', {"name": "Alice", "age": 25}), - (None, None), - ("", None), - ("invalid json", None), - ("{a=1, b=2}", {"a": "1", "b": "2"}), - ], - ) - def test_struct_conversion(self, input_value, expected): - converter = DefaultTypeConverter() - assert converter.convert("row", input_value) == expected - - @pytest.mark.parametrize( - ("input_value", "expected"), - [ - ("[1, 2, 3]", [1, 2, 3]), - ('["a", "b", "c"]', ["a", "b", "c"]), - (None, None), - ("", None), - ("invalid json", None), - ("[apple, banana]", ["apple", "banana"]), - ("[]", []), - ], - ) - def test_array_conversion(self, input_value, expected): - converter = DefaultTypeConverter() - assert converter.convert("array", input_value) == expected - - def test_array_varchar_keeps_strings(self): - converter = DefaultTypeConverter() - result = converter.convert("array", "[1234, 5678]", type_hint="array(varchar)") - assert result == ["1234", "5678"] - - def test_array_integer_converts_to_int(self): - converter = DefaultTypeConverter() - result = converter.convert("array", "[1, 2, 3]", type_hint="array(integer)") - assert result == [1, 2, 3] - - def test_array_boolean_converts(self): - converter = DefaultTypeConverter() - result = converter.convert("array", "[true, false]", type_hint="array(boolean)") - assert result == [True, False] - - def test_array_with_null(self): - converter = DefaultTypeConverter() - result = converter.convert("array", "[1, null, 3]", type_hint="array(integer)") - assert result == [1, None, 3] - - def test_map_varchar_integer(self): - converter = DefaultTypeConverter() - result = converter.convert( - "map", '{"key1": 1, "key2": 2}', type_hint="map(varchar, integer)" - ) - assert result == {"key1": 1, "key2": 2} - def test_map_native_format_with_hints(self): converter = DefaultTypeConverter() result = converter.convert("map", "{a=1, b=2}", type_hint="map(varchar, integer)") assert result == {"a": 1, "b": 2} - def test_row_type_hint(self): - converter = DefaultTypeConverter() - result = converter.convert( - "row", - '{"name": "Alice", "age": 25}', - type_hint="row(name varchar, age integer)", - ) - assert result == {"name": "Alice", "age": 25} - def test_row_native_format_with_hints(self): converter = DefaultTypeConverter() result = converter.convert( @@ -443,30 +363,18 @@ def test_nested_array_of_row(self): {"name": "Bob", "age": 30}, ] - def test_array_varchar_prevents_number_inference(self): + @pytest.mark.parametrize( + ("input_value", "expected"), + [("[1234, 5678]", ["1234", "5678"]), ("[1234, 5678, hello]", ["1234", "5678", "hello"])], + ) + def test_array_varchar_prevents_number_inference(self, input_value, expected): converter = DefaultTypeConverter() - result = converter.convert( - "array", - "[1234, 5678, hello]", - type_hint="array(varchar)", - ) - assert result == ["1234", "5678", "hello"] + assert converter.convert("array", input_value, type_hint="array(varchar)") == expected def test_none_value_with_type_hint(self): converter = DefaultTypeConverter() assert converter.convert("array", None, type_hint="array(varchar)") is None - def test_simple_type_hint(self): - converter = DefaultTypeConverter() - assert converter.convert("varchar", "hello", type_hint="varchar") == "hello" - - def test_type_hint_caching(self): - converter = DefaultTypeConverter() - converter.convert("array", "[1, 2]", type_hint="array(integer)") - assert "array(integer)" in converter._parsed_hints - converter.convert("array", "[3, 4]", type_hint="array(integer)") - assert len(converter._parsed_hints) == 1 - def test_empty_array_with_type_hint(self): converter = DefaultTypeConverter() assert converter.convert("array", "[]", type_hint="array(varchar)") == [] @@ -485,71 +393,12 @@ def test_row_with_nested_struct(self): ) assert result == {"header": {"seq": 123, "stamp": "2024"}, "x": 4.5} - def test_fallback_on_malformed_value(self): - """When typed conversion fails (returns None), fall back to untyped conversion.""" - converter = DefaultTypeConverter() - # "not-an-array" doesn't look like an array — typed converter returns None. - # Untyped _to_array also returns None for this input, which is correct. - result = converter.convert("array", "not-an-array", type_hint="array(integer)") - assert result is None - - def test_fallback_preserves_struct_value(self): - """Malformed struct with type_hint still falls back to untyped parsing.""" - converter = DefaultTypeConverter() - # Struct with no closing brace — typed converter returns None. - # Untyped _to_struct also returns None here. - result = converter.convert("row", "{unclosed", type_hint="row(a integer)") - assert result is None - def test_fallback_returns_untyped_result(self): - """When typed conversion returns None, untyped conversion is used.""" + """A value that the typed conversion cannot parse is converted without the hint.""" converter = DefaultTypeConverter() - # The typed converter returns None for a struct that doesn't start with "{". - # The untyped _to_struct also returns None for non-struct input. - # Use an array example where typed converter returns None (not a bracket-wrapped - # value), but untyped _to_array can still parse it via JSON. - result = converter.convert( - "row", - '{"a": 1}', - type_hint="row(a varchar)", - ) - # Typed conversion succeeds here — "a" is varchar so "1" stays a string - assert result == {"a": "1"} - - def test_hive_syntax_through_converter(self): - """Hive-style syntax works end-to-end through DefaultTypeConverter.""" - converter = DefaultTypeConverter() - result = converter.convert("array", "[1, 2, 3]", type_hint="array") - assert result == [1, 2, 3] - - def test_hive_syntax_struct_through_converter(self): - """Hive struct syntax works end-to-end.""" - converter = DefaultTypeConverter() - result = converter.convert( - "row", - "{name=Alice, age=25}", - type_hint="struct", - ) - assert result == {"name": "Alice", "age": 25} - - def test_hive_syntax_caching(self): - """Hive syntax is normalized before cache lookup.""" - converter = DefaultTypeConverter() - converter.convert("array", "[1]", type_hint="array") - converter.convert("array", "[2]", type_hint="array(integer)") - # Both should normalize to "array(integer)" in the cache - assert "array(integer)" in converter._parsed_hints - assert len(converter._parsed_hints) == 1 - - def test_normalize_hive_syntax_noop(self): - """Trino-style input passes through unchanged.""" - assert DefaultTypeConverter._normalize_hive_syntax("array(integer)") == "array(integer)" - - def test_normalize_hive_syntax_replaces(self): - assert ( - DefaultTypeConverter._normalize_hive_syntax("array>") - == "array(struct(a int))" - ) + # The untyped conversion keeps the original string of nested native arrays. + value = "[{a=[1, 2]}, {b=[3]}]" + assert converter.convert("array", value, type_hint="array(row(a array(integer)))") == value def test_normalize_hive_syntax_struct(self): converter = DefaultTypeConverter() @@ -571,11 +420,8 @@ def test_normalize_hive_syntax_nested(self): def test_normalize_hive_syntax_map(self): converter = DefaultTypeConverter() - result = converter.convert( - "map", - '{"x": 1, "y": 2}', - type_hint="map", - ) + # The native values are strings without the hint. + result = converter.convert("map", "{x=1, y=2}", type_hint="map") assert result == {"x": 1, "y": 2} def test_normalize_hive_syntax_mixed(self): diff --git a/tests/pyathena/test_cursor.py b/tests/pyathena/test_cursor.py index b5604e854..af0e77c7b 100644 --- a/tests/pyathena/test_cursor.py +++ b/tests/pyathena/test_cursor.py @@ -1,6 +1,4 @@ import contextlib -import json -import logging import random import re import string @@ -13,7 +11,7 @@ from datetime import UTC, date, datetime, timedelta from decimal import Decimal from random import randint -from unittest.mock import MagicMock, call, patch +from unittest.mock import MagicMock, patch import pytest from botocore.exceptions import ClientError @@ -31,7 +29,6 @@ ExecuteOptions, ) from pyathena.async_cursor import AsyncCursor -from pyathena.converter import _to_array, _to_map, _to_struct from pyathena.cursor import Cursor, DictCursor from pyathena.error import ( DatabaseError, @@ -52,15 +49,15 @@ CONVERTED_VALUES_QUERY, CONVERTED_VALUES_ROW, EVENT_TIMEOUT, + QueryExecutionHistory, ThrottledOperation, interrupt_start_waits, succeeded_query_execution, + succeeded_query_execution_data, throttle_metadata_api, unreachable_glue, ) -_logger = logging.getLogger(__name__) - def _offline_cursor(kill_on_interrupt, cursor_class=Cursor): """A cursor whose query requests go to a mocked Athena client. @@ -94,6 +91,24 @@ def _offline_cursor(kill_on_interrupt, cursor_class=Cursor): return cursor, cancel +def _cache_search_cursor(client): + """A cursor whose result cache search sends its requests to a client stand-in. + + Args: + client: The Athena client stand-in, such as ``QueryExecutionHistory``. + + Returns: + The cursor, for the schema ``this_schema``. + """ + cursor = Cursor.__new__(Cursor) # bypass __init__ to avoid AWS calls + cursor._connection = MagicMock(_client=client) + cursor._retry_config = RetryConfig(attempt=1) + cursor._work_group = None + cursor._schema_name = "this_schema" + cursor._catalog_name = None + return cursor + + def _block_start(cursor, response=None): """Make the cursor's StartQueryExecution request block until released. @@ -394,46 +409,42 @@ def execution(catalog): == "query_id_awsdatacatalog" ) - def test_cache_search_stops_at_expired_execution(self, caplog): + def test_cache_search_stops_at_expired_execution(self): query = "SELECT * FROM one_row" now = datetime.now(UTC) - cursor = Cursor.__new__(Cursor) - cursor._schema_name = "this_schema" - cursor._catalog_name = None - page = [ - succeeded_query_execution("expired", query, now - timedelta(hours=2)), - succeeded_query_execution("other", "SELECT 1", now), + executions = [ + succeeded_query_execution_data("other", "SELECT 1", now), + succeeded_query_execution_data("expired", query, now - timedelta(hours=2)), + *( + succeeded_query_execution_data(f"older_{i}", "SELECT 1", now - timedelta(hours=3)) + for i in range(48) + ), + # A search that went on to the second page would return this match. + succeeded_query_execution_data("next_page", query, now), ] + history = QueryExecutionHistory(executions) + cursor = _cache_search_cursor(history) - # A page read after the expired execution would return this match. - next_page = (None, [succeeded_query_execution("next_page", query, now)]) + # Without cache_size, cache_expiration_time alone bounds the search. + assert cursor._find_previous_query_id(query, None, cache_expiration_time=3600) is None + assert history.list_requests == 1 - with patch.object( - Cursor, "_list_query_executions", side_effect=[("next_token", page), next_page] - ) as list_mock: - # Without cache_size, cache_expiration_time alone bounds the search. - assert cursor._find_previous_query_id(query, None, cache_expiration_time=3600) is None - list_mock.assert_called_once_with(None, next_token=None, max_results=50) - # A failed search also returns None; the expiry must stop it without an error. - assert "Failed to check the cache" not in caplog.text - - def test_cache_search_reads_pages_up_to_cache_size(self): + @pytest.mark.parametrize( + ("position", "expected"), + [(59, "match"), (60, None)], + ids=["last_searched", "after_cache_size"], + ) + def test_cache_search_reads_pages_up_to_cache_size(self, position, expected): + """With cache_size=60, the search reads the 60 latest executions over two pages.""" query = "SELECT * FROM one_row" now = datetime.now(UTC) - cursor = Cursor.__new__(Cursor) - cursor._schema_name = "this_schema" - cursor._catalog_name = None - pages = [ - ("next_token", [succeeded_query_execution("other", "SELECT 1", now)]), - ("last_token", [succeeded_query_execution("match", query, now)]), + executions = [ + succeeded_query_execution_data(f"other_{i}", "SELECT 1", now) for i in range(70) ] + executions[position] = succeeded_query_execution_data("match", query, now) + cursor = _cache_search_cursor(QueryExecutionHistory(executions)) - with patch.object(Cursor, "_list_query_executions", side_effect=pages) as list_mock: - assert cursor._find_previous_query_id(query, "wg", cache_size=60) == "match" - assert list_mock.call_args_list == [ - call("wg", next_token=None, max_results=50), - call("wg", next_token="next_token", max_results=10), - ] + assert cursor._find_previous_query_id(query, "wg", cache_size=60) == expected def test_cache_search_prefers_latest_execution(self): query = "SELECT * FROM one_row" @@ -582,9 +593,9 @@ def test_null_param(self, cursor): cursor.execute("SELECT %(param)s FROM one_row", {"param": None}) assert cursor.fetchall() == [(None,)] - @pytest.mark.parametrize( - ("value", "expected"), - [ + def test_binary_parameter(self, cursor): + # Empty, special, all bytes, bytearray, memoryview, DB API Binary, and NULL. + values = [ (b"", b""), (b"\x00\xff'\\%", b"\x00\xff'\\%"), (bytes(range(256)), bytes(range(256))), @@ -592,25 +603,13 @@ def test_null_param(self, cursor): (memoryview(b"\x00\xff"), b"\x00\xff"), (Binary(bytearray(b"abc")), b"abc"), (None, None), - ], - ids=["empty", "special", "all_bytes", "bytearray", "memoryview", "dbapi_binary", "null"], - ) - def test_binary_parameter(self, cursor, value, expected): - cursor.execute("SELECT CAST(%(value)s AS VARBINARY)", {"value": value}) - assert cursor.fetchone() == (expected,) - assert cursor.description[0][1] == BINARY - - def test_timestamp_result_precision(self, cursor): + ] cursor.execute( - """ - SELECT TIMESTAMP '2020-01-01 00:00:00', - CAST(TIMESTAMP '2020-01-01 00:00:00.123456' AS TIMESTAMP(9)) - """ - ) - assert cursor.fetchone() == ( - datetime(2020, 1, 1), - datetime(2020, 1, 1, 0, 0, 0, 123456), + "SELECT " + ", ".join(f"CAST(%(v{i})s AS VARBINARY)" for i in range(len(values))), + {f"v{i}": value for i, (value, _) in enumerate(values)}, ) + assert cursor.fetchone() == tuple(expected for _, expected in values) + assert [d[1] for d in cursor.description] == [BINARY] * len(values) def test_datetime_parameter_microseconds(self, cursor): value = datetime(2017, 1, 1, 12, 0, 0, 789012) @@ -880,76 +879,44 @@ def test_complex(self, cursor): ] def test_complex_with_type_hints(self, cursor): - # 1. Basic complex columns from one_row_complex + # Without the hints, the numbers in the structs and the map are strings. cursor.execute( """ - SELECT col_array, col_map, col_struct + SELECT + col_array + ,col_map + ,col_struct + ,CAST( + ROW(ROW('2024-01-01', 123), 4.736, 0.583) + AS ROW(header ROW(stamp VARCHAR, seq INTEGER), x DOUBLE, y DOUBLE) + ) AS positions + ,CAST( + ARRAY[ + ROW(ROW(1, 2), ROW(CAST(0.5 AS DOUBLE))), + ROW(ROW(3, 4), ROW(CAST(1.5 AS DOUBLE))) + ] + AS ARRAY(ROW(pos ROW(x INTEGER, y INTEGER), vel ROW(x DOUBLE))) + ) AS data FROM one_row_complex """, result_set_type_hints={ "col_array": "array(integer)", "col_map": "map(integer, integer)", "col_struct": "row(a integer, b integer)", - }, - ) - row = cursor.fetchone() - assert row[0] == [1, 2] - assert row[1] == {"1": 2, "3": 4} - assert row[2] == {"a": 1, "b": 2} - # Without type hints col_struct values are strings; with hints they are ints - assert isinstance(row[2]["a"], int) - assert isinstance(row[2]["b"], int) - # Map values should also be ints - assert isinstance(row[1]["1"], int) - - # 2. Nested struct - cursor.execute( - """ - SELECT CAST( - ROW(ROW('2024-01-01', 123), 4.736, 0.583) - AS ROW(header ROW(stamp VARCHAR, seq INTEGER), x DOUBLE, y DOUBLE) - ) AS positions - """, - result_set_type_hints={ "positions": "row(header row(stamp varchar, seq integer), x double, y double)", - }, - ) - row = cursor.fetchone() - positions = row[0] - assert positions["header"]["stamp"] == "2024-01-01" - assert positions["header"]["seq"] == 123 - assert isinstance(positions["header"]["seq"], int) - assert positions["x"] == 4.736 - assert isinstance(positions["x"], float) - assert positions["y"] == 0.583 - assert isinstance(positions["y"], float) - - # 3. Array of nested structs - cursor.execute( - """ - SELECT CAST( - ARRAY[ - ROW(ROW(1, 2), ROW(CAST(0.5 AS DOUBLE))), - ROW(ROW(3, 4), ROW(CAST(1.5 AS DOUBLE))) - ] - AS ARRAY(ROW(pos ROW(x INTEGER, y INTEGER), vel ROW(x DOUBLE))) - ) AS data - """, - result_set_type_hints={ "data": "array(row(pos row(x integer, y integer), vel row(x double)))", }, ) - row = cursor.fetchone() - data = row[0] - assert len(data) == 2 - assert data[0]["pos"]["x"] == 1 - assert isinstance(data[0]["pos"]["x"], int) - assert data[0]["pos"]["y"] == 2 - assert isinstance(data[0]["pos"]["y"], int) - assert data[0]["vel"]["x"] == 0.5 - assert isinstance(data[0]["vel"]["x"], float) - assert data[1]["pos"]["x"] == 3 - assert data[1]["vel"]["x"] == 1.5 + assert cursor.fetchone() == ( + [1, 2], + {"1": 2, "3": 4}, + {"a": 1, "b": 2}, + {"header": {"stamp": "2024-01-01", "seq": 123}, "x": 4.736, "y": 0.583}, + [ + {"pos": {"x": 1, "y": 2}, "vel": {"x": 0.5}}, + {"pos": {"x": 3, "y": 4}, "vel": {"x": 1.5}}, + ], + ) def test_cancel(self, cursor): def cancel(c): @@ -1044,21 +1011,16 @@ def test_executemany(self, cursor, empty_table): cursor.execute(f"SELECT * FROM {empty_table}") assert sorted(cursor.fetchall()) == list(rows) + # The sum of the update counts does not depend on the statement; test_iceberg_table + # checks the update count of INSERT, UPDATE, and DELETE. @pytest.mark.parametrize( ("operation", "expected"), [ - pytest.param( - "INSERT INTO {table} SELECT id+10, group_id+10, value " - "FROM {table} WHERE group_id=%(group_id)d", - [(1, 10), (2, 20), (3, 30), (11, 10), (12, 20), (13, 30)], - id="insert", - ), pytest.param( "UPDATE {table} SET value=value+1 WHERE group_id=%(group_id)d", [(1, 11), (2, 21), (3, 31)], id="update", ), - pytest.param("DELETE FROM {table} WHERE group_id=%(group_id)d", [], id="delete"), ], ) def test_executemany_rowcount(self, cursor, executemany_table, operation, expected): @@ -1341,36 +1303,6 @@ def test_execute_internal_options_passthrough(self): cache_expiration_time=100, ) - def test_execute_qmark_parameters_skip_cache(self): - """A qmark query with parameters never searches the cache (no AWS, #941).""" - cursor = Cursor.__new__(Cursor) # bypass __init__ to avoid AWS calls - cursor._connection = MagicMock() - cursor._connection.client.start_query_execution.return_value = { - "QueryExecutionId": "test_query_id" - } - cursor._retry_config = RetryConfig() - cursor._kill_on_interrupt = True - - with ( - patch.object( - Cursor, - "_build_start_query_execution_request", - return_value={"ExecutionParameters": ["'1'"]}, - ) as request_mock, - patch.object(Cursor, "_find_previous_query_id", return_value="cached") as cache_mock, - ): - query_id = cursor._execute( - "SELECT ?", - ["'1'"], - options=ExecuteOptions( - paramstyle="qmark", cache_size=10, cache_expiration_time=100 - ), - ) - - assert query_id == "test_query_id" - assert request_mock.call_args.kwargs["execution_parameters"] == ["'1'"] - cache_mock.assert_not_called() - def test_connection_level_callback(self): """Test connection-level default callback.""" callback_results = [] @@ -1417,78 +1349,6 @@ def execute_callback(query_id): assert connection_callbacks[0] == execute_callbacks[0] # Same query_id assert connection_callbacks[0] == cursor.query_id - def test_callback_thread_safety(self, cursor): - """Test that callback can be used for query ID access from another thread.""" - query_started = threading.Event() - stored_query_id = [] - - def on_start(query_id): - stored_query_id.append(query_id) - query_started.set() - - def verify_query_id(): - # Wait for query to start - query_started.wait(timeout=10) - - # Verify query_id is accessible - if stored_query_id: - assert stored_query_id[0] == cursor.query_id - - # Use ThreadPoolExecutor for proper thread management - with ThreadPoolExecutor(max_workers=1) as executor: - # Start verification task - verify_future = executor.submit(verify_query_id) - - # Execute query with callback - cursor.execute("SELECT 5", on_start_query_execution=on_start) - - # Wait for verification to complete - verify_future.result(timeout=5) - - # Verify events occurred - assert query_started.is_set() - assert len(stored_query_id) == 1 - - # Verify query completed successfully - row = cursor.fetchone() - assert row == (5,) - - def test_on_poll_invoked_each_iteration(self): - """on_poll fires once per poll iteration with the current execution (no AWS).""" - states = [ - AthenaQueryExecution.STATE_QUEUED, - AthenaQueryExecution.STATE_RUNNING, - AthenaQueryExecution.STATE_SUCCEEDED, - ] - executions = [MagicMock(state=state) for state in states] - received = [] - - cursor = Cursor.__new__(Cursor) # bypass __init__ to avoid AWS calls - cursor._poll_interval = 0 - cursor._kill_on_interrupt = False - cursor._on_poll = received.append - - with patch.object(Cursor, "_get_query_execution", side_effect=executions): - result = cursor._poll("query_id") - - # Callback received every state in order, including the terminal one - assert [execution.state for execution in received] == states - assert result is executions[-1] - - def test_on_poll_none_is_noop(self): - """A None on_poll callback does not affect polling (no AWS).""" - execution = MagicMock(state=AthenaQueryExecution.STATE_SUCCEEDED) - - cursor = Cursor.__new__(Cursor) # bypass __init__ to avoid AWS calls - cursor._poll_interval = 0 - cursor._kill_on_interrupt = False - cursor._on_poll = None - - with patch.object(Cursor, "_get_query_execution", return_value=execution): - result = cursor._poll("query_id") - - assert result is execution - @pytest.mark.parametrize( "final_state", [AthenaQueryExecution.STATE_CANCELLED, AthenaQueryExecution.STATE_SUCCEEDED], @@ -1786,41 +1646,27 @@ def _metadata_view(metadata): ) def test_throttled_metadata_reads_glue(self, cursor, monkeypatch): + tables = ("one_row", "parquet_with_compression", "partition_table", "view_one_row") + def read(): return ( - [ - self._metadata_view(cursor.get_table_metadata(t)) - for t in ("one_row", "parquet_with_compression", "partition_table") - ], + [self._metadata_view(cursor.get_table_metadata(t)) for t in tables], sorted(self._metadata_view(m) for m in cursor.list_table_metadata()), + [m.name for m in cursor.list_table_metadata(expression="one_row")], ENV.schema in [d.name for d in cursor.list_databases("AwsDataCatalog")], ) expected = read() + # The listing filtered by the expression is not empty. + assert "one_row" in expected[2] calls = throttle_metadata_api(cursor.connection.client, monkeypatch) # Glue reports each table as Athena does, and nothing is retried first. assert read() == expected assert sorted(calls) == sorted( - ["get_table_metadata"] * 3 + ["list_table_metadata", "list_databases"] + ["get_table_metadata"] * len(tables) + ["list_table_metadata"] * 2 + ["list_databases"] ) - def test_wrapped_glue_throttling_reads_glue(self, cursor, monkeypatch): - # Athena reports Glue's own throttling inside a MetadataException, and - # the first attempt does not retry it either. - calls = throttle_metadata_api( - cursor.connection.client, - monkeypatch, - code="MetadataException", - message=( - "Rate exceeded (Service: AmazonDataCatalog; Status Code: 400; " - "Error Code: ThrottlingException; Request ID: example; Proxy: null)" - ), - ) - - assert cursor.get_table_metadata("one_row").name == "one_row" - assert calls == ["get_table_metadata"] - @staticmethod def _unreachable_glue(connection, monkeypatch): glue = unreachable_glue(connection) @@ -2026,428 +1872,3 @@ def __init__(self, **kwargs): cursor.execute("SELECT * FROM one_row") assert isinstance(cursor.result_set, CustomResultSet) assert type(cursor.fetchone()) is OrderedDict - - def test_null_vs_empty_string(self, dict_cursor): - """ - DictCursor should properly distinguish NULL from empty string. - See docs/null_handling.rst for details. - """ - query = """ - SELECT * FROM ( - VALUES - (1, '', 'empty_string'), - (2, CAST(NULL AS VARCHAR), 'null_value'), - (3, 'hello', 'normal_string'), - (4, 'N/A', 'na_string'), - (5, 'NULL', 'null_string_literal') - ) AS t(id, value, description) - ORDER BY id - """ - dict_cursor.execute(query) - rows = dict_cursor.fetchall() - - # Row 1: Empty string is preserved as empty string, not None - assert rows[0]["value"] == "" - assert rows[0]["value"] is not None - - # Row 2: NULL is properly returned as None - assert rows[1]["value"] is None - - # Row 3: Normal string - assert rows[2]["value"] == "hello" - - # Row 4: "N/A" string should be preserved as-is - assert rows[3]["value"] == "N/A" - - # Row 5: "NULL" string literal should be preserved as-is - assert rows[4]["value"] == "NULL" - - -class TestComplexDataTypes: - """Test complex data types (STRUCT, ARRAY, MAP) with actual Athena queries.""" - - @pytest.mark.parametrize( - ("query", "description"), - [ - ("SELECT ROW('John', 30) AS simple_struct", "simple_struct"), - ( - "SELECT CAST(ROW('Alice', 25) AS ROW(name VARCHAR, age INTEGER)) AS named_struct", - "named_struct", - ), - ("SELECT ROW('Hello, world', 'x=y+1') AS special_chars_struct", "special_chars_struct"), - ("SELECT ROW('He said \"hello\"', 'It''s working') AS quotes_struct", "quotes_struct"), - ("SELECT ROW('Alice', NULL, 'active') AS null_struct", "null_struct"), - ( - "SELECT ROW(ROW('John', 30), ROW('Engineer', 'Tech')) AS nested_struct", - "nested_struct", - ), - ("SELECT CAST(ROW('Alice', 25, 'Hello, world') AS JSON) AS json_struct", "json_struct"), - ], - ) - def test_struct_types(self, cursor, query, description): - """Test various STRUCT type scenarios to understand Athena's behavior.""" - _logger.info("=== STRUCT Type Test: %s ===", description) - cursor.execute(query) - result = cursor.fetchone() - struct_value = result[0] - _logger.info("%s: %r (type: %s)", description, struct_value, type(struct_value).__name__) - - # Validate struct value and converter behavior - assert struct_value is not None, f"STRUCT value should not be None for {description}" - - # Test struct conversion behavior - if isinstance(struct_value, str): - converted = _to_struct(struct_value) - _logger.info("%s: Converted %r -> %r", description, struct_value, converted) - # Conversion either fully parses to a dict or falls back to the - # original string for values too complex to parse losslessly - if converted is not None: - assert isinstance(converted, dict) or converted == struct_value, ( - f"Converted struct should be dict or the original string for {description}" - ) - elif isinstance(struct_value, dict): - # Already converted by the cursor converter - _logger.info("%s: Already converted to dict: %r", description, struct_value) - else: - # Log unexpected types for debugging but don't fail - _logger.warning( - "%s: Unexpected type %s: %r", - description, - type(struct_value).__name__, - struct_value, - ) - - @pytest.mark.parametrize( - ("query", "description"), - [ - ("SELECT ARRAY[1, 2, 3, 4, 5] AS simple_array", "simple_array"), - ("SELECT ARRAY['apple', 'banana', 'cherry'] AS string_array", "string_array"), - ( - "SELECT ARRAY['Hello, world', 'x=y+1', 'It''s working'] AS special_array", - "special_array", - ), - ("SELECT ARRAY[ROW('Alice', 25), ROW('Bob', 30)] AS struct_array", "struct_array"), - ("SELECT ARRAY[ARRAY[1, 2], ARRAY[3, 4]] AS nested_array", "nested_array"), - ( - "SELECT CAST(MAP(ARRAY['data'], ARRAY[ARRAY['Alice', 'Bob']]) AS JSON) " - "AS array_json", - "array_json", - ), - ], - ) - def test_array_types(self, cursor, query, description): - """Test various ARRAY type scenarios.""" - _logger.info("=== ARRAY Type Test: %s ===", description) - cursor.execute(query) - result = cursor.fetchone() - array_value = result[0] - _logger.info("%s: %r (type: %s)", description, array_value, type(array_value).__name__) - - # Validate array value - assert array_value is not None, f"ARRAY value should not be None for {description}" - _logger.info("%s: Array value type %s", description, type(array_value).__name__) - - @pytest.mark.parametrize( - ("query", "description"), - [ - ( - "SELECT MAP(ARRAY[1, 2, 3], ARRAY['one', 'two', 'three']) AS simple_map", - "simple_map", - ), - ( - "SELECT MAP(ARRAY['name', 'age', 'city'], ARRAY['John', '30', 'Tokyo']) " - "AS string_map", - "string_map", - ), - ( - "SELECT MAP(ARRAY['msg', 'formula'], ARRAY['Hello, world', 'x=y+1']) " - "AS special_map", - "special_map", - ), - ( - "SELECT CAST(MAP(ARRAY['person1', 'person2'], " - "ARRAY[ROW('Alice', 25), ROW('Bob', 30)]) AS JSON) AS struct_value_map", - "struct_value_map", - ), - ( - "SELECT CAST(MAP(ARRAY['name', 'age'], ARRAY['Alice', '25']) AS JSON) AS json_map", - "json_map", - ), - ], - ) - def test_map_types(self, cursor, query, description): - """Test various MAP type scenarios.""" - _logger.info("=== MAP Type Test: %s ===", description) - cursor.execute(query) - result = cursor.fetchone() - map_value = result[0] - _logger.info("%s: %r (type: %s)", description, map_value, type(map_value).__name__) - - # Validate map value and converter behavior - assert map_value is not None, f"MAP value should not be None for {description}" - - # Test map conversion behavior - if isinstance(map_value, str): - # For complex MAP structures, string is expected (JSON or native format) - if "ROW(" in map_value or "ARRAY[" in map_value: - # Complex structure, expect string format - _logger.info("%s: Complex MAP kept as string: %r", description, map_value) - else: - # Simple MAP, try conversion - converted = _to_map(map_value) - _logger.info("%s: Converted %r -> %r", description, map_value, converted) - # Conversion either fully parses to a dict or falls back to the - # original string for values too complex to parse losslessly - if converted is not None: - assert isinstance(converted, dict) or converted == map_value, ( - f"Converted map should be dict or the original string for {description}" - ) - elif isinstance(map_value, dict): - # Already converted by the cursor converter - _logger.info("%s: Already converted to dict: %r", description, map_value) - else: - # Log unexpected types for debugging but don't fail - _logger.warning( - "%s: Unexpected type %s: %r", - description, - type(map_value).__name__, - map_value, - ) - - @pytest.mark.parametrize( - ("query", "description"), - [ - ( - "SELECT CAST(ROW(ARRAY[1, 2, 3], MAP(ARRAY['a', 'b'], ARRAY[1, 2])) AS JSON) " - "AS struct_with_collections", - "struct_with_collections", - ), - ( - "SELECT CAST(ARRAY[MAP(ARRAY['name'], ARRAY['Alice']), " - "MAP(ARRAY['name'], ARRAY['Bob'])] AS JSON) AS array_of_maps", - "array_of_maps", - ), - ( - "SELECT CAST(MAP(ARRAY['numbers', 'letters'], " - "ARRAY[ARRAY['1', '2', '3'], ARRAY['a', 'b', 'c']]) AS JSON) AS map_with_arrays", - "map_with_arrays", - ), - ], - ) - def test_complex_combinations(self, cursor, query, description): - """Test complex combinations of data types.""" - _logger.info("=== Complex Combination Test: %s ===", description) - cursor.execute(query) - result = cursor.fetchone() - complex_value = result[0] - _logger.info("%s: %r (type: %s)", description, complex_value, type(complex_value).__name__) - - # For JSON cast results, expect string values that can be parsed as JSON - if isinstance(complex_value, str): - try: - # Test that the JSON string can be parsed - parsed = json.loads(complex_value) - _logger.info(" Parsed JSON: %r", parsed) - assert parsed is not None, f"Parsed JSON should not be None for {description}" - except json.JSONDecodeError as e: - raise AssertionError(f"JSON parsing failed for {description}: {e}") from e - else: - # If it's not a string, it should still be a valid value (not None) - assert complex_value is not None, f"Complex value should not be None for {description}" - _logger.info("%s: Complex value type %s", description, type(complex_value).__name__) - - @pytest.mark.parametrize( - ("query", "description"), - [ - ("SELECT ARRAY[1, 2, 3, 4, 5] AS simple_array", "simple_array"), - ("SELECT ARRAY['apple', 'banana', 'cherry'] AS string_array", "string_array"), - ("SELECT ARRAY[true, false, null] AS boolean_array", "boolean_array"), - ("SELECT ARRAY[1.5, 2.7, 3.14] AS float_array", "float_array"), - ("SELECT ARRAY[] AS empty_array", "empty_array"), - ("SELECT ARRAY[1, null, 3, null, 5] AS null_elements_array", "null_elements_array"), - ( - "SELECT ARRAY[CAST(1 AS VARCHAR), 'mixed', 'true', 'null', '2.5'] AS mixed_array", - "mixed_array", - ), - ( - "SELECT ARRAY[CAST(1.23 AS DECIMAL(10,2)), CAST(4.56 AS DECIMAL(10,2))] " - "AS decimal_array", - "decimal_array", - ), - ( - "SELECT ARRAY[DATE '2023-01-01', DATE '2023-12-31'] AS date_array", - "date_array", - ), - ( - "SELECT ARRAY[TIMESTAMP '2023-01-01 12:00:00'] AS timestamp_array", - "timestamp_array", - ), - ], - ) - def test_array_types_basic(self, cursor, query, description): - """Test basic ARRAY type scenarios.""" - _logger.info("=== ARRAY Type Test: %s ===", description) - cursor.execute(query) - result = cursor.fetchone() - array_value = result[0] - _logger.info("%s: %r (type: %s)", description, array_value, type(array_value).__name__) - - # Validate array value - assert array_value is not None or description == "empty_array", ( - f"ARRAY value should not be None for {description}" - ) - if description == "empty_array": - assert array_value == [], f"Empty array should be [] for {description}" - else: - assert isinstance(array_value, list), f"ARRAY value should be list for {description}" - _logger.info("%s: Array value type %s", description, type(array_value).__name__) - - @pytest.mark.parametrize( - ("query", "description"), - [ - ( - "SELECT ARRAY[ROW(1, 'Alice'), ROW(2, 'Bob'), ROW(3, 'Charlie')] AS struct_array", - "struct_array", - ), - ( - "SELECT ARRAY[ROW('name', 'John', 25), ROW('name', 'Jane', 30)] " - "AS unnamed_struct_array", - "unnamed_struct_array", - ), - ( - "SELECT ARRAY[CAST(ROW('John', 25) AS ROW(name VARCHAR, age INT)), " - "CAST(ROW('Jane', 30) AS ROW(name VARCHAR, age INT))] AS named_struct_array", - "named_struct_array", - ), - ], - ) - def test_array_types_with_structs(self, cursor, query, description): - """Test ARRAY types containing STRUCT elements.""" - _logger.info("=== ARRAY with STRUCT Test: %s ===", description) - cursor.execute(query) - result = cursor.fetchone() - array_value = result[0] - _logger.info("%s: %r (type: %s)", description, array_value, type(array_value).__name__) - - # Validate array value - assert array_value is not None, f"ARRAY value should not be None for {description}" - assert isinstance(array_value, list), f"ARRAY value should be list for {description}" - assert len(array_value) > 0, f"ARRAY should not be empty for {description}" - - # Check first element is a dict (converted struct) - first_element = array_value[0] - assert isinstance(first_element, dict), ( - f"First array element should be dict (struct) for {description}" - ) - _logger.info("%s: First element: %r", description, first_element) - - @pytest.mark.parametrize( - ("query", "description"), - [ - ( - "SELECT CAST(ARRAY[1, 2, 3] AS JSON) AS arr_json", - "arr_json", - ), - ( - "SELECT CAST(ARRAY['a', 'b', 'c'] AS JSON) AS str_arr_json", - "str_arr_json", - ), - ( - "SELECT CAST(ARRAY[ROW(1, 'Alice'), ROW(2, 'Bob')] AS JSON) AS struct_arr_json", - "struct_arr_json", - ), - ], - ) - def test_array_types_json_cast(self, cursor, query, description): - """Test ARRAY types with JSON casting.""" - _logger.info("=== ARRAY JSON Cast Test: %s ===", description) - cursor.execute(query) - result = cursor.fetchone() - array_value = result[0] - _logger.info("%s: %r (type: %s)", description, array_value, type(array_value).__name__) - - # Validate array value - assert array_value is not None, f"ARRAY value should not be None for {description}" - assert isinstance(array_value, list), ( - f"JSON cast ARRAY value should be list for {description}" - ) - _logger.info("%s: JSON cast array type %s", description, type(array_value).__name__) - - @pytest.mark.parametrize( - ("query", "description"), - [ - ("SELECT CARDINALITY(ARRAY[1, 2, 3, 4, 5]) AS array_size", "array_size"), - ("SELECT ARRAY[10, 20, 30, 40][2] AS array_element", "array_element"), - ("SELECT ARRAY[1, 2] || ARRAY[3, 4] AS array_concat", "array_concat"), - ("SELECT CONTAINS(ARRAY[1, 2, 3], 2) AS array_contains", "array_contains"), - ], - ) - def test_array_operations(self, cursor, query, description): - """Test ARRAY operations and functions.""" - _logger.info("=== ARRAY Operation Test: %s ===", description) - cursor.execute(query) - result = cursor.fetchone() - operation_result = result[0] - _logger.info( - "%s: %r (type: %s)", - description, - operation_result, - type(operation_result).__name__, - ) - - # Validate operation result - assert operation_result is not None, ( - f"ARRAY operation result should not be None for {description}" - ) - - # Type-specific validations - if description == "array_size": - assert operation_result == 5, f"Array size should be 5 for {description}" - elif description == "array_element": - assert operation_result == 20, f"Array element [2] should be 20 for {description}" - elif description == "array_concat": - assert operation_result == [1, 2, 3, 4], ( - f"Array concat should be [1,2,3,4] for {description}" - ) - elif description == "array_contains": - assert operation_result is True, f"Array contains should be True for {description}" - - def test_array_converter_behavior(self, cursor): - """Test ARRAY converter behavior with different formats.""" - _logger.info("=== ARRAY Converter Behavior Test ===") - - # Test simple array conversion - cursor.execute("SELECT ARRAY[1, 2, 3] AS simple") - result = cursor.fetchone() - simple_array = result[0] - _logger.info("Simple array: %r", simple_array) - assert simple_array == [1, 2, 3] - - # Test array with struct conversion - cursor.execute( - "SELECT ARRAY[CAST(ROW(1, 2) AS ROW(a INT, b INT)), " - "CAST(ROW(3, 4) AS ROW(a INT, b INT))] AS struct_array" - ) - result = cursor.fetchone() - struct_array = result[0] - _logger.info("Struct array: %r", struct_array) - assert isinstance(struct_array, list) - assert len(struct_array) == 2 - assert isinstance(struct_array[0], dict) - - # Test converter function directly - test_cases = [ - ("[1, 2, 3]", [1, 2, 3]), - ('["a", "b", "c"]', ["a", "b", "c"]), - ("[{a=1, b=2}]", [{"a": "1", "b": "2"}]), - ("[]", []), - (None, None), - ("invalid", None), - ] - - for test_input, expected in test_cases: - result = _to_array(test_input) - _logger.info("Converter test: %r -> %r", test_input, result) - assert result == expected, ( - f"Converter failed for {test_input}: expected {expected}, got {result}" - ) diff --git a/tests/pyathena/test_formatter.py b/tests/pyathena/test_formatter.py index 8290ee62d..2e40f4df4 100644 --- a/tests/pyathena/test_formatter.py +++ b/tests/pyathena/test_formatter.py @@ -12,7 +12,6 @@ import pytest from pyathena.error import ProgrammingError -from pyathena.formatter import _escape_trino # A value whose single quote must not be allowed to terminate the string # literal. Trino and Athena do not treat a backslash as an escape character @@ -599,7 +598,6 @@ def test_table_name_starting_with_as_is_still_hive_ddl(self, formatter): ) -def test_escape_trino_doubles_single_quotes(): - assert _escape_trino("a'b") == "'a''b'" +def test_backslash_before_quote_stays_literal_in_trino_statement(formatter): # A backslash is not an escape character in Trino/Presto; it stays literal. - assert _escape_trino("a\\'b") == "'a\\''b'" + assert formatter.format("SELECT %(v)s", {"v": "a\\'b"}) == "SELECT 'a\\''b'" diff --git a/tests/pyathena/test_glue.py b/tests/pyathena/test_glue.py index eecd3443b..1f59d7f2e 100644 --- a/tests/pyathena/test_glue.py +++ b/tests/pyathena/test_glue.py @@ -4,18 +4,14 @@ # See LICENSE or https://opensource.org/licenses/MIT. # # SPDX-License-Identifier: MIT -import threading import uuid -from concurrent.futures import ThreadPoolExecutor import pytest -from botocore.exceptions import ClientError, ParamValidationError -from botocore.exceptions import ConnectionError as BotoConnectionError +from botocore.exceptions import ParamValidationError from pyathena.glue import GlueMetadataClient from tests import ENV from tests.pyathena.conftest import connect -from tests.pyathena.util import unreachable_glue class TestGlueMetadataClient: @@ -43,20 +39,6 @@ def _view(metadata): metadata.parameters, ) - def test_reads_what_athena_reports(self, cursor): - glue = cursor.connection._glue - catalog = cursor.connection.catalog_name - - for table in ("one_row", "parquet_with_compression", "partition_table", "view_one_row"): - assert self._view(glue.get_table(catalog, ENV.schema, table)) == self._view( - cursor.get_table_metadata(table) - ) - assert sorted(self._view(m) for m in glue.list_tables(catalog, ENV.schema)) == sorted( - self._view(m) for m in cursor.list_table_metadata() - ) - assert [m.name for m in glue.list_tables(catalog, ENV.schema, "one_row")] == ["one_row"] - assert ENV.schema in [d.name for d in glue.list_databases(catalog)] - @pytest.mark.skipif( not ENV.s3tables_catalog, reason="AWS_ATHENA_S3_TABLES_CATALOG is not configured", @@ -91,28 +73,6 @@ def test_reads_s3_tables_catalog(self): finally: cursor.execute(f"DROP TABLE IF EXISTS {schema}.{table}") - def test_reports_a_missing_table(self, cursor): - with pytest.raises(ClientError) as caught: - cursor.connection._glue.get_table( - cursor.connection.catalog_name, ENV.schema, "no_such_table_786" - ) - - assert caught.value.response["Error"]["Code"] == "EntityNotFoundException" - assert cursor.connection._glue.reachable - - def test_rejects_an_unsupported_catalog(self, cursor): - with pytest.raises(ValueError, match="federated_catalog"): - cursor.connection._glue.get_table("federated_catalog", ENV.schema, "one_row") - - def test_stops_after_it_cannot_reach_glue(self, cursor): - glue = unreachable_glue(cursor.connection) - - with pytest.raises(BotoConnectionError): - glue.get_table("AwsDataCatalog", ENV.schema, "one_row") - - assert not glue.reachable - assert not glue.usable_for("AwsDataCatalog") - def test_keeps_going_after_a_request_it_rejects(self, cursor): # botocore rejects the request before sending it; Glue stays in use. glue = cursor.connection._glue @@ -129,28 +89,11 @@ def test_client_leaves_out_athena_endpoint(self): # Athena's API version, which Glue does not have. api_version="2017-05-18", ) as conn: - glue = conn._glue - created = [] - session_client = conn.session.client - - def contended_client(*args, **kwargs): - created.append((args, glue._lock.locked())) - return session_client(*args, **kwargs) - - conn._session.client = contended_client - barrier = threading.Barrier(8) - - def get_client(_): - barrier.wait() - return glue.client - - with ThreadPoolExecutor(max_workers=8) as executor: - clients = list(executor.map(get_client, range(8))) + client = conn._glue.client - assert created == [(("glue",), True)] - assert all(client is clients[0] for client in clients) - assert clients[0].meta.service_model.service_name == "glue" - assert clients[0].meta.endpoint_url == f"https://glue.{ENV.region_name}.amazonaws.com" + assert client.meta.service_model.service_name == "glue" + assert client.meta.service_model.api_version == "2017-03-31" + assert client.meta.endpoint_url == f"https://glue.{ENV.region_name}.amazonaws.com" def test_table_metadata_leaves_out_columns_iceberg_no_longer_has(self): # Glue's columns after DROP COLUMN b and CHANGE COLUMN a a2, as measured diff --git a/tests/pyathena/test_options.py b/tests/pyathena/test_options.py index 2ea323c09..731b3d2c6 100644 --- a/tests/pyathena/test_options.py +++ b/tests/pyathena/test_options.py @@ -37,39 +37,6 @@ def test_merge_ignores_none_overrides(self): assert merged.cache_size == 100 assert merged.s3_staging_dir is None - def test_merge_overrides_take_precedence(self): - options = ExecuteOptions( - work_group="primary", - cache_size=100, - result_reuse_enable=True, - paramstyle="pyformat", - ) - merged = options.merge( - work_group="adhoc", - cache_size=0, - result_reuse_enable=False, - paramstyle="qmark", - ) - assert merged.work_group == "adhoc" - assert merged.cache_size == 0 - assert merged.result_reuse_enable is False - assert merged.paramstyle == "qmark" - - def test_merge_does_not_mutate_original(self): - options = ExecuteOptions(work_group="primary") - merged = options.merge(work_group="adhoc") - assert options.work_group == "primary" - assert merged is not options - - def test_merge_without_applied_overrides_is_equal(self): - options = ExecuteOptions(work_group="primary") - assert options.merge() == options - assert options.merge(work_group=None) == options - - def test_merge_raises_on_unknown_field(self): - with pytest.raises(TypeError, match="unknown_field"): - ExecuteOptions().merge(unknown_field="value") - @pytest.mark.parametrize( ("base", "overrides", "expected"), [ @@ -100,10 +67,6 @@ def test_merge_precedence(self, base, overrides, expected): for name, value in expected.items(): assert getattr(merged, name) == value - def test_resolve_with_none_returns_defaults_with_overrides(self): - resolved = ExecuteOptions.resolve(None, work_group="wg") - assert resolved == ExecuteOptions(work_group="wg") - def test_resolve_applies_overrides_to_given_options(self): options = ExecuteOptions(work_group="primary", cache_size=100) resolved = ExecuteOptions.resolve(options, work_group="adhoc", cache_size=None) @@ -127,39 +90,11 @@ def test_resolve_without_base_preserves_override_values(self, overrides, expecte assert resolved.cache_expiration_time == 0 assert resolved.result_reuse_minutes is None - def test_resolve_returns_distinct_default_instances(self): - first = ExecuteOptions.resolve(None) - second = ExecuteOptions.resolve(None) - assert first == second - assert first is not second - - def test_resolve_preserves_callback_and_hint_references(self): - def callback(query_id): - pass - - first_hints = {"tags": "array(varchar)"} - second_hints = {"tags": "map(varchar, integer)"} - first = ExecuteOptions.resolve( - None, on_start_query_execution=callback, result_set_type_hints=first_hints - ) - second = ExecuteOptions.resolve(None, result_set_type_hints=second_hints) - assert first.on_start_query_execution is callback - assert first.result_set_type_hints is first_hints - assert second.result_set_type_hints is second_hints - first_hints["tags"] = "array(bigint)" - assert second.result_set_type_hints == {"tags": "map(varchar, integer)"} - @pytest.mark.parametrize("options", [None, ExecuteOptions(work_group="primary")]) def test_resolve_rejects_unknown_non_none_fields(self, options): with pytest.raises(TypeError, match="unknown_field"): ExecuteOptions.resolve(options, unknown_field="value") - def test_resolve_given_options_preserves_new_instance_behavior(self): - options = ExecuteOptions(work_group="primary") - resolved = ExecuteOptions.resolve(options, work_group=None) - assert resolved == options - assert resolved is not options - @pytest.mark.parametrize("provided", [False, True]) def test_resolve_retains_subclass_merge_dispatch(self, provided): @dataclasses.dataclass(frozen=True) diff --git a/tests/pyathena/test_parser.py b/tests/pyathena/test_parser.py index dd3b827f1..e1c776423 100644 --- a/tests/pyathena/test_parser.py +++ b/tests/pyathena/test_parser.py @@ -10,7 +10,6 @@ from pyathena.converter import _DEFAULT_CONVERTERS, _to_default, _to_struct from pyathena.parser import ( TypedValueConverter, - TypeNode, TypeSignatureParser, ) @@ -47,12 +46,6 @@ def test_array_type(self): assert len(node.children) == 1 assert node.children[0].type_name == "varchar" - def test_array_of_integer(self): - parser = TypeSignatureParser() - node = parser.parse("array(integer)") - assert node.type_name == "array" - assert node.children[0].type_name == "integer" - def test_map_type(self): parser = TypeSignatureParser() node = parser.parse("map(varchar, integer)") @@ -106,16 +99,6 @@ def test_deeply_nested(self): assert row_node.children[0].field_names == ["x", "y"] assert row_node.children[1].type_name == "varchar" - def test_parameterized_type(self): - parser = TypeSignatureParser() - node = parser.parse("decimal(10, 2)") - assert node.type_name == "decimal" - - def test_varchar_with_length(self): - parser = TypeSignatureParser() - node = parser.parse("varchar(255)") - assert node.type_name == "varchar" - def test_type_alias_int(self): parser = TypeSignatureParser() node = parser.parse("int") @@ -127,13 +110,6 @@ def test_type_alias_in_complex_type(self): assert node.type_name == "array" assert node.children[0].type_name == "integer" - def test_trailing_modifier_after_paren(self): - """Type with content after closing paren should not break parsing.""" - parser = TypeSignatureParser() - # Simulates a hypothetical "timestamp(3) with time zone" style input - node = parser.parse("decimal(10, 2) extra") - assert node.type_name == "decimal" - class TestTypedValueConverter: @pytest.fixture @@ -144,40 +120,11 @@ def converter(self): struct_parser=_to_struct, ) - def test_simple_varchar(self, converter): - node = TypeNode("varchar") - assert converter.convert("hello", node) == "hello" - - def test_simple_integer(self, converter): - node = TypeNode("integer") - assert converter.convert("42", node) == 42 - - def test_array_of_varchar(self, converter): - parser = TypeSignatureParser() - node = parser.parse("array(varchar)") - assert converter.convert("[1234, 5678]", node) == ["1234", "5678"] - - def test_array_of_integer(self, converter): - parser = TypeSignatureParser() - node = parser.parse("array(integer)") - assert converter.convert("[1, 2, 3]", node) == [1, 2, 3] - def test_map_varchar_integer(self, converter): parser = TypeSignatureParser() node = parser.parse("map(varchar, integer)") assert converter.convert('{"a": 1, "b": 2}', node) == {"a": 1, "b": 2} - def test_row_named_fields(self, converter): - parser = TypeSignatureParser() - node = parser.parse("row(name varchar, age integer)") - assert converter.convert("{name=Alice, age=25}", node) == {"name": "Alice", "age": 25} - - def test_nested_row(self, converter): - parser = TypeSignatureParser() - node = parser.parse("row(header row(seq integer, stamp varchar), x double)") - result = converter.convert("{header={seq=123, stamp=2024}, x=4.5}", node) - assert result == {"header": {"seq": 123, "stamp": "2024"}, "x": 4.5} - def test_array_of_row_json(self, converter): """JSON path: array(row(...)) with nested dict elements.""" parser = TypeSignatureParser() diff --git a/tests/pyathena/test_result_set.py b/tests/pyathena/test_result_set.py index 0612754f1..d901cdae5 100644 --- a/tests/pyathena/test_result_set.py +++ b/tests/pyathena/test_result_set.py @@ -9,12 +9,10 @@ import pytest -from pyathena.aio.common import AioBaseCursor, WithAsyncFetch -from pyathena.common import BaseCursor, CursorIterator from pyathena.converter import DefaultTypeConverter -from pyathena.error import OperationalError, ProgrammingError +from pyathena.error import OperationalError from pyathena.model import AthenaQueryExecution -from pyathena.result_set import AthenaResultSet, WithFetch, WithResultSet +from pyathena.result_set import AthenaResultSet from pyathena.util import RetryConfig from tests.pyathena.util import assert_fetches_fail, get_query_results_responses @@ -107,12 +105,6 @@ def test_fetch_after_close(self): assert result_set.fetchall() == [] assert result_set.rownumber is None - def test_get_query_results_after_close(self): - result_set = _fetching_result_set(("1", "2"), ("3",)) - result_set.close() - with pytest.raises(ProgrammingError, match=r"AthenaResultSet is closed\."): - result_set._get_query_results(2, "token1") - def test_fetchone_across_pages(self): result_set = _fetching_result_set(("1", "2"), ("3",)) assert [result_set.fetchone() for _ in range(4)] == [(1,), (2,), (3,), None] @@ -159,16 +151,3 @@ def test_fetch_all_rows_as_csv(self): with patch.object(AthenaResultSet, "_get_query_results", side_effect=pages): data = result_set._fetch_all_rows_as_csv() assert data == b'"v","w"""\n"a,""b""\nc",\n"v","w"""\n"","1"\n' - - -class TestWithResultSet: - def test_is_mixin(self): - assert WithResultSet.__bases__ == (object,) - - @pytest.mark.parametrize( - ("cursor_base", "base"), - [(WithFetch, BaseCursor), (WithAsyncFetch, AioBaseCursor)], - ) - def test_precedes_cursor_bases(self, cursor_base, base): - # Listed first, so that its members take precedence over the cursor bases'. - assert cursor_base.__bases__ == (WithResultSet, base, CursorIterator) diff --git a/tests/pyathena/test_util.py b/tests/pyathena/test_util.py index f3c93759e..062f67e9d 100644 --- a/tests/pyathena/test_util.py +++ b/tests/pyathena/test_util.py @@ -1,6 +1,5 @@ import threading import time -from typing import Any import pytest import tenacity.nap @@ -79,41 +78,12 @@ def test_strtobool_invalid(value): strtobool(value) -class _WithCodeError(Exception): - def __init__(self, code: int) -> None: - super().__init__(f"error:{code}") - self.response = {"Error": {"Code": code}} - - class _NoResponseError(Exception): def __init__(self) -> None: super().__init__("error") self.response = None -def _test_retry(ex: Exception) -> None: - calls = {"n": 0} - - def fn() -> Any: - calls["n"] += 1 - raise ex - - cfg = RetryConfig(attempt=1, max_delay=1) - - with pytest.raises(type(ex)): - retry_api_call(fn, config=cfg) - - assert calls["n"] == 1 - - -def test_retry_api_call(): - _test_retry(_WithCodeError(500)) - - -def test_retry_api_call_with_none_error(): - _test_retry(_NoResponseError()) - - @pytest.mark.parametrize( ("code", "message", "exceptions", "expected_calls"), [ @@ -251,23 +221,11 @@ def call(): retry_api_call(call, config) assert calls == 7 assert len(sleeps) == len(expected_bases) - for waited, base in zip(sleeps, expected_bases, strict=True): - assert base <= waited <= base + multiplier - - -def test_retry_api_call_jitter_varies(monkeypatch): - sleeps = [] - monkeypatch.setattr(tenacity.nap.time, "sleep", sleeps.append) - error = throttling_error("GetTableMetadata") - - def call(): - raise error - - config = RetryConfig(attempt=20, multiplier=1, max_delay=1) - with pytest.raises(ClientError): - retry_api_call(call, config) - assert all(1 <= waited <= 2 for waited in sleeps) - assert len(set(sleeps)) > 1 + jitters = [waited - base for waited, base in zip(sleeps, expected_bases, strict=True)] + assert all(0 <= jitter <= multiplier for jitter in jitters) + if multiplier: + # Random, so the six waits do not add the same jitter. + assert len(set(jitters)) > 1 @pytest.mark.parametrize( @@ -293,6 +251,7 @@ def test_is_retryable_error(code, message, expected): error = ClientError({"Error": {"Code": code, "Message": message}}, "GetTableMetadata") assert is_retryable_error(error, RetryConfig()) is expected assert is_retryable_error(ValueError("no response"), RetryConfig()) is False + assert is_retryable_error(_NoResponseError(), RetryConfig()) is False @pytest.mark.parametrize( diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 8d3a08412..7d9147787 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -266,6 +266,30 @@ def fail(**kwargs): return calls +def succeeded_query_execution_data(query_id, query, completion_date_time, schema="this_schema"): + """Build a succeeded DML query execution as ``BatchGetQueryExecution`` returns it. + + Args: + query_id: The query execution ID. + query: The query string. + completion_date_time: The completion time. + schema: The database the query ran against. + + Returns: + The ``QueryExecution`` dictionary. + """ + return { + "QueryExecutionId": query_id, + "Query": query, + "StatementType": AthenaQueryExecution.STATEMENT_TYPE_DML, + "QueryExecutionContext": {"Database": schema}, + "Status": { + "State": AthenaQueryExecution.STATE_SUCCEEDED, + "CompletionDateTime": completion_date_time, + }, + } + + def succeeded_query_execution(query_id, query, completion_date_time, schema="this_schema"): """Build a succeeded DML query execution, as the result cache search lists it. @@ -280,20 +304,67 @@ def succeeded_query_execution(query_id, query, completion_date_time, schema="thi """ return AthenaQueryExecution( { - "QueryExecution": { - "QueryExecutionId": query_id, - "Query": query, - "StatementType": AthenaQueryExecution.STATEMENT_TYPE_DML, - "QueryExecutionContext": {"Database": schema}, - "Status": { - "State": AthenaQueryExecution.STATE_SUCCEEDED, - "CompletionDateTime": completion_date_time, - }, - } + "QueryExecution": succeeded_query_execution_data( + query_id, query, completion_date_time, schema + ) } ) +class QueryExecutionHistory: + """An Athena client stand-in that lists the query executions of a work group. + + ``ListQueryExecutions`` returns up to ``MaxResults`` IDs per page in the given + order, latest first as Athena lists them, and ``BatchGetQueryExecution`` + returns their executions. + + Attributes: + executions: The ``QueryExecution`` dictionaries, latest first. + list_requests: The number of ``ListQueryExecutions`` requests. + """ + + def __init__(self, executions): + """Initialize the history. + + Args: + executions: The ``QueryExecution`` dictionaries, latest first. + """ + self.executions = executions + self.list_requests = 0 + + def list_query_executions(self, **request): + """List a page of query execution IDs. + + Args: + **request: The request parameters: ``MaxResults``, the ``NextToken`` + of a later page, and others such as ``WorkGroup``. + + Returns: + The response, with a ``NextToken`` unless the page is the last. + """ + self.list_requests += 1 + start = int(request.get("NextToken", 0)) + end = start + request["MaxResults"] + response = { + "QueryExecutionIds": [e["QueryExecutionId"] for e in self.executions[start:end]] + } + if end < len(self.executions): + response["NextToken"] = str(end) + return response + + def batch_get_query_execution(self, **request): + """Get the executions of query execution IDs. + + Args: + **request: The request parameters, with ``QueryExecutionIds``. + + Returns: + The response. + """ + executions = {e["QueryExecutionId"]: e for e in self.executions} + return {"QueryExecutions": [executions[i] for i in request["QueryExecutionIds"]]} + + def unreachable_glue(connection): """A Glue client for the connection that sends requests to a closed proxy port.""" return GlueMetadataClient( From cf0d76f8ad71266993f7edfc24c7268b75e420aa Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 21:34:24 +0900 Subject: [PATCH 2/5] Keep the checks that the reviewed covers did not make The typed-to-untyped fallback test converts a value that its hint does not describe instead of a value the typed path cannot parse yet. The cache search tests tell a stopped search from a failed one and check the work group of each request. The type-hint test compares reprs, so 123 and 123.0 differ. The Glue client test keeps its check that eight threads build one client. The no-response checks of is_retryable_error get their own test. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/test_converter.py | 9 +++++---- tests/pyathena/test_cursor.py | 19 +++++++++++++++---- tests/pyathena/test_glue.py | 28 ++++++++++++++++++++++++---- tests/pyathena/test_util.py | 7 +++++-- tests/pyathena/util.py | 6 +++--- 5 files changed, 52 insertions(+), 17 deletions(-) diff --git a/tests/pyathena/test_converter.py b/tests/pyathena/test_converter.py index ac334a80a..5f8507c7f 100644 --- a/tests/pyathena/test_converter.py +++ b/tests/pyathena/test_converter.py @@ -394,11 +394,12 @@ def test_row_with_nested_struct(self): assert result == {"header": {"seq": 123, "stamp": "2024"}, "x": 4.5} def test_fallback_returns_untyped_result(self): - """A value that the typed conversion cannot parse is converted without the hint.""" + """A value that the type hint does not describe is converted without the hint.""" converter = DefaultTypeConverter() - # The untyped conversion keeps the original string of nested native arrays. - value = "[{a=[1, 2]}, {b=[3]}]" - assert converter.convert("array", value, type_hint="array(row(a array(integer)))") == value + assert converter.convert("row", "{a=1, b=2}", type_hint="array(integer)") == { + "a": "1", + "b": "2", + } def test_normalize_hive_syntax_struct(self): converter = DefaultTypeConverter() diff --git a/tests/pyathena/test_cursor.py b/tests/pyathena/test_cursor.py index af0e77c7b..00566dc9e 100644 --- a/tests/pyathena/test_cursor.py +++ b/tests/pyathena/test_cursor.py @@ -427,7 +427,12 @@ def test_cache_search_stops_at_expired_execution(self): # Without cache_size, cache_expiration_time alone bounds the search. assert cursor._find_previous_query_id(query, None, cache_expiration_time=3600) is None - assert history.list_requests == 1 + assert len(history.list_requests) == 1 + # A failed search would also return None, but the same search reuses the + # execution once it is within cache_expiration_time. + assert ( + cursor._find_previous_query_id(query, None, cache_expiration_time=3 * 3600) == "expired" + ) @pytest.mark.parametrize( ("position", "expected"), @@ -442,9 +447,11 @@ def test_cache_search_reads_pages_up_to_cache_size(self, position, expected): succeeded_query_execution_data(f"other_{i}", "SELECT 1", now) for i in range(70) ] executions[position] = succeeded_query_execution_data("match", query, now) - cursor = _cache_search_cursor(QueryExecutionHistory(executions)) + history = QueryExecutionHistory(executions) + cursor = _cache_search_cursor(history) assert cursor._find_previous_query_id(query, "wg", cache_size=60) == expected + assert [request["WorkGroup"] for request in history.list_requests] == ["wg", "wg"] def test_cache_search_prefers_latest_execution(self): query = "SELECT * FROM one_row" @@ -907,7 +914,7 @@ def test_complex_with_type_hints(self, cursor): "data": "array(row(pos row(x integer, y integer), vel row(x double)))", }, ) - assert cursor.fetchone() == ( + expected = ( [1, 2], {"1": 2, "3": 4}, {"a": 1, "b": 2}, @@ -917,6 +924,10 @@ def test_complex_with_type_hints(self, cursor): {"pos": {"x": 3, "y": 4}, "vel": {"x": 1.5}}, ], ) + row = cursor.fetchone() + assert row == expected + # The reprs also tell 123 from 123.0. + assert repr(row) == repr(expected) def test_cancel(self, cursor): def cancel(c): @@ -1652,7 +1663,7 @@ def read(): return ( [self._metadata_view(cursor.get_table_metadata(t)) for t in tables], sorted(self._metadata_view(m) for m in cursor.list_table_metadata()), - [m.name for m in cursor.list_table_metadata(expression="one_row")], + sorted(m.name for m in cursor.list_table_metadata(expression="one_row")), ENV.schema in [d.name for d in cursor.list_databases("AwsDataCatalog")], ) diff --git a/tests/pyathena/test_glue.py b/tests/pyathena/test_glue.py index 1f59d7f2e..ab7d7693b 100644 --- a/tests/pyathena/test_glue.py +++ b/tests/pyathena/test_glue.py @@ -4,7 +4,9 @@ # See LICENSE or https://opensource.org/licenses/MIT. # # SPDX-License-Identifier: MIT +import threading import uuid +from concurrent.futures import ThreadPoolExecutor import pytest from botocore.exceptions import ParamValidationError @@ -89,11 +91,29 @@ def test_client_leaves_out_athena_endpoint(self): # Athena's API version, which Glue does not have. api_version="2017-05-18", ) as conn: - client = conn._glue.client + glue = conn._glue + created = [] + session_client = conn.session.client - assert client.meta.service_model.service_name == "glue" - assert client.meta.service_model.api_version == "2017-03-31" - assert client.meta.endpoint_url == f"https://glue.{ENV.region_name}.amazonaws.com" + def contended_client(*args, **kwargs): + created.append(args) + return session_client(*args, **kwargs) + + conn.session.client = contended_client + barrier = threading.Barrier(8) + + def get_client(_): + barrier.wait() + return glue.client + + with ThreadPoolExecutor(max_workers=8) as executor: + clients = list(executor.map(get_client, range(8))) + + assert created == [("glue",)] + assert all(client is clients[0] for client in clients) + assert clients[0].meta.service_model.service_name == "glue" + assert clients[0].meta.service_model.api_version == "2017-03-31" + assert clients[0].meta.endpoint_url == f"https://glue.{ENV.region_name}.amazonaws.com" def test_table_metadata_leaves_out_columns_iceberg_no_longer_has(self): # Glue's columns after DROP COLUMN b and CHANGE COLUMN a a2, as measured diff --git a/tests/pyathena/test_util.py b/tests/pyathena/test_util.py index 062f67e9d..d932b9a39 100644 --- a/tests/pyathena/test_util.py +++ b/tests/pyathena/test_util.py @@ -250,8 +250,11 @@ def call(): def test_is_retryable_error(code, message, expected): error = ClientError({"Error": {"Code": code, "Message": message}}, "GetTableMetadata") assert is_retryable_error(error, RetryConfig()) is expected - assert is_retryable_error(ValueError("no response"), RetryConfig()) is False - assert is_retryable_error(_NoResponseError(), RetryConfig()) is False + + +@pytest.mark.parametrize("error", [ValueError("no response"), _NoResponseError()]) +def test_is_retryable_error_without_response(error): + assert is_retryable_error(error, RetryConfig()) is False @pytest.mark.parametrize( diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 7d9147787..4c1bee3c5 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -320,7 +320,7 @@ class QueryExecutionHistory: Attributes: executions: The ``QueryExecution`` dictionaries, latest first. - list_requests: The number of ``ListQueryExecutions`` requests. + list_requests: The parameters of each ``ListQueryExecutions`` request. """ def __init__(self, executions): @@ -330,7 +330,7 @@ def __init__(self, executions): executions: The ``QueryExecution`` dictionaries, latest first. """ self.executions = executions - self.list_requests = 0 + self.list_requests = [] def list_query_executions(self, **request): """List a page of query execution IDs. @@ -342,7 +342,7 @@ def list_query_executions(self, **request): Returns: The response, with a ``NextToken`` unless the page is the last. """ - self.list_requests += 1 + self.list_requests.append(request) start = int(request.get("NextToken", 0)) end = start + request["MaxResults"] response = { From 9daa4b1243098f871379714e41ca12970442a19d Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 21:37:52 +0900 Subject: [PATCH 3/5] Leave out the response fields and comparisons the tests do not need Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/test_connection.py | 5 ----- tests/pyathena/test_cursor.py | 9 ++++----- 2 files changed, 4 insertions(+), 10 deletions(-) diff --git a/tests/pyathena/test_connection.py b/tests/pyathena/test_connection.py index 577fd8597..06ce2772e 100644 --- a/tests/pyathena/test_connection.py +++ b/tests/pyathena/test_connection.py @@ -187,11 +187,6 @@ def test_cursor_unknown_constructor_keyword(offline_connection, cursor_class): "Query": "SELECT 1", "StatementType": "DML", "Status": {"State": "FAILED", "StateChangeReason": "failed"}, - "Statistics": {}, - "ResultConfiguration": {}, - "QueryExecutionContext": {}, - "ResultReuseConfiguration": {"ResultReuseByAgeConfiguration": {"Enabled": False}}, - "EngineVersion": {}, } }, "StartCalculationExecution": {"CalculationExecutionId": "calculation-id", "State": "CREATING"}, diff --git a/tests/pyathena/test_cursor.py b/tests/pyathena/test_cursor.py index 00566dc9e..d4fbb7cee 100644 --- a/tests/pyathena/test_cursor.py +++ b/tests/pyathena/test_cursor.py @@ -412,12 +412,13 @@ def execution(catalog): def test_cache_search_stops_at_expired_execution(self): query = "SELECT * FROM one_row" now = datetime.now(UTC) + page_size = Cursor.LIST_QUERY_EXECUTIONS_MAX_RESULTS executions = [ succeeded_query_execution_data("other", "SELECT 1", now), succeeded_query_execution_data("expired", query, now - timedelta(hours=2)), *( succeeded_query_execution_data(f"older_{i}", "SELECT 1", now - timedelta(hours=3)) - for i in range(48) + for i in range(page_size - 2) ), # A search that went on to the second page would return this match. succeeded_query_execution_data("next_page", query, now), @@ -924,10 +925,8 @@ def test_complex_with_type_hints(self, cursor): {"pos": {"x": 3, "y": 4}, "vel": {"x": 1.5}}, ], ) - row = cursor.fetchone() - assert row == expected - # The reprs also tell 123 from 123.0. - assert repr(row) == repr(expected) + # Unlike ==, the reprs also tell 123 from 123.0. + assert repr(cursor.fetchone()) == repr(expected) def test_cancel(self, cursor): def cancel(c): From c94aac06ab1ffb58fee2781ab03605e373a1b909 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 21:58:28 +0900 Subject: [PATCH 4/5] Keep the checks that the independent review found unprotected Restore the DictCursor NULL and empty-string test, whose result set converts its own rows, and the test that merge() leaves the original options unchanged, as docs/usage.md states. Add native date and timestamp arrays to the live complex-value query and the parser test. Compare the type-hinted row with ==, plus the types of its values, so the order of map entries does not matter. Round the jitters, and let the build-once checks wait for a second build. Answer each stubbed request once, and reject a ListQueryExecutions page above 50, as Athena does. Pass connect() one positional argument, so a factory that accepts one is caught. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/test_connection.py | 43 +++++-------------- tests/pyathena/test_converter.py | 2 + tests/pyathena/test_cursor.py | 68 +++++++++++++++++++++++++++++-- tests/pyathena/test_glue.py | 7 ++++ tests/pyathena/test_options.py | 6 +++ tests/pyathena/test_util.py | 5 ++- tests/pyathena/util.py | 63 ++++++++++++++++++++++++++++ 7 files changed, 156 insertions(+), 38 deletions(-) diff --git a/tests/pyathena/test_connection.py b/tests/pyathena/test_connection.py index 06ce2772e..8d25856ef 100644 --- a/tests/pyathena/test_connection.py +++ b/tests/pyathena/test_connection.py @@ -35,7 +35,7 @@ from pyathena.spark.common import SparkBaseCursor from pyathena.spark.cursor import SparkCursor from pyathena.util import RetryConfig -from tests.pyathena.util import EVENT_TIMEOUT +from tests.pyathena.util import EVENT_TIMEOUT, answered_requests # Cursors whose arraysize is the GetQueryResults page size, capped at 1000. PAGED_CURSORS = [Cursor, DictCursor, AsyncCursor, AsyncDictCursor] @@ -95,7 +95,7 @@ def test_empty_settings_disable_both_environment_fallbacks(monkeypatch): @pytest.mark.parametrize("factory", [pyathena.connect, Connection]) def test_connect_keyword_only(factory): with pytest.raises(TypeError): - factory("s3://bucket/path/", "us-east-1") + factory("s3://bucket/path/") class RecordingCursor(Cursor): @@ -178,32 +178,6 @@ def test_cursor_unknown_constructor_keyword(offline_connection, cursor_class): offline_connection.cursor(cursor_class) -# Athena's answers to a query or calculation that fails at its first status request. -FAILED_EXECUTION_RESPONSES = { - "StartQueryExecution": {"QueryExecutionId": "query-id"}, - "GetQueryExecution": { - "QueryExecution": { - "QueryExecutionId": "query-id", - "Query": "SELECT 1", - "StatementType": "DML", - "Status": {"State": "FAILED", "StateChangeReason": "failed"}, - } - }, - "StartCalculationExecution": {"CalculationExecutionId": "calculation-id", "State": "CREATING"}, - "GetCalculationExecutionStatus": { - "Status": {"State": "FAILED"}, - "Statistics": {"DpuExecutionInMillis": 0}, - }, - "GetCalculationExecution": { - "CalculationExecutionId": "calculation-id", - "SessionId": "session", - "Status": {"State": "FAILED", "StateChangeReason": "failed"}, - "Statistics": {"DpuExecutionInMillis": 0}, - "Result": {}, - }, -} - - @pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) def test_cursor_accepts_connection_callbacks(offline_connection, cursor_class): started, polled = [], [] @@ -211,11 +185,7 @@ def test_cursor_accepts_connection_callbacks(offline_connection, cursor_class): offline_connection.on_poll = polled.append with ( - patch.object( - offline_connection.client, - "_make_api_call", - side_effect=lambda operation, _: FAILED_EXECUTION_RESPONSES[operation], - ), + answered_requests(offline_connection.client), _open_cursor(offline_connection, cursor_class) as cursor, ): try: @@ -350,8 +320,15 @@ def test_s3_client_built_once_across_threads(self): created = [] session_client = conn.session.client + second_build = threading.Event() + def contended_client(*args, **kwargs): created.append(args) + # Without the lock, another thread would build a client meanwhile. + if len(created) == 1: + second_build.wait(1) + else: + second_build.set() return session_client(*args, **kwargs) conn.session.client = contended_client diff --git a/tests/pyathena/test_converter.py b/tests/pyathena/test_converter.py index 5f8507c7f..f96e3eb03 100644 --- a/tests/pyathena/test_converter.py +++ b/tests/pyathena/test_converter.py @@ -260,6 +260,8 @@ def test_to_array_json_formats(input_value, expected): [{"name": "John", "age": "30"}, {"name": "Jane", "age": "25"}], ), ("[1, 2.5, hello]", ["1", "2.5", "hello"]), + ("[2023-01-01, 2023-12-31]", ["2023-01-01", "2023-12-31"]), + ("[2023-01-01 12:00:00.123]", ["2023-01-01 12:00:00.123"]), ], ) def test_to_array_athena_native_formats(input_value, expected): diff --git a/tests/pyathena/test_cursor.py b/tests/pyathena/test_cursor.py index d4fbb7cee..6401c07f6 100644 --- a/tests/pyathena/test_cursor.py +++ b/tests/pyathena/test_cursor.py @@ -109,6 +109,22 @@ def _cache_search_cursor(client): return cursor +def _types(value): + """The type of a value and of each value it contains, for comparing them. + + Args: + value: A fetched value, such as a row of lists and dictionaries. + + Returns: + The types, in the shape of the value. + """ + if isinstance(value, dict): + return {k: _types(v) for k, v in value.items()} + if isinstance(value, (list, tuple)): + return [_types(v) for v in value] + return type(value) + + def _block_start(cursor, response=None): """Make the cursor's StartQueryExecution request block until released. @@ -925,8 +941,10 @@ def test_complex_with_type_hints(self, cursor): {"pos": {"x": 3, "y": 4}, "vel": {"x": 1.5}}, ], ) - # Unlike ==, the reprs also tell 123 from 123.0. - assert repr(cursor.fetchone()) == repr(expected) + row = cursor.fetchone() + assert row == expected + # == does not tell 123 from 123.0. + assert _types(row) == _types(expected) def test_cancel(self, cursor): def cancel(c): @@ -1630,9 +1648,18 @@ def test_fetch_complex_values(self, cursor): ARRAY['x', '', 'y'] AS col_empty ,ARRAY['x', ' ', 'a,b'] AS col_space_comma ,ARRAY[json_parse('1234567890'), json_parse('"a,b"')] AS col_json + ,ARRAY[DATE '2023-01-01', DATE '2023-12-31'] AS col_date + ,ARRAY[CAST('2023-01-01 12:00:00.123' AS TIMESTAMP(3))] AS col_timestamp """ cursor.execute(query) - expected = (["x", "", "y"], ["x", " ", "a,b"], [1234567890, "a,b"]) + # Without a hint, the dates and timestamps in an array stay strings. + expected = ( + ["x", "", "y"], + ["x", " ", "a,b"], + [1234567890, "a,b"], + ["2023-01-01", "2023-12-31"], + ["2023-01-01 12:00:00.123"], + ) assert cursor.fetchall() == [expected] cursor.execute( query, @@ -1882,3 +1909,38 @@ def __init__(self, **kwargs): cursor.execute("SELECT * FROM one_row") assert isinstance(cursor.result_set, CustomResultSet) assert type(cursor.fetchone()) is OrderedDict + + def test_null_vs_empty_string(self, dict_cursor): + """ + DictCursor should properly distinguish NULL from empty string. + See docs/null_handling.rst for details. + """ + query = """ + SELECT * FROM ( + VALUES + (1, '', 'empty_string'), + (2, CAST(NULL AS VARCHAR), 'null_value'), + (3, 'hello', 'normal_string'), + (4, 'N/A', 'na_string'), + (5, 'NULL', 'null_string_literal') + ) AS t(id, value, description) + ORDER BY id + """ + dict_cursor.execute(query) + rows = dict_cursor.fetchall() + + # Row 1: Empty string is preserved as empty string, not None + assert rows[0]["value"] == "" + assert rows[0]["value"] is not None + + # Row 2: NULL is properly returned as None + assert rows[1]["value"] is None + + # Row 3: Normal string + assert rows[2]["value"] == "hello" + + # Row 4: "N/A" string should be preserved as-is + assert rows[3]["value"] == "N/A" + + # Row 5: "NULL" string literal should be preserved as-is + assert rows[4]["value"] == "NULL" diff --git a/tests/pyathena/test_glue.py b/tests/pyathena/test_glue.py index ab7d7693b..b9700e024 100644 --- a/tests/pyathena/test_glue.py +++ b/tests/pyathena/test_glue.py @@ -95,8 +95,15 @@ def test_client_leaves_out_athena_endpoint(self): created = [] session_client = conn.session.client + second_build = threading.Event() + def contended_client(*args, **kwargs): created.append(args) + # Without the lock, another thread would build a client meanwhile. + if len(created) == 1: + second_build.wait(1) + else: + second_build.set() return session_client(*args, **kwargs) conn.session.client = contended_client diff --git a/tests/pyathena/test_options.py b/tests/pyathena/test_options.py index 731b3d2c6..bfd3e9624 100644 --- a/tests/pyathena/test_options.py +++ b/tests/pyathena/test_options.py @@ -30,6 +30,12 @@ def test_is_frozen(self): with pytest.raises(dataclasses.FrozenInstanceError): options.work_group = "primary" # type: ignore[misc] + def test_merge_does_not_mutate_original(self): + options = ExecuteOptions(work_group="primary") + merged = options.merge(work_group="adhoc") + assert options.work_group == "primary" + assert merged is not options + def test_merge_ignores_none_overrides(self): options = ExecuteOptions(work_group="primary", cache_size=100) merged = options.merge(work_group=None, cache_size=None, s3_staging_dir=None) diff --git a/tests/pyathena/test_util.py b/tests/pyathena/test_util.py index d932b9a39..1ec39f296 100644 --- a/tests/pyathena/test_util.py +++ b/tests/pyathena/test_util.py @@ -224,8 +224,9 @@ def call(): jitters = [waited - base for waited, base in zip(sleeps, expected_bases, strict=True)] assert all(0 <= jitter <= multiplier for jitter in jitters) if multiplier: - # Random, so the six waits do not add the same jitter. - assert len(set(jitters)) > 1 + # Random, so the six waits do not add the same jitter. Rounded, so that + # the float residuals of a constant jitter do not differ. + assert len({round(jitter, 6) for jitter in jitters}) > 1 @pytest.mark.parametrize( diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 4c1bee3c5..733416c23 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -10,6 +10,7 @@ import threading import time from concurrent.futures import Future, wait +from contextlib import contextmanager from datetime import UTC, datetime, timedelta, timezone from pathlib import Path from unittest.mock import patch @@ -311,6 +312,62 @@ def succeeded_query_execution(query_id, query, completion_date_time, schema="thi ) +# Athena's answers to a query or calculation that fails at its first status request. +FAILED_EXECUTION_RESPONSES = { + "StartQueryExecution": {"QueryExecutionId": "query-id"}, + "GetQueryExecution": { + "QueryExecution": { + "QueryExecutionId": "query-id", + "Query": "SELECT 1", + "StatementType": "DML", + "Status": {"State": "FAILED", "StateChangeReason": "failed"}, + } + }, + "StartCalculationExecution": {"CalculationExecutionId": "calculation-id", "State": "CREATING"}, + "GetCalculationExecutionStatus": { + "Status": {"State": "FAILED"}, + "Statistics": {"DpuExecutionInMillis": 0}, + }, + "GetCalculationExecution": { + "CalculationExecutionId": "calculation-id", + "SessionId": "session", + "Status": {"State": "FAILED", "StateChangeReason": "failed"}, + "Statistics": {"DpuExecutionInMillis": 0}, + "Result": {}, + }, +} + + +@contextmanager +def answered_requests(client, responses=FAILED_EXECUTION_RESPONSES): + """Answer each request of a boto3 client once, without AWS, and record them. + + A repeated request raises ``AssertionError``, so that a cursor that keeps + polling stops instead of waiting for ever, and the block then raises it too, + also when the cursor caught it. + + Args: + client: The boto3 client. + responses: The response to each operation name. + + Yields: + The operation name and parameters of each request, in order. + """ + requests = [] + repeated = [] + + def answer(operation, parameters): + if any(operation == requested for requested, _ in requests): + repeated.append(operation) + raise AssertionError(f"repeated {operation} request") + requests.append((operation, parameters)) + return responses[operation] + + with patch.object(client, "_make_api_call", side_effect=answer): + yield requests + assert not repeated, f"repeated requests: {repeated}" + + class QueryExecutionHistory: """An Athena client stand-in that lists the query executions of a work group. @@ -343,6 +400,12 @@ def list_query_executions(self, **request): The response, with a ``NextToken`` unless the page is the last. """ self.list_requests.append(request) + # ListQueryExecutions rejects a MaxResults above 50. + if request["MaxResults"] > 50: + raise ClientError( + {"Error": {"Code": "InvalidRequestException", "Message": "MaxResults"}}, + "ListQueryExecutions", + ) start = int(request.get("NextToken", 0)) end = start + request["MaxResults"] response = { From 3fa2d7ef48da58ac8561916baa66862182bdb4e0 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 22:18:45 +0900 Subject: [PATCH 5/5] Check the spread of the jitters, and drop the offline temporal-array cases The live complex-value query already checks the native date and timestamp arrays, so the offline parser cases repeated its contract. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/test_converter.py | 2 -- tests/pyathena/test_util.py | 6 +++--- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/tests/pyathena/test_converter.py b/tests/pyathena/test_converter.py index f96e3eb03..5f8507c7f 100644 --- a/tests/pyathena/test_converter.py +++ b/tests/pyathena/test_converter.py @@ -260,8 +260,6 @@ def test_to_array_json_formats(input_value, expected): [{"name": "John", "age": "30"}, {"name": "Jane", "age": "25"}], ), ("[1, 2.5, hello]", ["1", "2.5", "hello"]), - ("[2023-01-01, 2023-12-31]", ["2023-01-01", "2023-12-31"]), - ("[2023-01-01 12:00:00.123]", ["2023-01-01 12:00:00.123"]), ], ) def test_to_array_athena_native_formats(input_value, expected): diff --git a/tests/pyathena/test_util.py b/tests/pyathena/test_util.py index 1ec39f296..6021212ba 100644 --- a/tests/pyathena/test_util.py +++ b/tests/pyathena/test_util.py @@ -224,9 +224,9 @@ def call(): jitters = [waited - base for waited, base in zip(sleeps, expected_bases, strict=True)] assert all(0 <= jitter <= multiplier for jitter in jitters) if multiplier: - # Random, so the six waits do not add the same jitter. Rounded, so that - # the float residuals of a constant jitter do not differ. - assert len({round(jitter, 6) for jitter in jitters}) > 1 + # Random, so the six jitters spread over more than the float residuals of a + # constant one would. + assert max(jitters) - min(jitters) > multiplier / 100 @pytest.mark.parametrize(