From 276ac4d8ab1f5a7cb29c6c76261b0574a8434b17 Mon Sep 17 00:00:00 2001
From: mhchia <kevin.mh.chia@gmail.com>
Date: Sun, 15 Sep 2019 14:59:59 +0800
Subject: [PATCH] Add initial test for `Swarm.close_peer`

---
 tests/factories.py          | 32 +++++++++++++++++++-------------
 tests/network/test_swarm.py |  4 ++--
 2 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/tests/factories.py b/tests/factories.py
index efa16c8..e9ec196 100644
--- a/tests/factories.py
+++ b/tests/factories.py
@@ -6,7 +6,6 @@ import factory
 from libp2p import generate_new_rsa_identity, initialize_default_swarm
 from libp2p.crypto.keys import KeyPair
 from libp2p.host.basic_host import BasicHost
-from libp2p.host.host_interface import IHost
 from libp2p.network.stream.net_stream_interface import INetStream
 from libp2p.network.swarm import Swarm
 from libp2p.pubsub.floodsub import FloodSub
@@ -34,19 +33,19 @@ def security_transport_factory(
         return {secio.ID: secio.Transport(key_pair)}
 
 
-class SwarmFactory(factory.Factory):
+def SwarmFactory(is_secure: bool) -> Swarm:
+    key_pair = generate_new_rsa_identity()
+    sec_opt = security_transport_factory(False, key_pair)
+    return initialize_default_swarm(key_pair, sec_opt=sec_opt)
+
+
+class ListeningSwarmFactory(factory.Factory):
     class Meta:
         model = Swarm
 
-    @classmethod
-    def _create(cls, is_secure=False):
-        key_pair = generate_new_rsa_identity()
-        sec_opt = security_transport_factory(is_secure, key_pair)
-        return initialize_default_swarm(key_pair, sec_opt=sec_opt)
-
     @classmethod
     async def create_and_listen(cls, is_secure: bool) -> Swarm:
-        swarm = cls._create(is_secure)
+        swarm = SwarmFactory(is_secure)
         await swarm.listen(LISTEN_MADDR)
         return swarm
 
@@ -69,9 +68,16 @@ class HostFactory(factory.Factory):
     network = factory.LazyAttribute(lambda o: SwarmFactory(o.is_secure))
 
     @classmethod
-    async def create_and_listen(cls, is_secure: bool) -> IHost:
-        swarm = await SwarmFactory.create_and_listen(is_secure)
-        return BasicHost(swarm)
+    async def create_and_listen(cls, is_secure: bool) -> BasicHost:
+        swarms = await ListeningSwarmFactory.create_batch_and_listen(is_secure, 1)
+        return BasicHost(swarms[0])
+
+    @classmethod
+    async def create_batch_and_listen(
+        cls, is_secure: bool, number: int
+    ) -> Tuple[BasicHost, ...]:
+        swarms = await ListeningSwarmFactory.create_batch_and_listen(is_secure, number)
+        return tuple(BasicHost(swarm) for swarm in range(swarms))
 
 
 class FloodsubFactory(factory.Factory):
@@ -106,7 +112,7 @@ class PubsubFactory(factory.Factory):
 
 
 async def swarm_pair_factory(is_secure: bool) -> Tuple[Swarm, Swarm]:
-    swarms = await SwarmFactory.create_batch_and_listen(2)
+    swarms = await ListeningSwarmFactory.create_batch_and_listen(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 e531de0..3b27f43 100644
--- a/tests/network/test_swarm.py
+++ b/tests/network/test_swarm.py
@@ -2,13 +2,13 @@ import asyncio
 
 import pytest
 
-from tests.factories import SwarmFactory
+from tests.factories import ListeningSwarmFactory
 from tests.utils import connect_swarm
 
 
 @pytest.mark.asyncio
 async def test_swarm_close_peer(is_host_secure):
-    swarms = await SwarmFactory.create_batch_and_listen(is_host_secure, 3)
+    swarms = await ListeningSwarmFactory.create_batch_and_listen(is_host_secure, 3)
     # 0 <> 1 <> 2
     await connect_swarm(swarms[0], swarms[1])
     await connect_swarm(swarms[1], swarms[2])