Add SubscriptionAPI

And `TrioSubscriptionAPI`, to make subscription io-agnostic.

tests/pubsub/test_floodsub.py

@@ -27,7 +27,7 @@ async def test_simple_two_nodes():
 
         await pubsubs_fsub[0].publish(topic, data)
 
-        res_b = await sub_b.receive()
+        res_b = await sub_b.get()
 
         # Check that the msg received by node_b is the same
         # as the message sent by node_a
@@ -75,12 +75,9 @@ async def test_lru_cache_two_nodes(monkeypatch):
         await trio.sleep(0.25)
 
         for index in expected_received_indices:
-            res_b = await sub_b.receive()
+            res_b = await sub_b.get()
             assert res_b.data == _make_testing_data(index)
 
-        with pytest.raises(trio.WouldBlock):
-            sub_b.receive_nowait()
-
 
 @pytest.mark.parametrize("test_case_obj", floodsub_protocol_pytest_params)
 @pytest.mark.trio

tests/pubsub/test_gossipsub.py

@@ -196,7 +196,7 @@ async def test_dense():
             await trio.sleep(0.5)
             # Assert that all blocking queues receive the message
             for queue in queues:
-                msg = await queue.receive()
+                msg = await queue.get()
                 assert msg.data == msg_content
 
 
@@ -229,7 +229,7 @@ async def test_fanout():
             await trio.sleep(0.5)
             # Assert that all blocking queues receive the message
             for sub in subs:
-                msg = await sub.receive()
+                msg = await sub.get()
                 assert msg.data == msg_content
 
         # Subscribe message origin
@@ -248,7 +248,7 @@ async def test_fanout():
             await trio.sleep(0.5)
             # Assert that all blocking queues receive the message
             for sub in subs:
-                msg = await sub.receive()
+                msg = await sub.get()
                 assert msg.data == msg_content
 
 
@@ -287,7 +287,7 @@ async def test_fanout_maintenance():
             await trio.sleep(0.5)
             # Assert that all blocking queues receive the message
             for queue in queues:
-                msg = await queue.receive()
+                msg = await queue.get()
                 assert msg.data == msg_content
 
         for sub in pubsubs_gsub:
@@ -319,7 +319,7 @@ async def test_fanout_maintenance():
             await trio.sleep(0.5)
             # Assert that all blocking queues receive the message
             for queue in queues:
-                msg = await queue.receive()
+                msg = await queue.get()
                 assert msg.data == msg_content
 
 
@@ -346,5 +346,5 @@ async def test_gossip_propagation():
         await trio.sleep(2)
 
         # should be able to read message
-        msg = await queue_1.receive()
+        msg = await queue_1.get()
         assert msg.data == msg_content

tests/pubsub/test_pubsub.py

@@ -384,7 +384,7 @@ async def test_handle_talk():
             len(pubsubs_fsub[0].topic_ids) == 1
             and sub == pubsubs_fsub[0].subscribed_topics_receive[TESTING_TOPIC]
         )
-        assert (await sub.receive()) == msg_0
+        assert (await sub.get()) == msg_0
 
 
 @pytest.mark.trio
@@ -486,7 +486,7 @@ async def test_push_msg(monkeypatch):
             with trio.fail_after(0.1):
                 await event.wait()
             # Test: Subscribers are notified when `push_msg` new messages.
-            assert (await sub.receive()) == msg_1
+            assert (await sub.get()) == msg_1
 
         with mock_router_publish() as event:
             # Test: add a topic validator and `push_msg` the message that

tests/pubsub/test_subscription.py

@@ -0,0 +1,77 @@
+import math
+
+import pytest
+import trio
+
+from libp2p.pubsub.pb import rpc_pb2
+from libp2p.pubsub.subscription import TrioSubscriptionAPI
+
+GET_TIMEOUT = 0.001
+
+
+def make_trio_subscription():
+    send_channel, receive_channel = trio.open_memory_channel(math.inf)
+    return send_channel, TrioSubscriptionAPI(receive_channel)
+
+
+def make_pubsub_msg():
+    return rpc_pb2.Message()
+
+
+async def send_something(send_channel):
+    msg = make_pubsub_msg()
+    await send_channel.send(msg)
+    return msg
+
+
+@pytest.mark.trio
+async def test_trio_subscription_get():
+    send_channel, sub = make_trio_subscription()
+    data_0 = await send_something(send_channel)
+    data_1 = await send_something(send_channel)
+    assert data_0 == await sub.get()
+    assert data_1 == await sub.get()
+    # No more message
+    with pytest.raises(trio.TooSlowError):
+        with trio.fail_after(GET_TIMEOUT):
+            await sub.get()
+
+
+@pytest.mark.trio
+async def test_trio_subscription_iter():
+    send_channel, sub = make_trio_subscription()
+    received_data = []
+
+    async def iter_subscriptions(subscription):
+        async for data in sub:
+            received_data.append(data)
+
+    async with trio.open_nursery() as nursery:
+        nursery.start_soon(iter_subscriptions, sub)
+        await send_something(send_channel)
+        await send_something(send_channel)
+        await send_channel.aclose()
+
+    assert len(received_data) == 2
+
+
+@pytest.mark.trio
+async def test_trio_subscription_cancel():
+    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 send_something(send_channel)
+    # Test: No side effect when cancelled twice.
+    await sub.cancel()
+
+
+@pytest.mark.trio
+async def test_trio_subscription_async_context_manager():
+    send_channel, sub = make_trio_subscription()
+    async with sub:
+        # 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):
+        await send_something(send_channel)