Allow to print the stack trace of the last retry error

Author: sergiitk

framework/errors.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import traceback
 from typing import Any
 
 
@@ -56,3 +57,16 @@ def note_blanket_error_info_below(
             f"{cls.note_blanket_error(reason)}"
             f"# Please inspect the information below:\n{info_below}"
         )
+
+
+def format_error_with_trace(error: Exception) -> str:
+    """Returns formatted exception its stack trace."""
+    return "".join(traceback.TracebackException.from_exception(error).format())
+
+
+def format_trace_only(error: Exception) -> str:
+    """Same as format_error_with_trace, but stack trace only."""
+    exception_tb = traceback.TracebackException.from_exception(error)
+    if not exception_tb.stack:
+        return ""
+    return "".join(exception_tb.stack.format())

framework/helpers/retryers.py

@@ -30,6 +30,9 @@
 from tenacity import stop
 from tenacity import wait
 from tenacity.retry import retry_base
+from typing_extensions import Self
+
+import framework.errors
 
 retryers_logger = logging.getLogger(__name__)
 # Type aliases
@@ -229,6 +232,8 @@ class RetryError(tenacity.RetryError):
 
     last_attempt: tenacity.Future
     note: str = ""
+    _message: str = ""
+    _print_last_trace: bool = True
 
     def __init__(
         self,
@@ -238,35 +243,55 @@ def __init__(
         attempts: int = 0,
         check_result: Optional[CheckResultFn] = None,
         note: str = "",
+        print_last_trace: bool = True,
     ):
         last_attempt: tenacity.Future = retry_state.outcome
         super().__init__(last_attempt)
+        self._print_last_trace = print_last_trace
+
+        message = f"Retry error"
 
-        self.message = f"Retry error"
         if retry_state.fn is None:
             # Context manager
-            self.message += f":"
+            message += f":"
         else:
             callback_name = tenacity_utils.get_callback_name(retry_state.fn)
-            self.message += f" calling {callback_name}:"
+            message += f" calling {callback_name}:"
+
         if timeout:
-            self.message += f" timeout {timeout} (h:mm:ss) exceeded"
+            message += f" timeout {timeout} (h:mm:ss) exceeded"
             if attempts:
-                self.message += " or"
+                message += " or"
+
         if attempts:
-            self.message += f" {attempts} attempts exhausted"
+            message += f" {attempts} attempts exhausted"
 
-        self.message += "."
+        message += "."
 
         if last_attempt.failed:
             err = last_attempt.exception()
-            self.message += f" Last exception: {type(err).__name__}: {err}"
+            message += f" Last Retry Attempt Error: {type(err).__name__}"
         elif check_result:
-            self.message += " Check result callback returned False."
+            message += " Check result callback returned False."
+
+        self._message = message
 
         if note:
             self.add_note(note)
 
+    @property
+    def message(self):
+        message = ""
+
+        # TODO(sergiitk): consider if we want to have print-by-default
+        # and/or ignore-by-default exception lists.
+        if cause := self.exception() and self._print_last_trace:
+            cause_trace = framework.errors.format_error_with_trace(cause)
+            message += f"Last Retry Attempt Traceback:\n{cause_trace}\n\n"
+
+        message += self._message
+        return message
+
     def result(self, *, default=None):
         return (
             self.last_attempt.result()
@@ -295,6 +320,10 @@ def reason_str(self):
     def _exception_str(cls, err: Optional[BaseException]) -> str:
         return f"{type(err).__name__}: {err}" if err else "???"
 
+    def with_last_trace(self, enabled: bool = True) -> Self:
+        self._print_last_trace = enabled
+        return self
+
     # TODO(sergiitk): Remove in py3.11, this will be built-in. See PEP 678.
     def add_note(self, note: str):
         self.note = note

framework/test_cases/base_testcase.py

@@ -13,13 +13,14 @@
 # limitations under the License.
 """Base test case used for xds test suites."""
 import inspect
-import traceback
 from typing import Optional, Union
 import unittest
 
 from absl import logging
 from absl.testing import absltest
 
+import framework.errors
+
 
 class BaseTestCase(absltest.TestCase):
     # @override
@@ -135,7 +136,7 @@ def _log_framed_test_failure(
     ) -> None:
         trace: str
         if isinstance(error, Exception):
-            trace = cls._format_error_with_trace(error)
+            trace = framework.errors.format_error_with_trace(error)
         else:
             trace = error
 
@@ -154,9 +155,3 @@ def _log_framed_test_failure(
                 "error": trace,
             },
         )
-
-    @classmethod
-    def _format_error_with_trace(cls, error: Exception) -> str:
-        return "".join(
-            traceback.TracebackException.from_exception(error).format()
-        )