Ensure all telemetry is always flushed

ascoderu · Jan 2, 2019 · 8c09f3e · 8c09f3e
1 parent 7055558
commit 8c09f3e
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 16 deletions.
diff --git a/opwen_email_server/utils/log.py b/opwen_email_server/utils/log.py
@@ -24,9 +24,24 @@
 from opwen_email_server.constants.logging import SEPARATOR
 from opwen_email_server.constants.logging import STDERR
 from opwen_email_server.utils.collections import append
+from opwen_email_server.utils.collections import singleton
+
+
+@singleton
+def _create_telemetry_channel() -> Optional[TelemetryChannel]:
+    if not APPINSIGHTS_KEY:
+        return None
+
+    sender = AsynchronousSender()
+    queue = AsynchronousQueue(sender)
+    context = TelemetryContext()
+    context.instrumentation_key = APPINSIGHTS_KEY
+    return TelemetryChannel(context, queue)
 
 
 class LogMixin(object):
+    _telemetry_channel = _create_telemetry_channel()
+
     @cached_property
     def _default_log_handlers(self) -> Iterable[Handler]:
         handlers = []
@@ -35,7 +50,7 @@ def _default_log_handlers(self) -> Iterable[Handler]:
         stderr.setFormatter(Formatter(STDERR))
         handlers.append(stderr)
 
-        if self._telemetry_channel:
+        if APPINSIGHTS_KEY:
             handlers.append(LoggingHandler(
                 APPINSIGHTS_KEY,
                 telemetry_channel=self._telemetry_channel))
@@ -50,20 +65,9 @@ def _logger(self) -> Logger:
         log.setLevel(LOG_LEVEL)
         return log
 
-    @cached_property
-    def _telemetry_channel(self) -> Optional[TelemetryChannel]:
-        if not APPINSIGHTS_KEY:
-            return None
-
-        sender = AsynchronousSender()
-        queue = AsynchronousQueue(sender)
-        context = TelemetryContext()
-        context.instrumentation_key = APPINSIGHTS_KEY
-        return TelemetryChannel(context, queue)
-
     @cached_property
     def _telemetry_client(self) -> Optional[TelemetryClient]:
-        if not self._telemetry_channel:
+        if not APPINSIGHTS_KEY:
             return None
 
         return TelemetryClient(APPINSIGHTS_KEY, self._telemetry_channel)
@@ -86,6 +90,7 @@ def log_exception(self, ex: Exception, message: str, *args: Any):
                 raise ex
             except Exception:
                 self._telemetry_client.track_exception()
+                self._telemetry_channel.flush()
 
     def _log(self, level: int, log_message: str, log_args: Iterable[Any]):
         if not self._logger.isEnabledFor(level):
@@ -103,4 +108,4 @@ def log_event(self, event_name: str, properties: Optional[dict] = None):
 
         if self._telemetry_client:
             self._telemetry_client.track_event(event_name, properties)
-            self._telemetry_client.flush()
+            self._telemetry_channel.flush()
diff --git a/tests/opwen_email_server/test_actions.py b/tests/opwen_email_server/test_actions.py
@@ -12,7 +12,8 @@
 
 class ActionTests(TestCase):
     @patch.object(actions._Action, '_telemetry_client')
-    def test_logs_exception(self, telemetry_mock):
+    @patch.object(actions._Action, '_telemetry_channel')
+    def test_logs_exception(self, mock_channel, mock_client):
         class TestAction(actions._Action):
             def _action(self):
                 int('not-a-number')
@@ -21,7 +22,8 @@ def _action(self):
             action = TestAction()
             action()
 
-        telemetry_mock.track_exception.assert_called_once_with()
+        mock_client.track_exception.assert_called_once_with()
+        mock_channel.flush.assert_called_once_with()
 
 
 class PingTests(TestCase):