Derive the one_row_complex expectations of the Python-object cursors

tests/pyathena/pandas/test_util.py

@@ -1,7 +1,6 @@
 import textwrap
 import uuid
 from datetime import date, datetime
-from decimal import Decimal
 
 import numpy as np
 import pandas as pd
@@ -16,6 +15,14 @@
     to_sql,
 )
 from tests import ENV
+from tests.pyathena.expected import (
+    ARRAY_JSON,
+    MAP_JSON,
+    PYTHON_VALUES,
+    TIME_OF_TIMESTAMP,
+    ExpectedResult,
+)
+from tests.pyathena.tables import ONE_ROW_COMPLEX
 
 
 def test_get_chunks():
@@ -52,77 +59,11 @@ def test_reset_index():
 
 
 def test_as_pandas(cursor):
-    cursor.execute(
-        """
-        SELECT
-          col_boolean
-          , col_tinyint
-          , col_smallint
-          , col_int
-          , col_bigint
-          , col_float
-          , col_double
-          , col_string
-          , col_timestamp
-          , CAST(col_timestamp AS time) AS col_time
-          , col_date
-          , col_binary
-          , col_array
-          , CAST(col_array AS json) AS col_array_json
-          , col_map
-          , CAST(col_map AS json) AS col_map_json
-          , col_struct
-          , col_decimal
-        FROM one_row_complex
-        """
-    )
+    expected = ExpectedResult(ONE_ROW_COMPLEX, casts=(TIME_OF_TIMESTAMP, ARRAY_JSON, MAP_JSON))
+    cursor.execute(expected.sql)
     df = as_pandas(cursor)
-    rows = [
-        (
-            row["col_boolean"],
-            row["col_tinyint"],
-            row["col_smallint"],
-            row["col_int"],
-            row["col_bigint"],
-            row["col_float"],
-            row["col_double"],
-            row["col_string"],
-            row["col_timestamp"],
-            row["col_time"],
-            row["col_date"],
-            row["col_binary"],
-            row["col_array"],
-            row["col_array_json"],
-            row["col_map"],
-            row["col_map_json"],
-            row["col_struct"],
-            row["col_decimal"],
-        )
-        for _, row in df.iterrows()
-    ]
-    expected = [
-        (
-            True,
-            127,
-            32767,
-            2147483647,
-            9223372036854775807,
-            0.5,
-            0.25,
-            "a string",
-            datetime(2017, 1, 1, 0, 0, 0),
-            datetime(2017, 1, 1, 0, 0, 0).time(),
-            date(2017, 1, 2),
-            b"123",
-            [1, 2],
-            [1, 2],
-            {"1": "2", "3": "4"},
-            {"1": 2, "3": 4},
-            {"a": "1", "b": "2"},
-            Decimal("0.1"),
-        )
-    ]
-    assert rows == expected
+    assert list(df.columns) == expected.names
+    assert [tuple(row) for _, row in df.iterrows()] == expected.rows(PYTHON_VALUES)
 
 
 def test_as_pandas_integer_na_values(cursor):

tests/pyathena/s3fs/test_async_cursor.py

@@ -8,8 +8,6 @@
 import contextlib
 import random
 import time
-from datetime import datetime
-from decimal import Decimal
 
 import pytest
 
@@ -19,6 +17,14 @@
 from pyathena.s3fs.result_set import AthenaS3FSResultSet
 from tests import ENV
 from tests.pyathena.conftest import connect
+from tests.pyathena.expected import (
+    ARRAY_JSON,
+    MAP_JSON,
+    PYTHON_VALUES,
+    TIME_OF_TIMESTAMP,
+    ExpectedResult,
+)
+from tests.pyathena.tables import ONE_ROW_COMPLEX
 
 
 class TestAsyncS3FSCursor:
@@ -61,76 +67,11 @@ def test_invalid_arraysize(self, async_s3fs_cursor):
             async_s3fs_cursor.arraysize = -1
 
     def test_complex(self, async_s3fs_cursor):
-        query_id, future = async_s3fs_cursor.execute(
-            """
-            SELECT
-              col_boolean
-              ,col_tinyint
-              ,col_smallint
-              ,col_int
-              ,col_bigint
-              ,col_float
-              ,col_double
-              ,col_string
-              ,col_varchar
-              ,col_timestamp
-              ,CAST(col_timestamp AS time) AS col_time
-              ,col_date
-              ,col_binary
-              ,col_array
-              ,CAST(col_array AS json) AS col_array_json
-              ,col_map
-              ,CAST(col_map AS json) AS col_map_json
-              ,col_struct
-              ,col_decimal
-            FROM one_row_complex
-            """
-        )
+        expected = ExpectedResult(ONE_ROW_COMPLEX, casts=(TIME_OF_TIMESTAMP, ARRAY_JSON, MAP_JSON))
+        query_id, future = async_s3fs_cursor.execute(expected.sql)
         result_set = future.result()
-        assert result_set.description == [
-            ("col_boolean", "boolean", None, None, 0, 0, "UNKNOWN"),
-            ("col_tinyint", "tinyint", None, None, 3, 0, "UNKNOWN"),
-            ("col_smallint", "smallint", None, None, 5, 0, "UNKNOWN"),
-            ("col_int", "integer", None, None, 10, 0, "UNKNOWN"),
-            ("col_bigint", "bigint", None, None, 19, 0, "UNKNOWN"),
-            ("col_float", "float", None, None, 17, 0, "UNKNOWN"),
-            ("col_double", "double", None, None, 17, 0, "UNKNOWN"),
-            ("col_string", "varchar", None, None, 2147483647, 0, "UNKNOWN"),
-            ("col_varchar", "varchar", None, None, 10, 0, "UNKNOWN"),
-            ("col_timestamp", "timestamp", None, None, 3, 0, "UNKNOWN"),
-            ("col_time", "time", None, None, 3, 0, "UNKNOWN"),
-            ("col_date", "date", None, None, 0, 0, "UNKNOWN"),
-            ("col_binary", "varbinary", None, None, 1073741824, 0, "UNKNOWN"),
-            ("col_array", "array", None, None, 0, 0, "UNKNOWN"),
-            ("col_array_json", "json", None, None, 0, 0, "UNKNOWN"),
-            ("col_map", "map", None, None, 0, 0, "UNKNOWN"),
-            ("col_map_json", "json", None, None, 0, 0, "UNKNOWN"),
-            ("col_struct", "row", None, None, 0, 0, "UNKNOWN"),
-            ("col_decimal", "decimal", None, None, 10, 1, "UNKNOWN"),
-        ]
-        assert result_set.fetchall() == [
-            (
-                True,
-                127,
-                32767,
-                2147483647,
-                9223372036854775807,
-                0.5,
-                0.25,
-                "a string",
-                "varchar",
-                datetime(2017, 1, 1, 0, 0, 0),
-                datetime(2017, 1, 1, 0, 0, 0).time(),
-                datetime(2017, 1, 2).date(),
-                b"123",
-                [1, 2],
-                [1, 2],
-                {"1": "2", "3": "4"},
-                {"1": 2, "3": 4},
-                {"a": "1", "b": "2"},
-                Decimal("0.1"),
-            )
-        ]
+        assert result_set.description == expected.description()
+        assert result_set.fetchall() == expected.rows(PYTHON_VALUES)
 
     def test_cancel(self, async_s3fs_cursor):
         query_id, future = async_s3fs_cursor.execute(

tests/pyathena/s3fs/test_cursor.py

@@ -3,8 +3,6 @@
 import string
 import time
 from concurrent.futures import ThreadPoolExecutor
-from datetime import datetime
-from decimal import Decimal
 
 import pytest
 
@@ -14,6 +12,14 @@
 from pyathena.s3fs.result_set import AthenaS3FSResultSet
 from tests import ENV
 from tests.pyathena.conftest import connect
+from tests.pyathena.expected import (
+    ARRAY_JSON,
+    MAP_JSON,
+    PYTHON_VALUES,
+    TIME_OF_TIMESTAMP,
+    ExpectedResult,
+)
+from tests.pyathena.tables import ONE_ROW_COMPLEX
 
 
 class TestS3FSCursor:
@@ -55,75 +61,10 @@ def test_invalid_arraysize(self, s3fs_cursor):
             s3fs_cursor.arraysize = -1
 
     def test_complex(self, s3fs_cursor):
-        s3fs_cursor.execute(
-            """
-            SELECT
-              col_boolean
-              ,col_tinyint
-              ,col_smallint
-              ,col_int
-              ,col_bigint
-              ,col_float
-              ,col_double
-              ,col_string
-              ,col_varchar
-              ,col_timestamp
-              ,CAST(col_timestamp AS time) AS col_time
-              ,col_date
-              ,col_binary
-              ,col_array
-              ,CAST(col_array AS json) AS col_array_json
-              ,col_map
-              ,CAST(col_map AS json) AS col_map_json
-              ,col_struct
-              ,col_decimal
-            FROM one_row_complex
-            """
-        )
-        assert s3fs_cursor.description == [
-            ("col_boolean", "boolean", None, None, 0, 0, "UNKNOWN"),
-            ("col_tinyint", "tinyint", None, None, 3, 0, "UNKNOWN"),
-            ("col_smallint", "smallint", None, None, 5, 0, "UNKNOWN"),
-            ("col_int", "integer", None, None, 10, 0, "UNKNOWN"),
-            ("col_bigint", "bigint", None, None, 19, 0, "UNKNOWN"),
-            ("col_float", "float", None, None, 17, 0, "UNKNOWN"),
-            ("col_double", "double", None, None, 17, 0, "UNKNOWN"),
-            ("col_string", "varchar", None, None, 2147483647, 0, "UNKNOWN"),
-            ("col_varchar", "varchar", None, None, 10, 0, "UNKNOWN"),
-            ("col_timestamp", "timestamp", None, None, 3, 0, "UNKNOWN"),
-            ("col_time", "time", None, None, 3, 0, "UNKNOWN"),
-            ("col_date", "date", None, None, 0, 0, "UNKNOWN"),
-            ("col_binary", "varbinary", None, None, 1073741824, 0, "UNKNOWN"),
-            ("col_array", "array", None, None, 0, 0, "UNKNOWN"),
-            ("col_array_json", "json", None, None, 0, 0, "UNKNOWN"),
-            ("col_map", "map", None, None, 0, 0, "UNKNOWN"),
-            ("col_map_json", "json", None, None, 0, 0, "UNKNOWN"),
-            ("col_struct", "row", None, None, 0, 0, "UNKNOWN"),
-            ("col_decimal", "decimal", None, None, 10, 1, "UNKNOWN"),
-        ]
-        assert s3fs_cursor.fetchall() == [
-            (
-                True,
-                127,
-                32767,
-                2147483647,
-                9223372036854775807,
-                0.5,
-                0.25,
-                "a string",
-                "varchar",
-                datetime(2017, 1, 1, 0, 0, 0),
-                datetime(2017, 1, 1, 0, 0, 0).time(),
-                datetime(2017, 1, 2).date(),
-                b"123",
-                [1, 2],
-                [1, 2],
-                {"1": "2", "3": "4"},
-                {"1": 2, "3": 4},
-                {"a": "1", "b": "2"},
-                Decimal("0.1"),
-            )
-        ]
+        expected = ExpectedResult(ONE_ROW_COMPLEX, casts=(TIME_OF_TIMESTAMP, ARRAY_JSON, MAP_JSON))
+        s3fs_cursor.execute(expected.sql)
+        assert s3fs_cursor.description == expected.description()
+        assert s3fs_cursor.fetchall() == expected.rows(PYTHON_VALUES)
 
     def test_cancel(self, s3fs_cursor):
         def cancel(c):

tests/pyathena/sqlalchemy/test_base.py

@@ -34,7 +34,9 @@
 )
 from pyathena.util import RetryConfig
 from tests.pyathena.conftest import ENV
-from tests.pyathena.util import decorated, throttle_metadata_api
+from tests.pyathena.expected import PYTHON_VALUES, ExpectedResult
+from tests.pyathena.tables import ONE_ROW_COMPLEX
+from tests.pyathena.util import assert_sqlalchemy_type, decorated, throttle_metadata_api
 
 # Amazon S3 Tables tests need a pre-provisioned table-bucket catalog; the session
 # creates its own namespace in it.
@@ -1394,53 +1396,12 @@ def test_filter_func(self, engine):
     def test_reflect_select(self, engine):
         engine, conn = engine
         one_row_complex = Table("one_row_complex", MetaData(schema=ENV.schema), autoload_with=conn)
-        assert len(one_row_complex.c) == 16
+        assert [c.name for c in one_row_complex.c] == [c.name for c in ONE_ROW_COMPLEX.columns]
         assert isinstance(one_row_complex.c.col_string, Column)
         rows = conn.execute(one_row_complex.select()).fetchall()
-        assert len(rows) == 1
-        assert list(rows[0]) == [
-            True,
-            127,
-            32767,
-            2147483647,
-            9223372036854775807,
-            0.5,
-            0.25,
-            "a string",
-            "varchar",
-            datetime(2017, 1, 1, 0, 0, 0),
-            date(2017, 1, 2),
-            b"123",
-            [1, 2],
-            {"1": "2", "3": "4"},  # map type now converted to dict
-            {"a": "1", "b": "2"},  # row type now converted to dict
-            Decimal("0.1"),
-        ]
-        assert isinstance(one_row_complex.c.col_boolean.type, types.BOOLEAN)
-        assert isinstance(one_row_complex.c.col_tinyint.type, TINYINT)
-        assert isinstance(one_row_complex.c.col_smallint.type, types.SMALLINT)
-        assert isinstance(one_row_complex.c.col_int.type, types.INTEGER)
-        assert isinstance(one_row_complex.c.col_bigint.type, types.BIGINT)
-        assert isinstance(one_row_complex.c.col_float.type, types.FLOAT)
-        assert isinstance(one_row_complex.c.col_double.type, types.DOUBLE)
-        assert isinstance(one_row_complex.c.col_string.type, types.String)
-        assert isinstance(one_row_complex.c.col_varchar.type, types.VARCHAR)
-        assert one_row_complex.c.col_varchar.type.length == 10
-        assert isinstance(one_row_complex.c.col_timestamp.type, types.TIMESTAMP)
-        assert isinstance(one_row_complex.c.col_date.type, types.DATE)
-        assert isinstance(one_row_complex.c.col_binary.type, types.BINARY)
-        assert isinstance(one_row_complex.c.col_array.type, AthenaArray)
-        assert isinstance(one_row_complex.c.col_array.type.item_type, types.INTEGER)
-        assert isinstance(one_row_complex.c.col_map.type, types.String)
-        # With struct support, col_struct should now be recognized as AthenaStruct
-
-        assert isinstance(one_row_complex.c.col_struct.type, AthenaStruct)
-        assert isinstance(
-            one_row_complex.c.col_decimal.type,
-            types.DECIMAL,
-        )
-        assert one_row_complex.c.col_decimal.type.precision == 10
-        assert one_row_complex.c.col_decimal.type.scale == 1
+        assert [tuple(row) for row in rows] == ExpectedResult(ONE_ROW_COMPLEX).rows(PYTHON_VALUES)
+        for column in ONE_ROW_COMPLEX.columns:
+            assert_sqlalchemy_type(one_row_complex.c[column.name].type, column)
 
     def test_select_offset_limit(self, engine):
         engine, conn = engine