[Framework] Allow to print the stack trace of the last retry error

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,19 @@ 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, only_if_tb_present=False) -> str:
+    """Returns formatted exception its stack trace."""
+    exception_tb = traceback.TracebackException.from_exception(error)
+    if only_if_tb_present and not exception_tb.stack:
+        return ""
+    return "".join(exception_tb.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,58 @@ 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__}({err})"
         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):
+        # TODO(sergiitk): consider if we want to have print-by-default
+        # and/or ignore-by-default exception lists.
+        tb_out = ""
+        if self._print_last_trace and (cause := self.exception()):
+            try:
+                if last_tb := framework.errors.format_error_with_trace(
+                    cause, only_if_tb_present=True
+                ):
+                    tb_out = f"^\nLast Retry Attempt Error {last_tb.rstrip()}"
+            except Exception as err_tb_format:
+                tb_out = f"<Error printing the traceback: {err_tb_format!r}>"
+
+        return f"{self._message}\n{tb_out}" if tb_out else self._message
+
     def result(self, *, default=None):
         return (
             self.last_attempt.result()
@@ -295,6 +323,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()
-        )

tests/unit/helpers/retryers_test.py

@@ -0,0 +1,124 @@
+# Copyright 2025 gRPC authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 datetime as dt
+from unittest import mock
+
+from absl.testing import absltest
+
+from framework.helpers import retryers
+
+
+class RetryErrorTest(absltest.TestCase):
+    def test_construct(self):
+        state_mock = mock.MagicMock()
+        state_mock.fn = lambda val: ...
+        state_mock.outcome.exception.return_value = OSError("my err")
+
+        retry_err = retryers.RetryError(state_mock, attempts=42)
+        want_msg = (
+            "Retry error calling"
+            " __main__.RetryErrorTest.test_construct.<locals>.<lambda>:"
+            " 42 attempts exhausted. Last Retry Attempt Error: OSError(my err)"
+        )
+        with self.assertRaisesWithLiteralMatch(retryers.RetryError, want_msg):
+            raise retry_err
+
+        self.assertNotEmpty(str(retry_err))
+        self.assertNotEmpty(repr(retry_err))
+
+    def test_retry_with_trace(self):
+        retryer = retryers.constant_retryer(
+            wait_fixed=dt.timedelta(microseconds=1),
+            timeout=dt.timedelta(seconds=1),
+            attempts=2,
+        )
+
+        errors = (ValueError, TypeError, AttributeError)
+        error_generator = self._error_generator(errors)
+
+        with self.assertRaises(retryers.RetryError) as capture:
+            retryer(self._raise_error, error_generator)
+
+        retry_err = capture.exception
+        retry_err_str = str(retry_err)
+        retry_err_lines = retry_err_str.split("\n")
+        self.assertEqual(
+            retry_err_lines[0],
+            "Retry error calling __main__.RetryErrorTest._raise_error:"
+            " timeout 0:00:01 (h:mm:ss) exceeded or 2 attempts exhausted."
+            " Last Retry Attempt Error:"
+            " TypeError(error msg test_retry_with_trace)",
+        )
+        self.assertContainsInOrder(
+            (
+                # First line see the exact match above
+                "Retry error calling",
+                "Last Retry Attempt Error: TypeError(error msg test_retry_with_trace)\n",
+                # our custom stack trace header
+                "\n^\nLast Retry Attempt Error Traceback (most recent call last):",
+                # ...
+                #   File "tests/unit/helpers/retryers_test.py", line 121, in _raise_error
+                #     raise err(err_msg)
+                # TypeError: error msg test_retry_with_trace'
+                'retryers_test.py", line',
+                ", in _raise_error\n",
+                "  raise err(err_msg)\n",
+                "TypeError: error msg test_retry_with_trace",
+            ),
+            retry_err_str,
+        )
+        self.assertEqual(
+            retry_err_lines[-1],
+            "TypeError: error msg test_retry_with_trace",
+        )
+
+    def test_retry_without_trace(self):
+        retryer = retryers.constant_retryer(
+            wait_fixed=dt.timedelta(microseconds=1),
+            timeout=dt.timedelta(seconds=1),
+            attempts=3,
+        )
+
+        will_throw = (ValueError, TypeError, AttributeError)
+        error_generator = self._error_generator(will_throw)
+
+        with self.assertRaises(retryers.RetryError) as capture:
+            # this will be a general pattern to disable last err trace.
+            try:
+                retryer(self._raise_error, error_generator)
+            except retryers.RetryError as retry_err:
+                retry_err.with_last_trace(False)
+                raise
+
+        # just a single line with no tracebacks
+        self.assertEqual(
+            str(capture.exception),
+            "Retry error calling __main__.RetryErrorTest._raise_error:"
+            " timeout 0:00:01 (h:mm:ss) exceeded or 3 attempts exhausted."
+            " Last Retry Attempt Error:"
+            " AttributeError(error msg test_retry_without_trace)",
+        )
+
+    def _raise_error(self, generator, msg: str = "error msg"):
+        for err in generator:
+            err_msg = f"{msg} {self.id().split('.')[-1]}"
+            raise err(err_msg)
+
+    def _error_generator(self, errors):
+        for err in errors:
+            yield err
+
+
+if __name__ == "__main__":
+    absltest.main()