From 30dee28ef2e876e73a2ccda8b16a8eacd5267784 Mon Sep 17 00:00:00 2001 From: Alex Stokes Date: Thu, 7 Nov 2019 20:59:01 -0800 Subject: [PATCH] add asynccontextmanager utility for a pair of connected hosts --- tests/factories.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/factories.py b/tests/factories.py index d59f2278..b1ac527f 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]: