Refactor tests to use fail_collect fixture for DataFrame collect method

Author: kosiew

python/tests/conftest.py

@@ -17,7 +17,7 @@
 
 import pyarrow as pa
 import pytest
-from datafusion import SessionContext
+from datafusion import DataFrame, SessionContext
 from pyarrow.csv import write_csv
 
 
@@ -49,3 +49,12 @@ def database(ctx, tmp_path):
         delimiter=",",
         schema_infer_max_records=10,
     )
+
+
+@pytest.fixture
+def fail_collect(monkeypatch):
+    def _fail_collect(self, *args, **kwargs):  # pragma: no cover - failure path
+        msg = "collect should not be called"
+        raise AssertionError(msg)
+
+    monkeypatch.setattr(DataFrame, "collect", _fail_collect)

python/tests/test_dataframe.py

@@ -1582,21 +1582,14 @@ def test_empty_to_arrow_table(df):
     assert set(pyarrow_table.column_names) == {"a", "b", "c"}
 
 
-def test_arrow_c_stream_to_table(monkeypatch):
+def test_arrow_c_stream_to_table(fail_collect):
     ctx = SessionContext()
 
     # Create a DataFrame with two separate record batches
     batch1 = pa.record_batch([pa.array([1])], names=["a"])
     batch2 = pa.record_batch([pa.array([2])], names=["a"])
     df = ctx.create_dataframe([[batch1], [batch2]])
 
-    # Fail if the DataFrame is pre-collected
-    def fail_collect(self):  # pragma: no cover - failure path
-        msg = "collect should not be called"
-        raise AssertionError(msg)
-
-    monkeypatch.setattr(DataFrame, "collect", fail_collect)
-
     table = pa.Table.from_batches(df)
     expected = pa.Table.from_batches([batch1, batch2])
 

python/tests/test_io.py

@@ -18,7 +18,7 @@
 
 import pyarrow as pa
 import pytest
-from datafusion import DataFrame, column
+from datafusion import column
 from datafusion._testing import range_table
 from datafusion.io import read_avro, read_csv, read_json, read_parquet
 
@@ -123,15 +123,9 @@ def test_arrow_c_stream_large_dataset(ctx):
         assert current_rss - start_rss < 50 * 1024 * 1024
 
 
-def test_table_from_batches_stream(ctx, monkeypatch):
+def test_table_from_batches_stream(ctx, fail_collect):
     df = range_table(ctx, 0, 10)
 
-    def fail_collect(self):  # pragma: no cover - failure path
-        msg = "collect should not be called"
-        raise AssertionError(msg)
-
-    monkeypatch.setattr(DataFrame, "collect", fail_collect)
-
     table = pa.Table.from_batches(df)
     assert table.shape == (10, 1)
     assert table.column_names == ["value"]