Reorganize factories

libp2p/__init__.py

@@ -77,7 +77,6 @@ def initialize_default_swarm(
     muxer_opt: TMuxerOptions = None,
     sec_opt: TSecurityOptions = None,
     peerstore_opt: IPeerStore = None,
-    disc_opt: IPeerRouting = None,
 ) -> Swarm:
     """
     initialize swarm when no swarm is passed in.
@@ -87,7 +86,6 @@ def initialize_default_swarm(
     :param muxer_opt: optional choice of stream muxer
     :param sec_opt: optional choice of security upgrade
     :param peerstore_opt: optional peerstore
-    :param disc_opt: optional discovery
     :return: return a default swarm instance
     """
 
@@ -147,7 +145,6 @@ async def new_node(
             muxer_opt=muxer_opt,
             sec_opt=sec_opt,
             peerstore_opt=peerstore_opt,
-            disc_opt=disc_opt,
         )
 
     # TODO enable support for other host type

tests/factories.py

@@ -3,12 +3,13 @@ from typing import Dict, Tuple
 
 import factory
 
-from libp2p import generate_new_rsa_identity, initialize_default_swarm
+from libp2p import generate_new_rsa_identity, generate_peer_id_from
 from libp2p.crypto.keys import KeyPair
 from libp2p.host.basic_host import BasicHost
 from libp2p.network.connection.swarm_connection import SwarmConn
 from libp2p.network.stream.net_stream_interface import INetStream
 from libp2p.network.swarm import Swarm
+from libp2p.peer.peerstore import PeerStore
 from libp2p.pubsub.floodsub import FloodSub
 from libp2p.pubsub.gossipsub import GossipSub
 from libp2p.pubsub.pubsub import Pubsub
@@ -17,7 +18,9 @@ from libp2p.security.insecure.transport import PLAINTEXT_PROTOCOL_ID, InsecureTr
 import libp2p.security.secio.transport as secio
 from libp2p.stream_muxer.mplex.mplex import MPLEX_PROTOCOL_ID, Mplex
 from libp2p.stream_muxer.mplex.mplex_stream import MplexStream
+from libp2p.transport.tcp.tcp import TCP
 from libp2p.transport.typing import TMuxerOptions
+from libp2p.transport.upgrader import TransportUpgrader
 from libp2p.typing import TProtocol
 from tests.configs import LISTEN_MADDR
 from tests.pubsub.configs import (
@@ -37,33 +40,47 @@ def security_transport_factory(
         return {secio.ID: secio.Transport(key_pair)}
 
 
-def SwarmFactory(
+class SwarmFactory(factory.Factory):
-    is_secure: bool, key_pair: KeyPair, muxer_opt: TMuxerOptions = None
-) -> Swarm:
-    sec_opt = security_transport_factory(is_secure, key_pair)
-    return initialize_default_swarm(key_pair, sec_opt=sec_opt, muxer_opt=muxer_opt)
-
-
-class ListeningSwarmFactory(factory.Factory):
     class Meta:
         model = Swarm
 
+    class Params:
+        is_secure = False
+        key_pair = factory.LazyFunction(generate_new_rsa_identity)
+        muxer_opt = {MPLEX_PROTOCOL_ID: Mplex}
+
+    peer_id = factory.LazyAttribute(lambda o: generate_peer_id_from(o.key_pair))
+    peerstore = factory.LazyFunction(PeerStore)
+    upgrader = factory.LazyAttribute(
+        lambda o: TransportUpgrader(
+            security_transport_factory(o.is_secure, o.key_pair), o.muxer_opt
+        )
+    )
+    transport = factory.LazyFunction(TCP)
+
     @classmethod
     async def create_and_listen(
-        cls, is_secure: bool, muxer_opt: TMuxerOptions = None
+        cls, is_secure: bool, key_pair: KeyPair = None, muxer_opt: TMuxerOptions = None
-    ) -> Tuple[Swarm, KeyPair]:
+    ) -> Swarm:
-        key_pair = generate_new_rsa_identity()
+        # `factory.Factory.__init__` does *not* prepare a *default value* if we pass
-        swarm = SwarmFactory(is_secure, key_pair, muxer_opt=muxer_opt)
+        # an argument explicitly with `None`. If an argument is `None`, we don't pass it to
+        # `factory.Factory.__init__`, in order to let the function initialize it.
+        optional_kwargs = {}
+        if key_pair is not None:
+            optional_kwargs["key_pair"] = key_pair
+        if muxer_opt is not None:
+            optional_kwargs["muxer_opt"] = muxer_opt
+        swarm = cls(is_secure=is_secure, **optional_kwargs)
         await swarm.listen(LISTEN_MADDR)
-        return swarm, key_pair
+        return swarm
 
     @classmethod
     async def create_batch_and_listen(
         cls, is_secure: bool, number: int, muxer_opt: TMuxerOptions = None
-    ) -> Tuple[Tuple[Swarm, KeyPair], ...]:
+    ) -> Tuple[Swarm, ...]:
         return await asyncio.gather(
             *[
-                cls.create_and_listen(is_secure, muxer_opt=muxer_opt)
+                cls.create_and_listen(is_secure=is_secure, muxer_opt=muxer_opt)
                 for _ in range(number)
             ]
         )
@@ -75,26 +92,27 @@ class HostFactory(factory.Factory):
 
     class Params:
         is_secure = False
+        key_pair = factory.LazyFunction(generate_new_rsa_identity)
 
-    network = factory.LazyAttribute(lambda o: SwarmFactory(o.is_secure, o.key_pair))
+    public_key = factory.LazyAttribute(lambda o: o.key_pair.public_key)
-
+    network = factory.LazyAttribute(
-    @classmethod
+        lambda o: SwarmFactory(is_secure=o.is_secure, key_pair=o.key_pair)
-    async def create_and_listen(cls, is_secure: bool) -> BasicHost:
+    )
-        swarms_and_keys = await ListeningSwarmFactory.create_batch_and_listen(
-            is_secure, 1
-        )
-        swarm, key_pair = swarms_and_keys[0]
-        return BasicHost(key_pair.public_key, swarm)
 
     @classmethod
     async def create_batch_and_listen(
         cls, is_secure: bool, number: int
     ) -> Tuple[BasicHost, ...]:
-        swarms_and_keys = await ListeningSwarmFactory.create_batch_and_listen(
+        key_pairs = [generate_new_rsa_identity() for _ in range(number)]
-            is_secure, number
+        swarms = await asyncio.gather(
+            *[
+                SwarmFactory.create_and_listen(is_secure, key_pair)
+                for key_pair in key_pairs
+            ]
         )
         return tuple(
-            BasicHost(key_pair.public_key, swarm) for swarm, key_pair in swarms_and_keys
+            BasicHost(key_pair.public_key, swarm)
+            for key_pair, swarm in zip(key_pairs, swarms)
         )
 
 
@@ -132,21 +150,15 @@ class PubsubFactory(factory.Factory):
 async def swarm_pair_factory(
     is_secure: bool, muxer_opt: TMuxerOptions = None
 ) -> Tuple[Swarm, Swarm]:
-    swarms_and_keys = await ListeningSwarmFactory.create_batch_and_listen(
+    swarms = await SwarmFactory.create_batch_and_listen(
         is_secure, 2, muxer_opt=muxer_opt
     )
-    swarms = tuple(swarm for swarm, _key_pair in swarms_and_keys)
     await connect_swarm(swarms[0], swarms[1])
     return swarms[0], swarms[1]
 
 
 async def host_pair_factory(is_secure) -> Tuple[BasicHost, BasicHost]:
-    hosts = await asyncio.gather(
+    hosts = await HostFactory.create_batch_and_listen(is_secure, 2)
-        *[
-            HostFactory.create_and_listen(is_secure),
-            HostFactory.create_and_listen(is_secure),
-        ]
-    )
     await connect(hosts[0], hosts[1])
     return hosts[0], hosts[1]
 

tests/network/test_notify.py

@@ -14,7 +14,6 @@ import enum
 
 import pytest
 
-from libp2p.crypto.rsa import create_new_key_pair
 from libp2p.network.notifee_interface import INotifee
 from tests.configs import LISTEN_MADDR
 from tests.factories import SwarmFactory
@@ -57,7 +56,7 @@ class MyNotifee(INotifee):
 
 @pytest.mark.asyncio
 async def test_notify(is_host_secure):
-    swarms = [SwarmFactory(is_host_secure, create_new_key_pair()) for _ in range(2)]
+    swarms = [SwarmFactory(is_secure=is_host_secure) for _ in range(2)]
 
     events_0_0 = []
     events_1_0 = []

tests/network/test_swarm.py

@@ -3,16 +3,13 @@ import asyncio
 import pytest
 
 from libp2p.network.exceptions import SwarmException
-from tests.factories import ListeningSwarmFactory
+from tests.factories import SwarmFactory
 from tests.utils import connect_swarm
 
 
 @pytest.mark.asyncio
 async def test_swarm_dial_peer(is_host_secure):
-    swarms_and_keys = await ListeningSwarmFactory.create_batch_and_listen(
+    swarms = await SwarmFactory.create_batch_and_listen(is_host_secure, 3)
-        is_host_secure, 3
-    )
-    swarms = tuple(swarm for swarm, _key_pair in swarms_and_keys)
     # Test: No addr found.
     with pytest.raises(SwarmException):
         await swarms[0].dial_peer(swarms[1].get_peer_id())
@@ -44,10 +41,7 @@ async def test_swarm_dial_peer(is_host_secure):
 
 @pytest.mark.asyncio
 async def test_swarm_close_peer(is_host_secure):
-    swarms_and_keys = await ListeningSwarmFactory.create_batch_and_listen(
+    swarms = await SwarmFactory.create_batch_and_listen(is_host_secure, 3)
-        is_host_secure, 3
-    )
-    swarms = tuple(swarm for swarm, _key_pair in swarms_and_keys)
     # 0 <> 1 <> 2
     await connect_swarm(swarms[0], swarms[1])
     await connect_swarm(swarms[1], swarms[2])