fix context manager

matrix_alertbot/alertmanager.py

@@ -6,7 +6,6 @@ from typing import Any, Dict, List
 import aiohttp
 import pytimeparse
 from aiohttp import ClientError
-from aiotools import AsyncContextManager
 from diskcache import Cache
 
 from matrix_alertbot.errors import (

matrix_alertbot/main.py

@@ -5,7 +5,6 @@ import sys
 from asyncio import TimeoutError
 from time import sleep
 
-import aiotools
 from aiohttp import ClientConnectionError, ServerDisconnectedError
 from diskcache import Cache
 from nio import (
@@ -52,19 +51,9 @@ def create_matrix_client(config: Config) -> AsyncClient:
     return client
 
 
-async def start_matrix_client(cache: Cache, config: Config) -> bool:
+async def start_matrix_client(
-    async with aiotools.closing_async(create_matrix_client(config)) as client:
+    client: AsyncClient, cache: Cache, config: Config
-        # Configure Alertmanager client
+) -> bool:
-        async with AlertmanagerClient(config.alertmanager_url, cache) as alertmanager:
-            # Set up event callbacks
-            callbacks = Callbacks(client, alertmanager, cache, config)
-            client.add_event_callback(callbacks.message, (RoomMessageText,))
-            client.add_event_callback(
-                callbacks.invite_event_filtered_callback, (InviteMemberEvent,)
-            )
-            client.add_event_callback(callbacks.decryption_failure, (MegolmEvent,))
-            client.add_event_callback(callbacks.unknown, (UnknownEvent,))
-
     # Keep trying to reconnect on failure (with some time in-between)
     while True:
         try:
@@ -85,9 +74,7 @@ async def start_matrix_client(cache: Cache, config: Config) -> bool:
 
                     # Check if login failed
                     if type(login_response) == LoginError:
-                                logger.error(
+                        logger.error("Failed to login: %s", login_response.message)
-                                    "Failed to login: %s", login_response.message
-                                )
                         return False
                 except LocalProtocolError as e:
                     # There's an edge case here where the user hasn't installed the correct C
@@ -104,20 +91,13 @@ async def start_matrix_client(cache: Cache, config: Config) -> bool:
 
             logger.info(f"Logged in as {config.user_id}")
             await client.sync_forever(timeout=30000, full_state=True)
-
         except (ClientConnectionError, ServerDisconnectedError, TimeoutError):
-                    logger.warning(
+            logger.warning("Unable to connect to homeserver, retrying in 15s...")
-                        "Unable to connect to homeserver, retrying in 15s..."
-                    )
 
             # Sleep so we don't bombard the server with login requests
             sleep(15)
-
+        finally:
-
+            client.close()
-async def start_webhook_server(cache: Cache, config: Config) -> None:
-    async with aiotools.closing_async(create_matrix_client(config)) as client:
-        async with Webhook(client, cache, config) as webhook_server:
-            await webhook_server.start()
 
 
 def main() -> None:
@@ -133,14 +113,35 @@ def main() -> None:
     # Read the parsed config file and create a Config object
     config = Config(config_path)
 
+    client = create_matrix_client(config)
+
     # Configure the cache
     cache = Cache(config.cache_dir)
 
+    # Configure Alertmanager client
+    alertmanager = AlertmanagerClient(config.alertmanager_url, cache)
+
+    # Set up event callbacks
+    callbacks = Callbacks(client, alertmanager, cache, config)
+    client.add_event_callback(callbacks.message, (RoomMessageText,))
+    client.add_event_callback(
+        callbacks.invite_event_filtered_callback, (InviteMemberEvent,)
+    )
+    client.add_event_callback(callbacks.decryption_failure, (MegolmEvent,))
+    client.add_event_callback(callbacks.unknown, (UnknownEvent,))
+
+    # Configure webhook server
+    webhook_server = Webhook(client, cache, config)
+
     loop = asyncio.get_event_loop()
-    loop.create_task(start_webhook_server(cache, config))
+    loop.create_task(webhook_server.start())
-    loop.create_task(start_matrix_client(cache, config))
+    loop.create_task(start_matrix_client(client, cache, config))
 
     try:
         loop.run_forever()
     except Exception as e:
         logger.error(e)
+    finally:
+        loop.run_until_complete(webhook_server.close())
+        loop.run_until_complete(alertmanager.close())
+        loop.run_until_complete(client.close())

matrix_alertbot/webhook.py

@@ -4,7 +4,6 @@ import logging
 from typing import Any
 
 from aiohttp import web, web_request
-from aiotools import AsyncContextManager
 from diskcache import Cache
 from nio import AsyncClient, SendRetryError