1 files changed, 183 insertions, 0 deletions
diff --git a/catapult/common/py_utils/py_utils/exc_util_unittest.py b/catapult/common/py_utils/py_utils/exc_util_unittest.py
new file mode 100644
index 00000000..31e3b57a
--- /dev/null
+++ b/catapult/common/py_utils/py_utils/exc_util_unittest.py
@@ -0,0 +1,183 @@
+# Copyright 2019 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+import re
+import sys
+import unittest
+
+from py_utils import exc_util
+
+
+class FakeConnectionError(Exception):
+  pass
+
+
+class FakeDisconnectionError(Exception):
+  pass
+
+
+class FakeProcessingError(Exception):
+  pass
+
+
+class FakeCleanupError(Exception):
+  pass
+
+
+class FaultyClient(object):
+  def __init__(self, *args):
+    self.failures = set(args)
+    self.called = set()
+
+  def Connect(self):
+    self.called.add('Connect')
+    if FakeConnectionError in self.failures:
+      raise FakeConnectionError('Oops!')
+
+  def Process(self):
+    self.called.add('Process')
+    if FakeProcessingError in self.failures:
+      raise FakeProcessingError('Oops!')
+
+  @exc_util.BestEffort
+  def Disconnect(self):
+    self.called.add('Disconnect')
+    if FakeDisconnectionError in self.failures:
+      raise FakeDisconnectionError('Oops!')
+
+  @exc_util.BestEffort
+  def Cleanup(self):
+    self.called.add('Cleanup')
+    if FakeCleanupError in self.failures:
+      raise FakeCleanupError('Oops!')
+
+
+class ReraiseTests(unittest.TestCase):
+  def assertLogMatches(self, pattern):
+    self.assertRegexpMatches(
+        sys.stderr.getvalue(), pattern)  # pylint: disable=no-member
+
+  def assertLogNotMatches(self, pattern):
+    self.assertNotRegexpMatches(
+        sys.stderr.getvalue(), pattern)  # pylint: disable=no-member
+
+  def testTryRaisesExceptRaises(self):
+    client = FaultyClient(FakeConnectionError, FakeDisconnectionError)
+
+    # The connection error reaches the top level, while the disconnection
+    # error is logged.
+    with self.assertRaises(FakeConnectionError):
+      try:
+        client.Connect()
+      except:
+        client.Disconnect()
+        raise
+
+    self.assertLogMatches(re.compile(
+        r'While handling a FakeConnectionError, .* was also raised:\n'
+        r'Traceback \(most recent call last\):\n'
+        r'.*\n'
+        r'FakeDisconnectionError: Oops!\n', re.DOTALL))
+    self.assertItemsEqual(client.called, ['Connect', 'Disconnect'])
+
+  def testTryRaisesExceptDoesnt(self):
+    client = FaultyClient(FakeConnectionError)
+
+    # The connection error reaches the top level, disconnecting did not raise
+    # an exception (so nothing is logged).
+    with self.assertRaises(FakeConnectionError):
+      try:
+        client.Connect()
+      except:
+        client.Disconnect()
+        raise
+
+    self.assertLogNotMatches('FakeDisconnectionError')
+    self.assertItemsEqual(client.called, ['Connect', 'Disconnect'])
+
+  def testTryPassesNoException(self):
+    client = FaultyClient(FakeDisconnectionError)
+
+    # If there is no connection error, the except clause is not called (even if
+    # it would have raised an exception).
+    try:
+      client.Connect()
+    except:
+      client.Disconnect()
+      raise
+
+    self.assertLogNotMatches('FakeConnectionError')
+    self.assertLogNotMatches('FakeDisconnectionError')
+    self.assertItemsEqual(client.called, ['Connect'])
+
+  def testTryRaisesFinallyRaises(self):
+    worker = FaultyClient(FakeProcessingError, FakeCleanupError)
+
+    # The processing error reaches the top level, the cleanup error is logged.
+    with self.assertRaises(FakeProcessingError):
+      try:
+        worker.Process()
+      except:
+        raise  # Needed for Cleanup to know if an exception is handled.
+      finally:
+        worker.Cleanup()
+
+    self.assertLogMatches(re.compile(
+        r'While handling a FakeProcessingError, .* was also raised:\n'
+        r'Traceback \(most recent call last\):\n'
+        r'.*\n'
+        r'FakeCleanupError: Oops!\n', re.DOTALL))
+    self.assertItemsEqual(worker.called, ['Process', 'Cleanup'])
+
+  def testTryRaisesFinallyDoesnt(self):
+    worker = FaultyClient(FakeProcessingError)
+
+    # The processing error reaches the top level, the cleanup code runs fine.
+    with self.assertRaises(FakeProcessingError):
+      try:
+        worker.Process()
+      except:
+        raise  # Needed for Cleanup to know if an exception is handled.
+      finally:
+        worker.Cleanup()
+
+    self.assertLogNotMatches('FakeProcessingError')
+    self.assertLogNotMatches('FakeCleanupError')
+    self.assertItemsEqual(worker.called, ['Process', 'Cleanup'])
+
+  def testTryPassesFinallyRaises(self):
+    worker = FaultyClient(FakeCleanupError)
+
+    # The processing code runs fine, the cleanup code raises an exception
+    # which reaches the top level.
+    with self.assertRaises(FakeCleanupError):
+      try:
+        worker.Process()
+      except:
+        raise  # Needed for Cleanup to know if an exception is handled.
+      finally:
+        worker.Cleanup()
+
+    self.assertLogNotMatches('FakeProcessingError')
+    self.assertLogNotMatches('FakeCleanupError')
+    self.assertItemsEqual(worker.called, ['Process', 'Cleanup'])
+
+  def testTryRaisesExceptRaisesFinallyRaises(self):
+    worker = FaultyClient(
+        FakeProcessingError, FakeDisconnectionError, FakeCleanupError)
+
+    # Chaining try-except-finally works fine. Only the processing error reaches
+    # the top level; the other two are logged.
+    with self.assertRaises(FakeProcessingError):
+      try:
+        worker.Process()
+      except:
+        worker.Disconnect()
+        raise
+      finally:
+        worker.Cleanup()
+
+    self.assertLogMatches('FakeDisconnectionError')
+    self.assertLogMatches('FakeCleanupError')
+    self.assertItemsEqual(worker.called, ['Process', 'Disconnect', 'Cleanup'])