From 705432c4a576e55061be695078f6419d920901c0 Mon Sep 17 00:00:00 2001
From: Blallo <blallo@autistici.org>
Date: Thu, 25 Aug 2022 13:11:53 +0200
Subject: [PATCH] Improve notifier connection

---
 latecomers/notifier.py | 29 +++++++++++++++++------------
 1 file changed, 17 insertions(+), 12 deletions(-)

diff --git a/latecomers/notifier.py b/latecomers/notifier.py
index 55c84f1..5df1c2a 100644
--- a/latecomers/notifier.py
+++ b/latecomers/notifier.py
@@ -1,4 +1,5 @@
 # -*- encoding: utf-8 -*-
+from contextlib import contextmanager
 from datetime import datetime
 from email import encoders
 from email.mime.base import MIMEBase
@@ -27,14 +28,7 @@ class Notifier(object):
         smtp_pass: T.Text,
         smtp_from: T.Optional[T.Text] = None,
     ) -> None:
-        context = ssl.create_default_context()
-        if smtp_starttls:
-            self._conn = smtplib.SMTP_SSL(smtp_addr, smtp_port)
-            self._greet = self._conn.ehlo
-        else:
-            self._conn = smtplib.SMTP(smtp_addr, smtp_port)
-            self._greet = lambda: self._conn.starttls(context=context)
-
+        self._starttls = smtp_starttls
         self._addr = smtp_addr
         self._port = smtp_port
         self._user = smtp_user
@@ -42,14 +36,25 @@ class Notifier(object):
         self._from = smtp_from if smtp_from else smtp_user
         logger.info("%s initialized", self)
 
+    @contextmanager
+    def connect(self):
+        context = ssl.create_default_context()
+        if self._starttls:
+            with smtplib.SMTP(self._addr, self._port) as conn:
+                conn.starttls(context=context)
+                conn.login(self._user, self._pass)
+                yield conn
+        else:
+            with smtplib.SMTP_SSL(self._addr, self._port) as conn:
+                conn.login(self._user, self._pass)
+                yield conn
+
     def __str__(self) -> T.Text:
         return f"Notifier<{self._addr}:{self._port},from={self._from}>"
 
     def send(self, to: T.List[T.Text], email: T.Text) -> None:
-        self._greet()
-        self._conn.login(self._user, self._pass)
-        self._conn.sendmail(self._from, to, email)
-        self._conn.close()
+        with self.connect() as conn:
+            conn.sendmail(self._from, to, email)
 
     @retry_and_log(logger, RETRIES)
     def send_result(self, to: T.List[T.Text], result: bytes) -> None: