Add clean-up logics into TrioSubscriptionAPI

Register an `unsubscribe_fn` when initializing the TrioSubscriptionAPI.
`unsubscribe_fn` is called when subscription is unsubscribed.

tests/pubsub/test_pubsub.py

@@ -413,7 +413,39 @@ async def test_message_all_peers(monkeypatch, is_host_secure):
 
 
 @pytest.mark.trio
-async def test_publish(monkeypatch):
+async def test_subscribe_and_publish():
+    async with PubsubFactory.create_batch_with_floodsub(1) as pubsubs_fsub:
+        pubsub = pubsubs_fsub[0]
+
+        list_data = [b"d0", b"d1"]
+        event_receive_data_started = trio.Event()
+
+        async def publish_data(topic):
+            await event_receive_data_started.wait()
+            for data in list_data:
+                await pubsub.publish(topic, data)
+
+        async def receive_data(topic):
+            i = 0
+            event_receive_data_started.set()
+            assert topic not in pubsub.topic_ids
+            subscription = await pubsub.subscribe(topic)
+            async with subscription:
+                assert topic in pubsub.topic_ids
+                async for msg in subscription:
+                    assert msg.data == list_data[i]
+                    i += 1
+                    if i == len(list_data):
+                        break
+            assert topic not in pubsub.topic_ids
+
+        async with trio.open_nursery() as nursery:
+            nursery.start_soon(receive_data, TESTING_TOPIC)
+            nursery.start_soon(publish_data, TESTING_TOPIC)
+
+
+@pytest.mark.trio
+async def test_publish_push_msg_is_called(monkeypatch):
     msg_forwarders = []
     msgs = []
 

tests/pubsub/test_subscription.py

@@ -11,7 +11,14 @@ GET_TIMEOUT = 0.001
 
 def make_trio_subscription():
     send_channel, receive_channel = trio.open_memory_channel(math.inf)
-    return send_channel, TrioSubscriptionAPI(receive_channel)
+
+    async def unsubscribe_fn():
+        await send_channel.aclose()
+
+    return (
+        send_channel,
+        TrioSubscriptionAPI(receive_channel, unsubscribe_fn=unsubscribe_fn),
+    )
 
 
 def make_pubsub_msg():
@@ -56,14 +63,14 @@ async def test_trio_subscription_iter():
 
 
 @pytest.mark.trio
-async def test_trio_subscription_cancel():
+async def test_trio_subscription_unsubscribe():
     send_channel, sub = make_trio_subscription()
-    await sub.cancel()
-    # Test: If the subscription is cancelled, `send_channel` should be broken.
-    with pytest.raises(trio.BrokenResourceError):
+    await sub.unsubscribe()
+    # Test: If the subscription is unsubscribed, `send_channel` should be closed.
+    with pytest.raises(trio.ClosedResourceError):
         await send_something(send_channel)
     # Test: No side effect when cancelled twice.
-    await sub.cancel()
+    await sub.unsubscribe()
 
 
 @pytest.mark.trio
@@ -73,5 +80,5 @@ async def test_trio_subscription_async_context_manager():
         # Test: `sub` is not cancelled yet, so `send_something` works fine.
         await send_something(send_channel)
     # Test: `sub` is cancelled, `send_something` fails
-    with pytest.raises(trio.BrokenResourceError):
+    with pytest.raises(trio.ClosedResourceError):
         await send_something(send_channel)