From 0356380996a880ea7f8555b39441914d8d04d28e Mon Sep 17 00:00:00 2001
From: mhchia <kevin.mh.chia@gmail.com>
Date: Sun, 15 Sep 2019 20:44:48 +0800
Subject: [PATCH] Add tests for swarm, and debug

Fix `swarm_pair_factory`
---
 libp2p/network/swarm.py     | 11 ++++++++++
 tests/factories.py          |  2 +-
 tests/network/test_swarm.py | 44 +++++++++++++++++++++++++++++++++++++
 3 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/libp2p/network/swarm.py b/libp2p/network/swarm.py
index 3c83846..272a8a9 100644
--- a/libp2p/network/swarm.py
+++ b/libp2p/network/swarm.py
@@ -261,12 +261,18 @@ class Swarm(INetwork):
     async def close_peer(self, peer_id: ID) -> None:
         if peer_id not in self.connections:
             return
+        # TODO: Should be changed to close multisple connections,
+        #   if we have several connections per peer in the future.
         connection = self.connections[peer_id]
         await connection.close()
 
         logger.debug("successfully close the connection to peer %s", peer_id)
 
     async def add_conn(self, muxed_conn: IMuxedConn) -> SwarmConn:
+        """
+        Add a `IMuxedConn` to `Swarm` as a `SwarmConn`, notify "connected",
+        and start to monitor the connection for its new streams and disconnection.
+        """
         swarm_conn = SwarmConn(muxed_conn, self)
         # Store muxed_conn with peer id
         self.connections[muxed_conn.peer_id] = swarm_conn
@@ -278,7 +284,12 @@ class Swarm(INetwork):
         return swarm_conn
 
     def remove_conn(self, swarm_conn: SwarmConn) -> None:
+        """
+        Simply remove the connection from Swarm's records, without closing the connection.
+        """
         peer_id = swarm_conn.conn.peer_id
+        if peer_id not in self.connections:
+            return
         # TODO: Should be changed to remove the exact connection,
         #   if we have several connections per peer in the future.
         del self.connections[peer_id]
diff --git a/tests/factories.py b/tests/factories.py
index e9ec196..e39b12d 100644
--- a/tests/factories.py
+++ b/tests/factories.py
@@ -112,7 +112,7 @@ class PubsubFactory(factory.Factory):
 
 
 async def swarm_pair_factory(is_secure: bool) -> Tuple[Swarm, Swarm]:
-    swarms = await ListeningSwarmFactory.create_batch_and_listen(2)
+    swarms = await ListeningSwarmFactory.create_batch_and_listen(is_secure, 2)
     await connect_swarm(swarms[0], swarms[1])
     return swarms[0], swarms[1]
 
diff --git a/tests/network/test_swarm.py b/tests/network/test_swarm.py
index 3b27f43..cf8eadf 100644
--- a/tests/network/test_swarm.py
+++ b/tests/network/test_swarm.py
@@ -2,10 +2,43 @@ import asyncio
 
 import pytest
 
+from libp2p.network.exceptions import SwarmException
 from tests.factories import ListeningSwarmFactory
 from tests.utils import connect_swarm
 
 
+@pytest.mark.asyncio
+async def test_swarm_dial_peer(is_host_secure):
+    swarms = await ListeningSwarmFactory.create_batch_and_listen(is_host_secure, 3)
+    # Test: No addr found.
+    with pytest.raises(SwarmException):
+        await swarms[0].dial_peer(swarms[1].get_peer_id())
+
+    # Test: len(addr) in the peerstore is 0.
+    swarms[0].peerstore.add_addrs(swarms[1].get_peer_id(), [], 10000)
+    with pytest.raises(SwarmException):
+        await swarms[0].dial_peer(swarms[1].get_peer_id())
+
+    # Test: Succeed if addrs of the peer_id are present in the peerstore.
+    addrs = tuple(
+        addr
+        for transport in swarms[1].listeners.values()
+        for addr in transport.get_addrs()
+    )
+    swarms[0].peerstore.add_addrs(swarms[1].get_peer_id(), addrs, 10000)
+    await swarms[0].dial_peer(swarms[1].get_peer_id())
+    assert swarms[0].get_peer_id() in swarms[1].connections
+    assert swarms[1].get_peer_id() in swarms[0].connections
+
+    # Test: Reuse connections when we already have ones with a peer.
+    conn_to_1 = swarms[0].connections[swarms[1].get_peer_id()]
+    conn = await swarms[0].dial_peer(swarms[1].get_peer_id())
+    assert conn is conn_to_1
+
+    # Clean up
+    await asyncio.gather(*[swarm.close() for swarm in swarms])
+
+
 @pytest.mark.asyncio
 async def test_swarm_close_peer(is_host_secure):
     swarms = await ListeningSwarmFactory.create_batch_and_listen(is_host_secure, 3)
@@ -47,3 +80,14 @@ async def test_swarm_close_peer(is_host_secure):
 
     # Clean up
     await asyncio.gather(*[swarm.close() for swarm in swarms])
+
+
+@pytest.mark.asyncio
+async def test_swarm_remove_conn(swarm_pair):
+    swarm_0, swarm_1 = swarm_pair
+    conn_0 = swarm_0.connections[swarm_1.get_peer_id()]
+    swarm_0.remove_conn(conn_0)
+    assert swarm_1.get_peer_id() not in swarm_0.connections
+    # Test: Remove twice. There should not be errors.
+    swarm_0.remove_conn(conn_0)
+    assert swarm_1.get_peer_id() not in swarm_0.connections