recipe_modules/utils/api.py

# Copyright 2020 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Helper functions for common boilerplate tasks.

Top-level functions stored in __init__.py can be imported and used directly
(without needing a `RecipeApi` object) using `RECIPE_MODULES` imports. E.g.:

  from RECIPE_MODULES.fuchsia.utils import pluralize
"""

import traceback

from recipe_engine import recipe_api


class UtilsApi(recipe_api.RecipeApi):
    def retry(self, func, max_attempts, sleep=5.0, backoff_factor=1.5):
        """Retry the given function with exponential backoff.

        Args:
            func (callable): A function that performs the action that should
                be retried on failure. If it raises a `StepFailure` it will
                be retried. Any other exception will end the retry loop and
                bubble up.
            max_attempts (int): How many times to try before giving up.
            sleep (int or float): The time to sleep after the first attempt.
            backoff_factor (int or float): The factor by which the sleep time
                will be multiplied after each attempt.
        """
        for attempt in range(max_attempts):
            try:
                return func()
            except self.m.step.StepFailure as ex:
                # We shouldn't retry if the step was cancelled because the build
                # is likely being shut down.
                #
                # As of 2021-12-22 there's no way to simulate step cancellation
                # in tests, so disable coverage.
                if ex.was_cancelled:  # pragma: no cover
                    raise
                if attempt == max_attempts - 1:
                    raise

            self.m.time.sleep(sleep)
            sleep *= backoff_factor

    def traceback_format_exc(self):
        """Simplify expectation file representation of exceptions.

        If in recipe unit testing, reduce the output of the exception to just
        the first line, cutting off the stack trace. This means expectation
        files don't change just because an unrelated line was added at another
        point in a source file. This is a no-op in production.

        Must be called from within an "except:" block.
        """

        text = traceback.format_exc()
        if self._test_data.enabled:
            # Prune the traceback so that the checked-in test expectations
            # don't contain source code line numbers, which often change.
            lines = text.splitlines()
            text = (
                "%s\n... <traceback text omitted to simplify test expectations>\n%s"
                % (lines[0], lines[-1])
            )
        return text