Add tests for all Receiver.close() implementations

Signed-off-by: Sahas Subramanian <[email protected]>

Author: shsms

tests/test_anycast.py

@@ -217,3 +217,30 @@ async def test_anycast_filter() -> None:
 
     assert (await receiver.receive()) == 12
     assert (await receiver.receive()) == 15
+
+
+async def test_anycast_close_receiver() -> None:
+    """Ensure closing a receiver stops the receiver."""
+    chan = Anycast[int](name="input-chan")
+    sender = chan.new_sender()
+
+    receiver_1 = chan.new_receiver()
+    receiver_2 = chan.new_receiver()
+
+    await sender.send(1)
+
+    assert (await receiver_1.receive()) == 1
+
+    receiver_1.close()
+
+    await sender.send(2)
+
+    with pytest.raises(ReceiverStoppedError):
+        _ = await receiver_1.receive()
+
+    assert (await receiver_2.receive()) == 2
+
+    receiver_2.close()
+
+    with pytest.raises(ReceiverStoppedError):
+        _ = await receiver_2.receive()

tests/test_broadcast.py

@@ -232,6 +232,42 @@ async def test_broadcast_map() -> None:
     assert (await receiver.receive()) is True
 
 
+async def test_broadcast_map_close_receiver() -> None:
+    """Ensure closing a map stops the receiver."""
+    chan = Broadcast[int](name="input-chan")
+    sender = chan.new_sender()
+
+    receiver_1 = chan.new_receiver()
+    receiver_2 = chan.new_receiver()
+    plus_100_rx = receiver_1.map(lambda num: num + 100)
+
+    await sender.send(1)
+
+    assert (await plus_100_rx.receive()) == 101
+    assert (await receiver_2.receive()) == 1
+
+    plus_100_rx.close()
+
+    await sender.send(2)
+
+    with pytest.raises(ReceiverStoppedError):
+        _ = await plus_100_rx.receive()
+
+    with pytest.raises(ReceiverStoppedError):
+        _ = await receiver_1.receive()
+
+    assert (await receiver_2.receive()) == 2
+
+    await sender.send(3)
+
+    assert (await receiver_2.receive()) == 3
+
+    receiver_2.close()
+
+    with pytest.raises(ReceiverStoppedError):
+        _ = await receiver_2.receive()
+
+
 async def test_broadcast_filter() -> None:
     """Ensure filter keeps only the messages that pass the filter."""
     chan = Broadcast[int](name="input-chan")
@@ -249,6 +285,43 @@ async def test_broadcast_filter() -> None:
     assert (await receiver.receive()) == 15
 
 
+async def test_broadcast_filter_close_receiver() -> None:
+    """Ensure closing a filter stops the receiver."""
+    chan = Broadcast[int](name="input-chan")
+    sender = chan.new_sender()
+
+    receiver_1 = chan.new_receiver()
+    receiver_2 = chan.new_receiver()
+
+    gt_10_rx = receiver_1.filter(lambda num: num > 10)
+
+    await sender.send(1)
+    assert (await receiver_2.receive()) == 1
+
+    await sender.send(100)
+    assert (await gt_10_rx.receive()) == 100
+    assert (await receiver_2.receive()) == 100
+
+    gt_10_rx.close()
+
+    await sender.send(2)
+
+    with pytest.raises(ReceiverStoppedError):
+        _ = await gt_10_rx.receive()
+    with pytest.raises(ReceiverStoppedError):
+        _ = await receiver_1.receive()
+
+    assert (await receiver_2.receive()) == 2
+
+    await sender.send(3)
+    assert (await receiver_2.receive()) == 3
+
+    receiver_2.close()
+
+    with pytest.raises(ReceiverStoppedError):
+        _ = await receiver_2.receive()
+
+
 async def test_broadcast_filter_type_guard() -> None:
     """Ensure filter type guard works."""
     chan = Broadcast[int | str](name="input-chan")
@@ -320,3 +393,35 @@ class Narrower(Actual):
 
     await sender.send(Narrower(10))
     assert (await receiver.receive()).value == 10
+
+
+async def test_broadcast_close_receiver() -> None:
+    """Ensure closing a receiver stops the receiver."""
+    chan = Broadcast[int](name="input-chan")
+    sender = chan.new_sender()
+
+    receiver_1 = chan.new_receiver()
+    receiver_2 = chan.new_receiver()
+
+    await sender.send(1)
+
+    assert (await receiver_1.receive()) == 1
+    assert (await receiver_2.receive()) == 1
+
+    receiver_1.close()
+
+    await sender.send(2)
+
+    with pytest.raises(ReceiverStoppedError):
+        _ = await receiver_1.receive()
+
+    assert (await receiver_2.receive()) == 2
+
+    await sender.send(3)
+
+    assert (await receiver_2.receive()) == 3
+
+    receiver_2.close()
+
+    with pytest.raises(ReceiverStoppedError):
+        _ = await receiver_2.receive()

tests/test_event.py

@@ -57,3 +57,51 @@ async def wait_for_event() -> None:
     assert not event.is_set
 
     await event_task
+
+
+async def test_event_close_receiver() -> None:
+    """Ensure that closing an event stops the receiver."""
+    event = Event()
+    assert not event.is_set
+    assert not event.is_stopped
+
+    is_ready = False
+
+    async def wait_for_event() -> None:
+        nonlocal is_ready
+        await event.ready()
+        is_ready = True
+
+    event_task = _asyncio.create_task(wait_for_event())
+
+    await _asyncio.sleep(0)  # Yield so the wait_for_event task can run.
+
+    assert not is_ready
+    assert not event.is_set
+    assert not event.is_stopped
+
+    event.set()
+
+    await _asyncio.sleep(0)  # Yield so the wait_for_event task can run.
+    assert is_ready
+    assert event.is_set
+    assert not event.is_stopped
+
+    event.consume()
+    assert not event.is_set
+    assert not event.is_stopped
+    assert event_task.done()
+    assert event_task.result() is None
+    assert not event_task.cancelled()
+
+    event.close()
+    assert not event.is_set
+    assert event.is_stopped
+
+    await event.ready()
+    with _pytest.raises(ReceiverStoppedError):
+        event.consume()
+    assert event.is_stopped
+    assert not event.is_set
+
+    await event_task

tests/test_file_watcher_integration.py

@@ -150,3 +150,40 @@ async def test_file_watcher_exit_iterator(tmp_path: pathlib.Path) -> None:
         file_watcher.consume()
 
     assert number_of_writes == expected_number_of_writes
+
+
+@pytest.mark.integration
+async def test_file_watcher_close_receiver(tmp_path: pathlib.Path) -> None:
+    """Ensure closing the file watcher stops the receiver.
+
+    Args:
+        tmp_path: A tmp directory to run the file watcher on. Created by pytest.
+    """
+    filename = tmp_path / "test-file"
+
+    number_of_writes = 0
+    expected_number_of_writes = 3
+
+    file_watcher = FileWatcher(
+        paths=[str(tmp_path)],
+        force_polling=True,
+        polling_interval=timedelta(seconds=0.05),
+    )
+    timer = Timer(timedelta(seconds=0.1), SkipMissedAndDrift())
+
+    async for selected in select(file_watcher, timer):
+        if selected_from(selected, timer):
+            filename.write_text(f"{selected.message}")
+        elif selected_from(selected, file_watcher):
+            number_of_writes += 1
+            if number_of_writes == expected_number_of_writes:
+                file_watcher.close()
+                break
+
+    ready = await file_watcher.ready()
+    assert ready is False
+
+    with pytest.raises(ReceiverStoppedError):
+        file_watcher.consume()
+
+    assert number_of_writes == expected_number_of_writes

tests/test_merge_integration.py

@@ -8,6 +8,8 @@
 import pytest
 
 from frequenz.channels import Anycast, Sender, merge
+from frequenz.channels._broadcast import Broadcast
+from frequenz.channels._receiver import ReceiverStoppedError
 
 
 @pytest.mark.integration
@@ -39,3 +41,44 @@ async def send(ch1: Sender[int], ch2: Sender[int]) -> None:
         # succession.
         assert set(results[idx : idx + 2]) == {ctr + 1, ctr + 101}
     assert results[-1] == 1000
+
+
+async def test_merge_close_receiver() -> None:
+    """Ensure merge() closes when a receiver is closed."""
+    chan1 = Broadcast[int](name="chan1")
+    chan2 = Broadcast[int](name="chan2")
+
+    async def send(ch1: Sender[int], ch2: Sender[int]) -> None:
+        for ctr in range(5):
+            await ch1.send(ctr + 1)
+            await ch2.send(ctr + 101)
+        await chan1.close()
+        await chan2.close()
+
+    rx1 = chan1.new_receiver()
+    rx2 = chan2.new_receiver()
+    closing_merge = merge(rx1, rx2)
+    prx1 = chan1.new_receiver()
+    prx2 = chan2.new_receiver()
+    completing_merge = merge(prx1, prx2)
+
+    senders = asyncio.create_task(send(chan1.new_sender(), chan2.new_sender()))
+
+    results: list[int] = []
+    async for item in closing_merge:
+        results.append(item)
+        if item == 3:
+            closing_merge.close()
+    await senders
+    assert set(results) == {1, 101, 2, 102, 3, 103}
+
+    with pytest.raises(ReceiverStoppedError):
+        _ = await rx1.receive()
+
+    with pytest.raises(ReceiverStoppedError):
+        _ = await rx2.receive()
+
+    comp_results: set[int] = set()
+    async for item in completing_merge:
+        comp_results.add(item)
+    assert comp_results == {1, 101, 2, 102, 3, 103, 4, 104, 5, 105}

tests/test_timer.py

@@ -13,6 +13,7 @@
 import pytest
 from hypothesis import strategies as st
 
+from frequenz.channels import ReceiverStoppedError
 from frequenz.channels.timer import (
     SkipMissedAndDrift,
     SkipMissedAndResync,
@@ -331,6 +332,21 @@ async def test_timer_construction_wrong_args() -> None:
         )
 
 
+async def test_timer_close_receiver() -> None:
+    """Test the autostart of a periodic timer."""
+    event_loop = asyncio.get_running_loop()
+
+    timer = Timer(timedelta(seconds=1.0), TriggerAllMissed())
+
+    drift = await timer.receive()
+    assert drift == pytest.approx(timedelta(seconds=0.0))
+    assert event_loop.time() == pytest.approx(1.0)
+
+    timer.close()
+    with pytest.raises(ReceiverStoppedError):
+        await timer.receive()
+
+
 async def test_timer_autostart() -> None:
     """Test the autostart of a periodic timer."""
     event_loop = asyncio.get_running_loop()