Add chunksize tests for PolarsCursor fetch methods

- Add tests for fetchone, fetchmany, fetchall, and iterator with chunksize
- Add tests for fetch methods with UNLOAD mode and chunksize
- Remove redundant iter_chunks tests from AsyncPolarsCursor since both
  cursor types share the same AthenaPolarsResultSet implementation

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: laughingman7743

tests/pyathena/polars/test_async_cursor.py

@@ -324,111 +324,3 @@ def test_empty_result_unload(self, async_polars_cursor):
         df = future.result().as_polars()
         assert df.height == 0
         assert df.width == 0
-
-    def test_iter_chunks(self):
-        """Test chunked iteration over query results."""
-        with contextlib.closing(connect(schema_name=ENV.schema)) as conn:
-            cursor = conn.cursor(AsyncPolarsCursor, chunksize=5)
-            query_id, future = cursor.execute("SELECT * FROM many_rows LIMIT 15")
-            assert query_id is not None
-            result_set = future.result()
-            chunks = list(result_set.iter_chunks())
-            assert len(chunks) > 0
-            total_rows = sum(chunk.height for chunk in chunks)
-            assert total_rows == 15
-            for chunk in chunks:
-                assert isinstance(chunk, pl.DataFrame)
-
-    def test_iter_chunks_without_chunksize(self, async_polars_cursor):
-        """Test that iter_chunks works without chunksize, yielding entire DataFrame."""
-        query_id, future = async_polars_cursor.execute("SELECT * FROM one_row")
-        assert query_id is not None
-        result_set = future.result()
-        chunks = list(result_set.iter_chunks())
-        # Without chunksize, yields entire DataFrame as single chunk
-        assert len(chunks) == 1
-        assert isinstance(chunks[0], pl.DataFrame)
-        assert chunks[0].height == 1
-
-    def test_iter_chunks_many_rows(self):
-        """Test chunked iteration with many rows."""
-        with contextlib.closing(connect(schema_name=ENV.schema)) as conn:
-            cursor = conn.cursor(AsyncPolarsCursor, chunksize=1000)
-            query_id, future = cursor.execute("SELECT * FROM many_rows")
-            assert query_id is not None
-            result_set = future.result()
-            chunks = list(result_set.iter_chunks())
-            total_rows = sum(chunk.height for chunk in chunks)
-            assert total_rows == 10000
-            assert len(chunks) >= 10  # At least 10 chunks with chunksize=1000
-
-    @pytest.mark.parametrize(
-        "async_polars_cursor",
-        [
-            {
-                "cursor_kwargs": {"unload": True, "chunksize": 5},
-            },
-        ],
-        indirect=["async_polars_cursor"],
-    )
-    def test_iter_chunks_unload(self, async_polars_cursor):
-        """Test chunked iteration with UNLOAD (Parquet)."""
-        query_id, future = async_polars_cursor.execute("SELECT * FROM many_rows LIMIT 15")
-        assert query_id is not None
-        result_set = future.result()
-        chunks = list(result_set.iter_chunks())
-        assert len(chunks) > 0
-        total_rows = sum(chunk.height for chunk in chunks)
-        assert total_rows == 15
-        for chunk in chunks:
-            assert isinstance(chunk, pl.DataFrame)
-
-    def test_iter_chunks_data_consistency(self):
-        """Test that chunked and regular reading produce the same data."""
-        with contextlib.closing(connect(schema_name=ENV.schema)) as conn:
-            # Regular reading (no chunksize)
-            regular_cursor = conn.cursor(AsyncPolarsCursor)
-            query_id, future = regular_cursor.execute("SELECT * FROM many_rows LIMIT 100")
-            assert query_id is not None
-            regular_df = future.result().as_polars()
-
-            # Chunked reading
-            chunked_cursor = conn.cursor(AsyncPolarsCursor, chunksize=25)
-            query_id, future = chunked_cursor.execute("SELECT * FROM many_rows LIMIT 100")
-            assert query_id is not None
-            result_set = future.result()
-            chunked_dfs = list(result_set.iter_chunks())
-
-            # Combine chunks
-            combined_df = pl.concat(chunked_dfs)
-
-            # Should have the same data (sort for comparison)
-            assert regular_df.sort("a").equals(combined_df.sort("a"))
-
-            # Should have multiple chunks
-            assert len(chunked_dfs) > 1
-
-    def test_iter_chunks_chunk_sizes(self):
-        """Test that chunks have correct sizes."""
-        with contextlib.closing(connect(schema_name=ENV.schema)) as conn:
-            cursor = conn.cursor(AsyncPolarsCursor, chunksize=10)
-            query_id, future = cursor.execute("SELECT * FROM many_rows LIMIT 50")
-            assert query_id is not None
-            result_set = future.result()
-
-            chunk_sizes = []
-            total_rows = 0
-
-            for chunk in result_set.iter_chunks():
-                chunk_size = chunk.height
-                chunk_sizes.append(chunk_size)
-                total_rows += chunk_size
-
-                # Each chunk should not exceed chunksize
-                assert chunk_size <= 10
-
-            # Should have processed all 50 rows
-            assert total_rows == 50
-
-            # Should have multiple chunks
-            assert len(chunk_sizes) > 1

tests/pyathena/polars/test_cursor.py

@@ -542,3 +542,85 @@ def test_iter_chunks_chunk_sizes(self):
 
             # Should have multiple chunks
             assert len(chunk_sizes) > 1
+
+    def test_fetchone_with_chunksize(self):
+        """Test that fetchone works correctly with chunksize enabled."""
+        with contextlib.closing(connect(schema_name=ENV.schema)) as conn:
+            cursor = conn.cursor(PolarsCursor, chunksize=5)
+            cursor.execute("SELECT * FROM many_rows LIMIT 15")
+
+            rows = []
+            while True:
+                row = cursor.fetchone()
+                if row is None:
+                    break
+                rows.append(row)
+
+            assert len(rows) == 15
+
+    def test_fetchmany_with_chunksize(self):
+        """Test that fetchmany works correctly with chunksize enabled."""
+        with contextlib.closing(connect(schema_name=ENV.schema)) as conn:
+            cursor = conn.cursor(PolarsCursor, chunksize=5)
+            cursor.execute("SELECT * FROM many_rows LIMIT 15")
+
+            batch1 = cursor.fetchmany(10)
+            batch2 = cursor.fetchmany(10)
+
+            assert len(batch1) == 10
+            assert len(batch2) == 5
+
+    def test_fetchall_with_chunksize(self):
+        """Test that fetchall works correctly with chunksize enabled."""
+        with contextlib.closing(connect(schema_name=ENV.schema)) as conn:
+            cursor = conn.cursor(PolarsCursor, chunksize=5)
+            cursor.execute("SELECT * FROM many_rows LIMIT 15")
+
+            rows = cursor.fetchall()
+            assert len(rows) == 15
+
+    def test_iterator_with_chunksize(self):
+        """Test that cursor iteration works correctly with chunksize enabled."""
+        with contextlib.closing(connect(schema_name=ENV.schema)) as conn:
+            cursor = conn.cursor(PolarsCursor, chunksize=5)
+            cursor.execute("SELECT * FROM many_rows LIMIT 15")
+
+            rows = list(cursor)
+            assert len(rows) == 15
+
+    @pytest.mark.parametrize(
+        "polars_cursor",
+        [
+            {
+                "cursor_kwargs": {"unload": True, "chunksize": 5},
+            },
+        ],
+        indirect=["polars_cursor"],
+    )
+    def test_fetchone_with_chunksize_unload(self, polars_cursor):
+        """Test that fetchone works correctly with chunksize and unload enabled."""
+        polars_cursor.execute("SELECT * FROM many_rows LIMIT 15")
+
+        rows = []
+        while True:
+            row = polars_cursor.fetchone()
+            if row is None:
+                break
+            rows.append(row)
+
+        assert len(rows) == 15
+
+    @pytest.mark.parametrize(
+        "polars_cursor",
+        [
+            {
+                "cursor_kwargs": {"unload": True, "chunksize": 5},
+            },
+        ],
+        indirect=["polars_cursor"],
+    )
+    def test_iterator_with_chunksize_unload(self, polars_cursor):
+        """Test that cursor iteration works with chunksize and unload enabled."""
+        polars_cursor.execute("SELECT * FROM many_rows LIMIT 15")
+        rows = list(polars_cursor)
+        assert len(rows) == 15