diff --git a/tests/test_stream_ops.py b/tests/test_stream_ops.py index 0d827fc..2e34e94 100644 --- a/tests/test_stream_ops.py +++ b/tests/test_stream_ops.py @@ -702,10 +702,11 @@ async def test_read_session_termination(self, stream: S2Stream): ) batches = [] - async for batch in stream.read_session( + async with stream.read_session( start=SeqNum(0), limit=ReadLimit(count=2) - ): - batches.append(batch) + ) as session: + async for batch in session: + batches.append(batch) assert len(batches) >= 1 assert len(batches[0].records) == 2 @@ -717,8 +718,9 @@ async def test_read_session_read_existing_then_tails(self, stream: S2Stream): with pytest.raises(asyncio.TimeoutError): async with asyncio.timeout(2): - async for batch in stream.read_session(start=SeqNum(0)): - received.extend(batch.records) + async with stream.read_session(start=SeqNum(0)) as session: + async for batch in session: + received.extend(batch.records) assert len(received) == 2 assert received[0].body == b"a" @@ -734,8 +736,9 @@ async def append_later(): try: with pytest.raises(asyncio.TimeoutError): async with asyncio.timeout(5): - async for batch in stream.read_session(start=SeqNum(0)): - received.extend(batch.records) + async with stream.read_session(start=SeqNum(0)) as session: + async for batch in session: + received.extend(batch.records) finally: append_later_task.cancel() with suppress(asyncio.CancelledError): @@ -746,24 +749,27 @@ async def append_later(): async def test_read_session_tails(self, stream: S2Stream): with pytest.raises(asyncio.TimeoutError): async with asyncio.timeout(1): - async for _ in stream.read_session(start=SeqNum(0)): - pass + async with stream.read_session(start=SeqNum(0)) as session: + async for _ in session: + pass async def test_read_session_clamp_to_tail_tails(self, stream: S2Stream): await stream.append(AppendInput(records=[Record(body=b"data")])) received = [] with pytest.raises(asyncio.TimeoutError): async with asyncio.timeout(1): - async for batch in stream.read_session( + async with stream.read_session( start=SeqNum(100), clamp_to_tail=True - ): - received.extend(batch.records) + ) as session: + async for batch in session: + received.extend(batch.records) assert received == [] async def test_read_session_beyond_tail_errors(self, stream: S2Stream): with pytest.raises(ReadUnwrittenError) as exc_info: - async for _ in stream.read_session(start=SeqNum(100)): - pass + async with stream.read_session(start=SeqNum(100)) as session: + async for _ in session: + pass assert exc_info.value.tail.seq_num == 0