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])