diff --git a/src/calendar_queue/calendar.py b/src/calendar_queue/calendar.py index 50e2ad5..5a5dd1d 100644 --- a/src/calendar_queue/calendar.py +++ b/src/calendar_queue/calendar.py @@ -148,10 +148,26 @@ async def __anext__(self) -> tuple[float, CalendarEvent]: return event def stop(self) -> None: - """Stop the execution of the calendar""" + """Stop the execution of the calendar. + + Once called, any active ``async for`` loop over this instance will + raise ``StopAsyncIteration`` at the next iteration. The effect is + permanent: further iteration will also stop immediately until + ``reset()`` is called. + """ self._stop_event.set() + def reset(self) -> None: + """Re-enable iteration after ``stop()`` has been called. + + Clears the internal stop flag so that the calendar can be iterated + again with ``async for``. Scheduled events that were preserved when + ``stop()`` was called (see ``__anext__``) will be emitted normally. + """ + + self._stop_event.clear() + def clear(self) -> list[tuple[float, CalendarEvent]]: """Cancel all scheduled events. diff --git a/src/calendar_queue/calendar_queue.py b/src/calendar_queue/calendar_queue.py index c1930a9..22141e9 100644 --- a/src/calendar_queue/calendar_queue.py +++ b/src/calendar_queue/calendar_queue.py @@ -56,7 +56,7 @@ def _update_timer(self): return prev_timer_ts = self._getter_timer.when() - if self.qsize() > 1 and prev_timer_ts > loop.time(): + if prev_timer_ts > loop.time(): should_update = not math.isclose( prev_timer_ts, loop.time() + ts - time(), abs_tol=0.00001 ) @@ -258,6 +258,12 @@ def delete_items( # Restore heap invariant after arbitrary deletions. if del_items: heapify(self._queue) + # Mirror the _unfinished_tasks increments done by put_nowait so + # that join() does not block indefinitely for deleted items. + n = min(len(del_items), self._unfinished_tasks) + self._unfinished_tasks -= n + if self._unfinished_tasks == 0: + self._finished.set() self._update_timer() diff --git a/tests/__init__.py b/tests/__init__.py index 6c12b18..96484e8 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -1 +1 @@ -ABS_TOLERANCE = 0.0001 +ABS_TOLERANCE = 0.001 diff --git a/tests/test_calendar.py b/tests/test_calendar.py index 2133d8f..cb9350d 100644 --- a/tests/test_calendar.py +++ b/tests/test_calendar.py @@ -77,6 +77,40 @@ async def test_events_generator(): c.stop() +@pytest.mark.asyncio +async def test_reset(): + """Test that reset() re-enables iteration after stop() has been called.""" + + c: Calendar[int] = Calendar() + + ts = datetime.now() + + for i in range(4): + c.schedule(i, ts + timedelta(seconds=i * 1)) + + # Consume the first two events then stop. + collected = [] + async for _, event in c: + collected.append(event) + if event == 1: + c.stop() + + assert collected == [0, 1] + + # Without reset, iteration stops immediately. + async for _, event in c: + pytest.fail(f"Should not have received event {event} after stop()") + + # After reset, the remaining events are emitted normally. + c.reset() + async for _, event in c: + collected.append(event) + if len(c.events()) == 0: + c.stop() + + assert collected == [0, 1, 2, 3] + + @pytest.mark.asyncio async def test_cancel_events(): """Test that cancelling events remove them from the calendar diff --git a/tests/test_calendar_queue.py b/tests/test_calendar_queue.py index f461210..66a6d45 100644 --- a/tests/test_calendar_queue.py +++ b/tests/test_calendar_queue.py @@ -2,8 +2,6 @@ import sys from datetime import timedelta from time import time -from unittest.mock import patch - import pytest from calendar_queue import CalendarQueue @@ -149,15 +147,12 @@ async def test_delete_restores_heap_property(): @pytest.mark.asyncio async def test_far_schedule(): """Test putting an event scheduled far in time. - We use LoopTimeTravel to simulate the time traveling - and verify that the elements are returned at the correct - timestamp. + Verifies that the item blocks while not yet due, and is returned + once its time arrives. """ cq = CalendarQueue() - cq_loop = cq._get_loop() # type: ignore - delta = timedelta(days=30) cq.put_nowait((time() + delta.total_seconds(), "foo")) @@ -166,9 +161,18 @@ async def test_far_schedule(): await asyncio.wait_for(cq.get(), 2) pytest.fail("Item was returned immediately. It makes no sense.") - with patch.object(cq_loop, "time", return_value=cq_loop.time() + delta.total_seconds() + 10): - item = await asyncio.wait_for(cq.get(), 5) - assert item[-1] == "foo" + # The item is still in the queue (the timed-out get() was cancelled + # before consuming it). Replace the far-future timer with one that + # fires immediately to simulate the scheduled time arriving, without + # globally patching loop.time() which breaks asyncio's own internals. + assert cq._getter_timer is not None # type: ignore + cq._getter_timer.cancel() # type: ignore + cq._getter_timer = cq._get_loop().call_later( # type: ignore + 0, cq._calendar_alarm + ) + + item = await asyncio.wait_for(cq.get(), 5) + assert item[-1] == "foo" @pytest.mark.asyncio