Fix subscriber not properly unsubscribing when an exception is raised…

… inside the context manager (#112)

* add testcases for unsubsribe

* fix the context manager of the function "subscribe" not removing the associated queue from the channel in case of a raised exception inside the context manager (e.g. being used inside a generator that gets closed raising GeneratorExit)

* activate broadcaster via context manager

---------

Co-authored-by: alex.oleshkevich <[email protected]>

broadcaster/_base.py

@@ -90,12 +90,11 @@ async def subscribe(self, channel: str) -> AsyncIterator["Subscriber"]:
                 self._subscribers[channel].add(queue)
 
             yield Subscriber(queue)
-
+        finally:
             self._subscribers[channel].remove(queue)
             if not self._subscribers.get(channel):
                 del self._subscribers[channel]
                 await self._backend.unsubscribe(channel)
-        finally:
             await queue.put(None)
 
 

tests/test_unsubscribe.py

@@ -0,0 +1,25 @@
+import pytest
+from broadcaster import Broadcast
+
+
+@pytest.mark.asyncio
+async def test_unsubscribe():
+    """The queue should be removed when the context manager is left."""
+    async with Broadcast("memory://") as broadcast:
+        async with broadcast.subscribe("chatroom"):
+            pass
+
+        assert "chatroom" not in broadcast._subscribers
+
+
+@pytest.mark.asyncio
+async def test_unsubscribe_w_exception():
+    """In case an exception is raised inside the context manager, the queue should be removed."""
+    async with Broadcast("memory://") as broadcast:
+        try:
+            async with broadcast.subscribe("chatroom"):
+                raise RuntimeError("MyException")
+        except RuntimeError:
+            pass
+
+        assert "chatroom" not in broadcast._subscribers