Skip to content

Remove and rewrite unnecessary aio and Spark cursor tests #1258

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
laughingman7743 merged 4 commits into from
Oct 18, 2026
Merged

Remove and rewrite unnecessary aio and Spark cursor tests #1258

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
21c24eb
Remove and rewrite unnecessary aio and Spark cursor tests
laughingman7743 Oct 17, 2026
18d208e
Check the options, the Glue thread, and the session that the rewrites…
laughingman7743 Oct 17, 2026
078c134
Keep the checks that the independent review found unprotected
laughingman7743 Oct 17, 2026
a5b5b34
Check the end of async with for AioSparkCursor, and the result set af…
laughingman7743 Oct 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
110 changes: 2 additions & 108 deletions tests/pyathena/aio/arrow/test_cursor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,145 +10,39 @@
import pytest

from pyathena.aio.arrow.cursor import AioArrowCursor
from pyathena.arrow.result_set import AthenaArrowResultSet
from pyathena.error import ProgrammingError
from pyathena.model import AthenaQueryExecution
from pyathena.util import RetryConfig
from tests import ENV
from tests.pyathena.aio.conftest import _aio_connect


class TestAioArrowCursor:
async def test_binary_null_vs_empty(self, aio_arrow_cursor):
query = """SELECT * FROM (VALUES
(1, CAST(NULL AS VARBINARY), 'null', CAST(NULL AS VARCHAR)),
(2, X'', 'empty', ''),
(3, X'00ff275c25', 'comma, quote" and' || chr(10) || 'newline', 'NULL')
) AS t(id, value, label, text_value) ORDER BY id"""
await aio_arrow_cursor.execute(query)
result = aio_arrow_cursor
rows = await result.fetchall()
assert [row[:3] for row in rows] == [
(1, None, "null"),
(2, b"", "empty"),
(3, b"\x00\xff'\\%", 'comma, quote" and\nnewline'),
]
assert [row[3] for row in rows] == ["", "", "NULL"]

async def test_fetchone(self, aio_arrow_cursor):
await aio_arrow_cursor.execute("SELECT * FROM one_row")
assert await aio_arrow_cursor.execute("SELECT * FROM one_row") is aio_arrow_cursor
assert aio_arrow_cursor.rownumber == 0
assert await aio_arrow_cursor.fetchone() == (1,)
assert aio_arrow_cursor.rownumber == 1
assert await aio_arrow_cursor.fetchone() is None

async def test_fetchmany(self, aio_arrow_cursor):
await aio_arrow_cursor.execute("SELECT * FROM many_rows LIMIT 15")
assert len(await aio_arrow_cursor.fetchmany(10)) == 10
assert len(await aio_arrow_cursor.fetchmany(10)) == 5

async def test_fetchall(self, aio_arrow_cursor):
await aio_arrow_cursor.execute("SELECT * FROM one_row")
assert await aio_arrow_cursor.fetchall() == [(1,)]
await aio_arrow_cursor.execute("SELECT a FROM many_rows ORDER BY a")
assert await aio_arrow_cursor.fetchall() == [(i,) for i in range(10000)]

async def test_as_arrow(self, aio_arrow_cursor):
await aio_arrow_cursor.execute("SELECT * FROM one_row")
table = aio_arrow_cursor.as_arrow()
assert table.num_rows == 1
assert table.num_columns == 1
assert table.column_names == ["number_of_rows"]

async def test_as_polars(self, aio_arrow_cursor):
await aio_arrow_cursor.execute("SELECT * FROM one_row")
df = aio_arrow_cursor.as_polars()
assert df.height == 1
assert df.width == 1

async def test_execute_returns_self(self, aio_arrow_cursor):
result = await aio_arrow_cursor.execute("SELECT * FROM one_row")
assert result is aio_arrow_cursor

async def test_no_result_set_raises(self, aio_arrow_cursor):
with pytest.raises(ProgrammingError):
await aio_arrow_cursor.fetchone()
with pytest.raises(ProgrammingError):
await aio_arrow_cursor.fetchmany()
with pytest.raises(ProgrammingError):
await aio_arrow_cursor.fetchall()
with pytest.raises(ProgrammingError):
aio_arrow_cursor.as_arrow()
with pytest.raises(ProgrammingError):
aio_arrow_cursor.as_polars()

async def test_context_manager(self):
from pyathena.aio.arrow.cursor import AioArrowCursor

conn = await _aio_connect(schema_name=ENV.schema, cursor_class=AioArrowCursor)
try:
async with conn.cursor() as cursor:
await cursor.execute("SELECT * FROM one_row")
assert await cursor.fetchone() == (1,)
finally:
conn.close()

async def test_arraysize_default(self, aio_arrow_cursor):
assert aio_arrow_cursor.arraysize == AthenaArrowResultSet.DEFAULT_FETCH_SIZE

async def test_invalid_arraysize(self, aio_arrow_cursor):
aio_arrow_cursor.arraysize = 10000
assert aio_arrow_cursor.arraysize == 10000
with pytest.raises(ProgrammingError):
aio_arrow_cursor.arraysize = -1

async def test_description(self, aio_arrow_cursor):
await aio_arrow_cursor.execute("SELECT CAST(1 AS INT) AS foobar FROM one_row")
assert await aio_arrow_cursor.fetchall() == [(1,)]
assert aio_arrow_cursor.description == [("foobar", "integer", None, None, 10, 0, "UNKNOWN")]

async def test_description_initial(self, aio_arrow_cursor):
assert aio_arrow_cursor.description is None

async def test_cancel_initial(self, aio_arrow_cursor):
with pytest.raises(ProgrammingError):
await aio_arrow_cursor.cancel()

async def test_executemany_fetch(self, aio_arrow_cursor):
await aio_arrow_cursor.executemany(
"SELECT %(x)d FROM one_row", [{"x": i} for i in range(1, 2)]
)
with pytest.raises(ProgrammingError):
await aio_arrow_cursor.fetchall()
with pytest.raises(ProgrammingError):
await aio_arrow_cursor.fetchmany()
with pytest.raises(ProgrammingError):
await aio_arrow_cursor.fetchone()
with pytest.raises(ProgrammingError):
aio_arrow_cursor.as_arrow()
with pytest.raises(ProgrammingError):
aio_arrow_cursor.as_polars()

@pytest.mark.parametrize(
"aio_arrow_cursor",
[{"cursor_kwargs": {"unload": True}}],
indirect=["aio_arrow_cursor"],
)
async def test_fetchone_unload(self, aio_arrow_cursor):
await aio_arrow_cursor.execute("SELECT * FROM one_row")
assert aio_arrow_cursor.as_arrow().num_rows == 1
assert await aio_arrow_cursor.fetchone() == (1,)
assert await aio_arrow_cursor.fetchone() is None

@pytest.mark.parametrize(
"aio_arrow_cursor",
[{"cursor_kwargs": {"unload": True}}],
indirect=["aio_arrow_cursor"],
)
async def test_as_arrow_unload(self, aio_arrow_cursor):
await aio_arrow_cursor.execute("SELECT * FROM one_row")
table = aio_arrow_cursor.as_arrow()
assert table.num_rows == 1

@pytest.mark.parametrize(
"execute_kwargs", [{}, {"connect_timeout": 3.0, "request_timeout": 4.0}]
)
Expand Down
96 changes: 2 additions & 94 deletions tests/pyathena/aio/polars/test_cursor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,130 +10,38 @@
import pytest

from pyathena.aio.polars.cursor import AioPolarsCursor
from pyathena.error import ProgrammingError
from pyathena.model import AthenaQueryExecution
from pyathena.polars.result_set import AthenaPolarsResultSet
from pyathena.util import RetryConfig
from tests import ENV
from tests.pyathena.aio.conftest import _aio_connect


class TestAioPolarsCursor:
async def test_fetchone(self, aio_polars_cursor):
await aio_polars_cursor.execute("SELECT * FROM one_row")
assert await aio_polars_cursor.execute("SELECT * FROM one_row") is aio_polars_cursor
assert aio_polars_cursor.rownumber == 0
assert await aio_polars_cursor.fetchone() == (1,)
assert aio_polars_cursor.rownumber == 1
assert await aio_polars_cursor.fetchone() is None

async def test_fetchmany(self, aio_polars_cursor):
await aio_polars_cursor.execute("SELECT * FROM many_rows LIMIT 15")
assert len(await aio_polars_cursor.fetchmany(10)) == 10
assert len(await aio_polars_cursor.fetchmany(10)) == 5

async def test_fetchall(self, aio_polars_cursor):
await aio_polars_cursor.execute("SELECT * FROM one_row")
assert await aio_polars_cursor.fetchall() == [(1,)]
await aio_polars_cursor.execute("SELECT a FROM many_rows ORDER BY a")
assert await aio_polars_cursor.fetchall() == [(i,) for i in range(10000)]

async def test_as_polars(self, aio_polars_cursor):
await aio_polars_cursor.execute("SELECT * FROM one_row")
df = aio_polars_cursor.as_polars()
assert df.height == 1
assert df.width == 1

async def test_as_arrow(self, aio_polars_cursor):
await aio_polars_cursor.execute("SELECT * FROM one_row")
table = aio_polars_cursor.as_arrow()
assert table.num_rows == 1
assert table.num_columns == 1

async def test_execute_returns_self(self, aio_polars_cursor):
result = await aio_polars_cursor.execute("SELECT * FROM one_row")
assert result is aio_polars_cursor

async def test_no_result_set_raises(self, aio_polars_cursor):
with pytest.raises(ProgrammingError):
await aio_polars_cursor.fetchone()
with pytest.raises(ProgrammingError):
await aio_polars_cursor.fetchmany()
with pytest.raises(ProgrammingError):
await aio_polars_cursor.fetchall()
with pytest.raises(ProgrammingError):
aio_polars_cursor.as_polars()
with pytest.raises(ProgrammingError):
aio_polars_cursor.as_arrow()

async def test_context_manager(self):
from pyathena.aio.polars.cursor import AioPolarsCursor

conn = await _aio_connect(schema_name=ENV.schema, cursor_class=AioPolarsCursor)
try:
async with conn.cursor() as cursor:
await cursor.execute("SELECT * FROM one_row")
assert await cursor.fetchone() == (1,)
finally:
conn.close()

async def test_arraysize_default(self, aio_polars_cursor):
assert aio_polars_cursor.arraysize == AthenaPolarsResultSet.DEFAULT_FETCH_SIZE

async def test_invalid_arraysize(self, aio_polars_cursor):
aio_polars_cursor.arraysize = 10000
assert aio_polars_cursor.arraysize == 10000
with pytest.raises(ProgrammingError):
aio_polars_cursor.arraysize = -1

async def test_description(self, aio_polars_cursor):
await aio_polars_cursor.execute("SELECT CAST(1 AS INT) AS foobar FROM one_row")
assert await aio_polars_cursor.fetchall() == [(1,)]
assert aio_polars_cursor.description == [
("foobar", "integer", None, None, 10, 0, "UNKNOWN")
]

async def test_description_initial(self, aio_polars_cursor):
assert aio_polars_cursor.description is None

async def test_cancel_initial(self, aio_polars_cursor):
with pytest.raises(ProgrammingError):
await aio_polars_cursor.cancel()

async def test_executemany_fetch(self, aio_polars_cursor):
await aio_polars_cursor.executemany(
"SELECT %(x)d FROM one_row", [{"x": i} for i in range(1, 2)]
)
with pytest.raises(ProgrammingError):
await aio_polars_cursor.fetchall()
with pytest.raises(ProgrammingError):
await aio_polars_cursor.fetchmany()
with pytest.raises(ProgrammingError):
await aio_polars_cursor.fetchone()
with pytest.raises(ProgrammingError):
aio_polars_cursor.as_polars()
with pytest.raises(ProgrammingError):
aio_polars_cursor.as_arrow()

@pytest.mark.parametrize(
"aio_polars_cursor",
[{"cursor_kwargs": {"unload": True}}],
indirect=["aio_polars_cursor"],
)
async def test_fetchone_unload(self, aio_polars_cursor):
await aio_polars_cursor.execute("SELECT * FROM one_row")
assert aio_polars_cursor.as_polars().height == 1
assert await aio_polars_cursor.fetchone() == (1,)
assert await aio_polars_cursor.fetchone() is None

@pytest.mark.parametrize(
"aio_polars_cursor",
[{"cursor_kwargs": {"unload": True}}],
indirect=["aio_polars_cursor"],
)
async def test_as_polars_unload(self, aio_polars_cursor):
await aio_polars_cursor.execute("SELECT * FROM one_row")
df = aio_polars_cursor.as_polars()
assert df.height == 1

@pytest.mark.parametrize(
"execute_kwargs",
[{}, {"block_size": 2048, "cache_type": "none", "s3_max_workers": 3, "chunksize": 20}],
Expand Down
81 changes: 1 addition & 80 deletions tests/pyathena/aio/s3fs/test_cursor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,98 +10,19 @@
import pytest

from pyathena.aio.s3fs.cursor import AioS3FSCursor
from pyathena.error import ProgrammingError
from pyathena.model import AthenaQueryExecution
from pyathena.s3fs.reader import AthenaCSVReader, EmptyStringAsNullCSVReader
from pyathena.s3fs.result_set import AthenaS3FSResultSet
from pyathena.util import RetryConfig
from tests import ENV
from tests.pyathena.aio.conftest import _aio_connect


class TestAioS3FSCursor:
@pytest.mark.parametrize(
("aio_s3fs_cursor", "expected_empty"),
[
({}, ""),
({"cursor_kwargs": {"csv_reader": EmptyStringAsNullCSVReader}}, None),
],
indirect=["aio_s3fs_cursor"],
)
async def test_null_vs_empty_string(self, aio_s3fs_cursor, expected_empty):
await aio_s3fs_cursor.execute(
"SELECT NULL AS null_col, '' AS empty_col, 'text' AS value_col"
)
assert await aio_s3fs_cursor.fetchall() == [(None, expected_empty, "text")]

async def test_fetchone(self, aio_s3fs_cursor):
await aio_s3fs_cursor.execute("SELECT * FROM one_row")
assert await aio_s3fs_cursor.execute("SELECT * FROM one_row") is aio_s3fs_cursor
assert aio_s3fs_cursor.rownumber == 0
assert await aio_s3fs_cursor.fetchone() == (1,)
assert aio_s3fs_cursor.rownumber == 1
assert await aio_s3fs_cursor.fetchone() is None

async def test_fetchmany(self, aio_s3fs_cursor):
await aio_s3fs_cursor.execute("SELECT * FROM many_rows LIMIT 15")
assert len(await aio_s3fs_cursor.fetchmany(10)) == 10
assert len(await aio_s3fs_cursor.fetchmany(10)) == 5

async def test_fetchall(self, aio_s3fs_cursor):
await aio_s3fs_cursor.execute("SELECT * FROM one_row")
assert await aio_s3fs_cursor.fetchall() == [(1,)]
await aio_s3fs_cursor.execute("SELECT a FROM many_rows ORDER BY a")
assert await aio_s3fs_cursor.fetchall() == [(i,) for i in range(10000)]

async def test_description(self, aio_s3fs_cursor):
await aio_s3fs_cursor.execute("SELECT 1 AS foobar FROM one_row")
assert await aio_s3fs_cursor.fetchall() == [(1,)]
assert aio_s3fs_cursor.description == [("foobar", "integer", None, None, 10, 0, "UNKNOWN")]

async def test_description_initial(self, aio_s3fs_cursor):
assert aio_s3fs_cursor.description is None

async def test_cancel_initial(self, aio_s3fs_cursor):
with pytest.raises(ProgrammingError):
await aio_s3fs_cursor.cancel()

async def test_executemany_fetch(self, aio_s3fs_cursor):
await aio_s3fs_cursor.executemany(
"SELECT %(x)d FROM one_row", [{"x": i} for i in range(1, 2)]
)
with pytest.raises(ProgrammingError):
await aio_s3fs_cursor.fetchall()
with pytest.raises(ProgrammingError):
await aio_s3fs_cursor.fetchmany()
with pytest.raises(ProgrammingError):
await aio_s3fs_cursor.fetchone()

async def test_arraysize_default(self, aio_s3fs_cursor):
assert aio_s3fs_cursor.arraysize == AthenaS3FSResultSet.DEFAULT_FETCH_SIZE

async def test_invalid_arraysize(self, aio_s3fs_cursor):
aio_s3fs_cursor.arraysize = 10000
assert aio_s3fs_cursor.arraysize == 10000
with pytest.raises(ProgrammingError):
aio_s3fs_cursor.arraysize = -1

async def test_async_iterator(self, aio_s3fs_cursor):
await aio_s3fs_cursor.execute("SELECT * FROM one_row")
rows = [row async for row in aio_s3fs_cursor]
assert rows == [(1,)]

async def test_context_manager(self):
conn = await _aio_connect(schema_name=ENV.schema)
try:
async with conn.cursor(AioS3FSCursor) as cursor:
await cursor.execute("SELECT * FROM one_row")
assert await cursor.fetchone() == (1,)
finally:
conn.close()

async def test_execute_returns_self(self, aio_s3fs_cursor):
result = await aio_s3fs_cursor.execute("SELECT * FROM one_row")
assert result is aio_s3fs_cursor

@pytest.mark.parametrize("execute_kwargs", [{}, {"csv_reader": AthenaCSVReader}])
async def test_read_options(self, execute_kwargs):
"""The cursor's read options reach the result set, and execute() overrides them.
Expand Down
Loading
Loading