diff --git a/python/pytests/execution_test.py b/python/pytests/execution_test.py index 27a3fb768..6a41d9f32 100644 --- a/python/pytests/execution_test.py +++ b/python/pytests/execution_test.py @@ -4,6 +4,7 @@ import kaskada as kd import pandas as pd import pytest +import asyncio @pytest.fixture @@ -139,3 +140,33 @@ async def test_history(golden, source_int64) -> None: ) ) ) + +async def test_iter_pandas_async_live_timeout(golden, source_int64) -> None: + data2 = "\n".join( + [ + "time,key,m,n", + "1996-12-20T16:39:57,A,5,10", + "1996-12-20T16:39:58,B,24,3", + "1996-12-20T16:39:59,A,17,6", + "1996-12-20T16:40:00,C,,9", + "1996-12-20T16:40:01,A,12,", + "1996-12-20T16:40:02,A,,", + ] + ) + + execution = source_int64.run_iter(mode="live") + + # Await the first batch. + golden.jsonl(await execution.__anext__()) + + # Add data and await the second batch. + await source_int64.add_string(data2) + golden.jsonl(await execution.__anext__()) + + with pytest.raises(TimeoutError): + async with asyncio.timeout(1): + await execution.__anext__() + + execution.stop() + with pytest.raises(StopAsyncIteration): + print(await execution.__anext__()) \ No newline at end of file diff --git a/python/pytests/golden/execution_test/test_iter_pandas_async_live_timeout.jsonl b/python/pytests/golden/execution_test/test_iter_pandas_async_live_timeout.jsonl new file mode 100644 index 000000000..c1d8ce147 --- /dev/null +++ b/python/pytests/golden/execution_test/test_iter_pandas_async_live_timeout.jsonl @@ -0,0 +1,6 @@ +{"_time":"1996-12-19T16:39:57.000000000","_key":"A","time":"1996-12-19T16:39:57.000000000","key":"A","m":5.0,"n":10.0} +{"_time":"1996-12-19T16:39:58.000000000","_key":"B","time":"1996-12-19T16:39:58.000000000","key":"B","m":24.0,"n":3.0} +{"_time":"1996-12-19T16:39:59.000000000","_key":"A","time":"1996-12-19T16:39:59.000000000","key":"A","m":17.0,"n":6.0} +{"_time":"1996-12-19T16:40:00.000000000","_key":"A","time":"1996-12-19T16:40:00.000000000","key":"A","m":null,"n":9.0} +{"_time":"1996-12-19T16:40:01.000000000","_key":"A","time":"1996-12-19T16:40:01.000000000","key":"A","m":12.0,"n":null} +{"_time":"1996-12-19T16:40:02.000000000","_key":"A","time":"1996-12-19T16:40:02.000000000","key":"A","m":null,"n":null} diff --git a/python/pytests/golden/execution_test/test_iter_pandas_async_live_timeout_1.jsonl b/python/pytests/golden/execution_test/test_iter_pandas_async_live_timeout_1.jsonl new file mode 100644 index 000000000..922338e6f --- /dev/null +++ b/python/pytests/golden/execution_test/test_iter_pandas_async_live_timeout_1.jsonl @@ -0,0 +1,6 @@ +{"_time":"1996-12-20T16:39:57.000000000","_key":"A","time":"1996-12-20T16:39:57.000000000","key":"A","m":5.0,"n":10.0} +{"_time":"1996-12-20T16:39:58.000000000","_key":"B","time":"1996-12-20T16:39:58.000000000","key":"B","m":24.0,"n":3.0} +{"_time":"1996-12-20T16:39:59.000000000","_key":"A","time":"1996-12-20T16:39:59.000000000","key":"A","m":17.0,"n":6.0} +{"_time":"1996-12-20T16:40:00.000000000","_key":"C","time":"1996-12-20T16:40:00.000000000","key":"C","m":null,"n":9.0} +{"_time":"1996-12-20T16:40:01.000000000","_key":"A","time":"1996-12-20T16:40:01.000000000","key":"A","m":12.0,"n":null} +{"_time":"1996-12-20T16:40:02.000000000","_key":"A","time":"1996-12-20T16:40:02.000000000","key":"A","m":null,"n":null}