From 498a60c1d5e45e0406a07322d7e87a22a5caf45a Mon Sep 17 00:00:00 2001
From: Ross Williams <rossjrw@gmail.com>
Date: Thu, 9 May 2024 20:09:43 +0100
Subject: [PATCH] Move log dump catch to context manager

---
 notifier/notify.py | 43 +++++++++++++++++++++++++++----------------
 1 file changed, 27 insertions(+), 16 deletions(-)

diff --git a/notifier/notify.py b/notifier/notify.py
index 79154db..9e806ce 100644
--- a/notifier/notify.py
+++ b/notifier/notify.py
@@ -1,7 +1,8 @@
+from contextlib import contextmanager
 import logging
 import re
 from smtplib import SMTPAuthenticationError
-from typing import FrozenSet, Iterable, List, Optional, Set, Tuple
+from typing import FrozenSet, Iterable, Iterator, List, Optional, Set, Tuple
 
 from notifier.config.remote import get_global_config
 from notifier.config.user import get_user_config
@@ -70,6 +71,28 @@ def pick_channels_to_notify(
     return channels
 
 
+@contextmanager
+def activation_log_dump_context(
+    config: LocalConfig, database: BaseDatabaseDriver, dry_run: bool
+) -> Iterator[LogDumpCacher]:
+    """Creates a log dump context that ends the long if the wrapped process fails."""
+    activation_log_dump = LogDumpCacher[ActivationLogDump](
+        {"start_timestamp": timestamp()},
+        database.store_activation_log_dump,
+        dry_run,
+    )
+    try:
+        yield activation_log_dump
+
+    finally:
+        # Even if the run failed, record the end timestamp and upload if possible
+        activation_log_dump.update({"end_timestamp": timestamp()})
+
+        if not dry_run:
+            logger.info("Uploading log dumps...")
+            record_activation_log(config, database)
+
+
 def notify(
     *,
     config: LocalConfig,
@@ -88,13 +111,9 @@ def notify(
     schedules.
     """
 
-    activation_log_dump = LogDumpCacher[ActivationLogDump](
-        {"start_timestamp": timestamp()},
-        database.store_activation_log_dump,
-        dry_run,
-    )
-
-    try:
+    with activation_log_dump_context(
+        config, database, dry_run
+    ) as activation_log_dump:
         # If there are no active channels, which shouldn't happen, there is
         # nothing to do
         if len(active_channels) == 0:
@@ -169,14 +188,6 @@ def notify(
         delete_prepared_invalid_user_pages(config, connection)
         rename_invalid_user_config_pages(config, connection)
 
-    finally:
-        # Even if the run failed, record the end timestamp and upload if possible
-        activation_log_dump.update({"end_timestamp": timestamp()})
-
-        assert not dry_run
-        logger.info("Uploading log dumps...")
-        record_activation_log(config, database)
-
 
 def notify_active_channels(
     active_channels: Iterable[str],