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,72 +51,53 @@ 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:
+    # Keep trying to reconnect on failure (with some time in-between)
-            # Set up event callbacks
+    while True:
-            callbacks = Callbacks(client, alertmanager, cache, config)
+        try:
-            client.add_event_callback(callbacks.message, (RoomMessageText,))
+            if config.user_token:
-            client.add_event_callback(
+                # Use token to log in
-                callbacks.invite_event_filtered_callback, (InviteMemberEvent,)
+                client.load_store()
-            )
-            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)
+                # Sync encryption keys with the server
-            while True:
+                if client.should_upload_keys:
+                    await client.keys_upload()
+            else:
+                # Try to login with the configured username/password
                 try:
-                    if config.user_token:
+                    login_response = await client.login(
-                        # Use token to log in
+                        password=config.user_password,
-                        client.load_store()
+                        device_name=config.device_name,
-
-                        # Sync encryption keys with the server
-                        if client.should_upload_keys:
-                            await client.keys_upload()
-                    else:
-                        # Try to login with the configured username/password
-                        try:
-                            login_response = await client.login(
-                                password=config.user_password,
-                                device_name=config.device_name,
-                            )
-
-                            # Check if login failed
-                            if type(login_response) == LoginError:
-                                logger.error(
-                                    "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
-                            # dependencies. In that case, a LocalProtocolError is raised on login.
-                            logger.fatal(
-                                "Failed to login. Have you installed the correct dependencies? "
-                                "https://github.com/poljar/matrix-nio#installation "
-                                "Error: %s",
-                                e,
-                            )
-                            return False
-
-                        # Login succeeded!
-
-                    logger.info(f"Logged in as {config.user_id}")
-                    await client.sync_forever(timeout=30000, full_state=True)
-
-                except (ClientConnectionError, ServerDisconnectedError, TimeoutError):
-                    logger.warning(
-                        "Unable to connect to homeserver, retrying in 15s..."
                     )
 
-                    # Sleep so we don't bombard the server with login requests
+                    # Check if login failed
-                    sleep(15)
+                    if type(login_response) == LoginError:
+                        logger.error("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
+                    # dependencies. In that case, a LocalProtocolError is raised on login.
+                    logger.fatal(
+                        "Failed to login. Have you installed the correct dependencies? "
+                        "https://github.com/poljar/matrix-nio#installation "
+                        "Error: %s",
+                        e,
+                    )
+                    return False
 
+                # Login succeeded!
 
-async def start_webhook_server(cache: Cache, config: Config) -> None:
+            logger.info(f"Logged in as {config.user_id}")
-    async with aiotools.closing_async(create_matrix_client(config)) as client:
+            await client.sync_forever(timeout=30000, full_state=True)
-        async with Webhook(client, cache, config) as webhook_server:
+        except (ClientConnectionError, ServerDisconnectedError, TimeoutError):
-            await webhook_server.start()
+            logger.warning("Unable to connect to homeserver, retrying in 15s...")
+
+            # Sleep so we don't bombard the server with login requests
+            sleep(15)
+        finally:
+            client.close()
 
 
 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