diff --git a/tests/factories.py b/tests/factories.py index d59f227..b1ac527 100644 --- a/tests/factories.py +++ b/tests/factories.py @@ -1,4 +1,5 @@ import asyncio +from contextlib import asynccontextmanager from typing import Dict, Tuple import factory @@ -163,6 +164,14 @@ async def host_pair_factory(is_secure) -> Tuple[BasicHost, BasicHost]: return hosts[0], hosts[1] +@asynccontextmanager +async def pair_of_connected_hosts(is_secure=True): + a, b = await host_pair_factory(is_secure) + yield a, b + close_tasks = (a.close(), b.close()) + await asyncio.gather(*close_tasks) + + async def swarm_conn_pair_factory( is_secure: bool, muxer_opt: TMuxerOptions = None ) -> Tuple[SwarmConn, Swarm, SwarmConn, Swarm]: