Merge pull request singer-io#62 from b-ryan/feature/main-critical-wrapper

nick-mccoy · web-flow · commit 452d9eda258e · 2017-12-20T15:57:16.000-05:00
Decorator for wrapping main functions with critical log
diff --git a/singer/utils.py b/singer/utils.py
@@ -169,3 +169,18 @@ def exception_is_4xx(exception):
         return False
 
     return 400 <= exception.response.status_code < 500
+
+
+def handle_top_exception(logger):
+    """A decorator that will catch exceptions and log the exception's message
+    as a CRITICAL log."""
+    def decorator(fn):
+        @functools.wraps(fn)
+        def wrapped(*args, **kwargs):
+            try:
+                return fn(*args, **kwargs)
+            except Exception as exc:
+                logger.critical(exc)
+                raise
+        return wrapped
+    return decorator
diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -1,10 +1,28 @@
 import unittest
 from datetime import datetime as dt
 from datetime import timezone as tz
+import logging
 import singer.utils as u
 
 
 class TestFormat(unittest.TestCase):
     def test_small_years(self):
         self.assertEqual(u.strftime(dt(90, 1, 1, tzinfo=tz.utc), '%04Y-%m-%dT%H:%M:%S.%fZ'),
                          "0090-01-01T00:00:00.000000Z")
+
+
+class TestHandleException(unittest.TestCase):
+    def setUp(self):
+        self.logger = logging.getLogger(__name__)
+
+    def test_successful_fn(self):
+        @u.handle_top_exception(self.logger)
+        def foo():
+            return 3
+        self.assertEqual(foo(), 3)
+
+    def test_exception_fn(self):
+        @u.handle_top_exception(self.logger)
+        def foo():
+            raise RuntimeError("foo")
+        self.assertRaises(RuntimeError, foo)
