From 2e437e5b8b00b5acbb896827617c1cb162f1acd5 Mon Sep 17 00:00:00 2001
From: zixuanzh <zixuanzh@seas.upenn.edu>
Date: Sun, 17 Mar 2019 21:15:14 -0400
Subject: [PATCH] add test for listen event

---
 tests/libp2p/test_notify.py | 79 +++++++++++++++++++++++++++++++------
 tests/utils.py              | 10 +++++
 2 files changed, 77 insertions(+), 12 deletions(-)

diff --git a/tests/libp2p/test_notify.py b/tests/libp2p/test_notify.py
index a2fecaf..edaba2f 100644
--- a/tests/libp2p/test_notify.py
+++ b/tests/libp2p/test_notify.py
@@ -12,8 +12,9 @@ features are implemented in swarm
 import pytest
 
 from tests.utils import *
-from libp2p import new_node
+from libp2p import new_node, initialize_default_swarm
 from libp2p.network.notifee_interface import INotifee
+from libp2p.host.basic_host import BasicHost
 
 # pylint: disable=too-many-locals
 
@@ -39,7 +40,8 @@ class MyNotifee(INotifee):
         pass
 
     async def listen(self, network, multiaddr):
-        pass
+        self.events.append(["listened" + self.val_to_append_to_event,\
+            multiaddr])
 
     async def listen_close(self, network, multiaddr):
         pass
@@ -65,16 +67,6 @@ class InvalidNotifee():
     async def listen(self):
         assert False
 
-async def perform_two_host_set_up_custom_handler(handler):
-    transport_opt_list = [["/ip4/127.0.0.1/tcp/0"], ["/ip4/127.0.0.1/tcp/0"]]
-    (node_a, node_b) = await set_up_nodes_by_transport_opt(transport_opt_list) 
-
-    node_b.set_stream_handler("/echo/1.0.0", handler)
-
-    # Associate the peer with local ip address (see default parameters of Libp2p())
-    node_a.get_peerstore().add_addrs(node_b.get_id(), node_b.get_addrs(), 10)
-    return node_a, node_b
-
 @pytest.mark.asyncio
 async def test_one_notifier():
     node_a, node_b = await perform_two_host_set_up_custom_handler(echo_stream_handler)
@@ -146,6 +138,69 @@ async def test_one_notifier_on_two_nodes():
     # Success, terminate pending tasks.
     await cleanup()
 
+@pytest.mark.asyncio
+async def test_one_notifier_on_two_nodes_with_listen():
+    events_b = []
+
+    node_a_transport_opt = ["/ip4/127.0.0.1/tcp/0"]
+    node_a = await new_node(transport_opt=node_a_transport_opt)
+    await node_a.get_network().listen(multiaddr.Multiaddr(node_a_transport_opt[0]))
+
+    # Set up node_b swarm to pass into host
+    node_b_transport_opt = ["/ip4/127.0.0.1/tcp/0"]
+    node_b_multiaddr = multiaddr.Multiaddr(node_b_transport_opt[0])
+    node_b_swarm = initialize_default_swarm(transport_opt=node_b_transport_opt)
+    node_b = BasicHost(node_b_swarm)
+
+    async def my_stream_handler(stream):
+        # Ensure the listened, connected and opened_stream events were hit in Notifee obj
+        # and that the stream passed into opened_stream matches the stream created on
+        # node_b
+        assert events_b == [
+            ["listenedb", node_b_multiaddr], \
+            ["connectedb", stream.mplex_conn], \
+            ["opened_streamb", stream]
+        ]
+        while True:
+            read_string = (await stream.read()).decode()
+
+            resp = "ack:" + read_string
+            await stream.write(resp.encode())
+
+    # Add notifee for node_a
+    events_a = []
+    assert node_a.get_network().notify(MyNotifee(events_a, "a"))
+
+    # Add notifee for node_b
+    assert node_b.get_network().notify(MyNotifee(events_b, "b"))
+
+    # start listen on node_b_swarm
+    await node_b.get_network().listen(node_b_multiaddr)
+
+    node_b.set_stream_handler("/echo/1.0.0", my_stream_handler)
+    # Associate the peer with local ip address (see default parameters of Libp2p())
+    node_a.get_peerstore().add_addrs(node_b.get_id(), node_b.get_addrs(), 10)
+    stream = await node_a.new_stream(node_b.get_id(), ["/echo/1.0.0"])
+
+    # Ensure the connected and opened_stream events were hit in MyNotifee obj
+    # and that stream passed into opened_stream matches the stream created on
+    # node_a
+    assert events_a == [
+        ["connecteda", stream.mplex_conn], \
+        ["opened_streama", stream]
+    ]
+
+    messages = ["hello", "hello"]
+    for message in messages:
+        await stream.write(message.encode())
+
+        response = (await stream.read()).decode()
+
+        assert response == ("ack:" + message)
+
+    # Success, terminate pending tasks.
+    await cleanup()
+
 @pytest.mark.asyncio
 async def test_two_notifiers():
     node_a, node_b = await perform_two_host_set_up_custom_handler(echo_stream_handler)
diff --git a/tests/utils.py b/tests/utils.py
index e7f9b85..28f217e 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -29,3 +29,13 @@ async def echo_stream_handler(stream):
 
         resp = "ack:" + read_string
         await stream.write(resp.encode())
+
+async def perform_two_host_set_up_custom_handler(handler):
+    transport_opt_list = [["/ip4/127.0.0.1/tcp/0"], ["/ip4/127.0.0.1/tcp/0"]]
+    (node_a, node_b) = await set_up_nodes_by_transport_opt(transport_opt_list) 
+
+    node_b.set_stream_handler("/echo/1.0.0", handler)
+
+    # Associate the peer with local ip address (see default parameters of Libp2p())
+    node_a.get_peerstore().add_addrs(node_b.get_id(), node_b.get_addrs(), 10)
+    return node_a, node_b
\ No newline at end of file