Test the auto-close feature of broadcast channels

Signed-off-by: Sahas Subramanian <sahas.subramanian@proton.me>

Author: shsms

tests/test_broadcast.py

@@ -17,6 +17,7 @@
     ReceiverStoppedError,
     Sender,
     SenderError,
+    broadcast,
 )
 
 
@@ -425,3 +426,50 @@ async def test_broadcast_close_receiver() -> None:
 
     with pytest.raises(ReceiverStoppedError):
         _ = await receiver_2.receive()
+
+
+async def test_broadcast_auto_close_1() -> None:
+    """Ensure broadcast auto close works when all receivers are closed."""
+    sender, receiver = broadcast(int, name="auto-close-test")
+
+    receiver_2 = sender.subscribe()
+
+    await sender.send(1)
+
+    assert (await receiver.receive()) == 1
+    assert (await receiver_2.receive()) == 1
+
+    receiver.close()
+
+    await sender.send(2)
+
+    assert (await receiver_2.receive()) == 2
+
+    receiver_2.close()
+
+    with pytest.raises(SenderError) as excinfo:
+        await sender.send(3)
+    assert isinstance(excinfo.value.__cause__, ChannelClosedError)
+
+
+async def test_broadcast_auto_close_2() -> None:
+    """Ensure broadcast auto close works when all senders are closed."""
+    sender, receiver = broadcast(int, name="auto-close-test")
+
+    await sender.send(1)
+
+    assert (await receiver.receive()) == 1
+
+    sender_2 = sender.clone()
+
+    sender.close()
+
+    await sender_2.send(2)
+
+    sender_2.close()
+
+    assert (await receiver.receive()) == 2
+
+    with pytest.raises(ReceiverStoppedError) as excinfo:
+        await receiver.receive()
+    assert isinstance(excinfo.value.__cause__, ChannelClosedError)