Add tests for synchronous and asynchronous record batch stream retrieval

Author: kosiew

python/tests/test_record_batch_stream.py

@@ -0,0 +1,18 @@
+import pytest
+
+
+def test_record_batch_stream_next(ctx):
+    stream = ctx.sql("SELECT 1 as a").execute_stream()
+    batch = next(stream)
+    assert batch.to_pyarrow().num_rows == 1
+    with pytest.raises(StopIteration):
+        next(stream)
+
+
+@pytest.mark.asyncio
+async def test_record_batch_stream_anext(ctx):
+    stream = ctx.sql("SELECT 1 as a").execute_stream()
+    batch = await stream.__anext__()
+    assert batch.to_pyarrow().num_rows == 1
+    with pytest.raises(StopAsyncIteration):
+        await stream.__anext__()

src/record_batch.rs

@@ -61,11 +61,7 @@ impl PyRecordBatchStream {
 #[pymethods]
 impl PyRecordBatchStream {
     fn next(&mut self, py: Python) -> PyResult<PyRecordBatch> {
-        let mut stream = wait_for_future(py, self.stream.lock())?;
-        match wait_for_stream_next(py, &mut stream).map_err(PyDataFusionError::from)? {
-            Some(batch) => Ok(batch.into()),
-            None => Err(PyStopIteration::new_err("stream exhausted")),
-        }
+        wait_for_future(py, next_stream(self.stream.clone(), true))?
     }
 
     fn __next__(&mut self, py: Python) -> PyResult<PyRecordBatch> {