From cf1d082628e94e241dea92a125ee1ae48e46bb65 Mon Sep 17 00:00:00 2001
From: Hugo Osvaldo Barrera <hugo@barrera.io>
Date: Mon, 16 Aug 2021 19:15:18 +0200
Subject: [PATCH] Use context managers for aio connectors

Not sure why we didn't do this initially, but this ensures that we
always close all connectors properly, and also gives much clearer scope
regarding their life-cycles.
---
 tests/conftest.py          |   5 +-
 vdirsyncer/cli/__init__.py | 110 +++++++++++++++++--------------------
 2 files changed, 52 insertions(+), 63 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index f4985b7..2807473 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -63,8 +63,5 @@ async def aio_session(event_loop):
 
 @pytest.fixture
 async def aio_connector(event_loop):
-    conn = aiohttp.TCPConnector(limit_per_host=16)
-    try:
+    async with aiohttp.TCPConnector(limit_per_host=16) as conn:
         yield conn
-    finally:
-        await conn.close()
diff --git a/vdirsyncer/cli/__init__.py b/vdirsyncer/cli/__init__.py
index 664aa9e..4e27bc5 100644
--- a/vdirsyncer/cli/__init__.py
+++ b/vdirsyncer/cli/__init__.py
@@ -127,27 +127,25 @@ def sync(ctx, collections, force_delete):
     from .tasks import sync_collection
 
     async def main(collections):
-        conn = aiohttp.TCPConnector(limit_per_host=16)
-
-        tasks = []
-        for pair_name, collections in collections:
-            async for collection, config in prepare_pair(
-                pair_name=pair_name,
-                collections=collections,
-                config=ctx.config,
-                connector=conn,
-            ):
-                tasks.append(
-                    sync_collection(
-                        collection=collection,
-                        general=config,
-                        force_delete=force_delete,
-                        connector=conn,
+        async with aiohttp.TCPConnector(limit_per_host=16) as conn:
+            tasks = []
+            for pair_name, collections in collections:
+                async for collection, config in prepare_pair(
+                    pair_name=pair_name,
+                    collections=collections,
+                    config=ctx.config,
+                    connector=conn,
+                ):
+                    tasks.append(
+                        sync_collection(
+                            collection=collection,
+                            general=config,
+                            force_delete=force_delete,
+                            connector=conn,
+                        )
                     )
-                )
 
-        await asyncio.gather(*tasks)
-        await conn.close()
+            await asyncio.gather(*tasks)
 
     asyncio.run(main(collections))
 
@@ -166,28 +164,26 @@ def metasync(ctx, collections):
     from .tasks import prepare_pair
 
     async def main(collections):
-        conn = aiohttp.TCPConnector(limit_per_host=16)
+        async with aiohttp.TCPConnector(limit_per_host=16) as conn:
 
-        for pair_name, collections in collections:
-            collections = prepare_pair(
-                pair_name=pair_name,
-                collections=collections,
-                config=ctx.config,
-                connector=conn,
-            )
+            for pair_name, collections in collections:
+                collections = prepare_pair(
+                    pair_name=pair_name,
+                    collections=collections,
+                    config=ctx.config,
+                    connector=conn,
+                )
 
-            await asyncio.gather(
-                *[
-                    metasync_collection(
-                        collection=collection,
-                        general=config,
-                        connector=conn,
-                    )
-                    async for collection, config in collections
-                ]
-            )
-
-        await conn.close()
+                await asyncio.gather(
+                    *[
+                        metasync_collection(
+                            collection=collection,
+                            general=config,
+                            connector=conn,
+                        )
+                        async for collection, config in collections
+                    ]
+                )
 
     asyncio.run(main(collections))
 
@@ -213,18 +209,15 @@ def discover(ctx, pairs, list):
     config = ctx.config
 
     async def main():
-        conn = aiohttp.TCPConnector(limit_per_host=16)
-
-        for pair_name in pairs or config.pairs:
-            await discover_collections(
-                status_path=config.general["status_path"],
-                pair=config.get_pair(pair_name),
-                from_cache=False,
-                list_collections=list,
-                connector=conn,
-            )
-
-        await conn.close()
+        async with aiohttp.TCPConnector(limit_per_host=16) as conn:
+            for pair_name in pairs or config.pairs:
+                await discover_collections(
+                    status_path=config.general["status_path"],
+                    pair=config.get_pair(pair_name),
+                    from_cache=False,
+                    list_collections=list,
+                    connector=conn,
+                )
 
     asyncio.run(main())
 
@@ -267,14 +260,13 @@ def repair(ctx, collection, repair_unsafe_uid):
     click.confirm("Do you want to continue?", abort=True)
 
     async def main():
-        conn = aiohttp.TCPConnector(limit_per_host=16)
-        await repair_collection(
-            ctx.config,
-            collection,
-            repair_unsafe_uid=repair_unsafe_uid,
-            connector=conn,
-        )
-        await conn.close()
+        async with aiohttp.TCPConnector(limit_per_host=16) as conn:
+            await repair_collection(
+                ctx.config,
+                collection,
+                repair_unsafe_uid=repair_unsafe_uid,
+                connector=conn,
+            )
 
     asyncio.run(main())