fix failing test

kylebarron · kylebarron · commit e620b8205709 · 2024-12-17T13:11:29.000-05:00
diff --git a/python/tests/test_dataframe.py b/python/tests/test_dataframe.py
@@ -761,8 +761,8 @@ def test_execution_plan(aggregate_df):
     batch = stream.next()
     assert batch is not None
     # there should be no more batches
-    batch = stream.next()
-    assert batch is None
+    with pytest.raises(StopIteration):
+        stream.next()
 
 
 def test_repartition(df):
