From 145bf41c138e7b0f14fb0615c90c2400e8c6b2ab Mon Sep 17 00:00:00 2001 From: Yih-Dar <2521628+ydshieh@users.noreply.github.com> Date: Thu, 2 Feb 2023 17:41:22 +0100 Subject: [PATCH] Allow to add more information in `is_flaky` (#21426) * Allow to add more information * fix style --------- Co-authored-by: ydshieh --- src/transformers/testing_utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/transformers/testing_utils.py b/src/transformers/testing_utils.py index d80536fc2a..3a8ee66218 100644 --- a/src/transformers/testing_utils.py +++ b/src/transformers/testing_utils.py @@ -1690,7 +1690,7 @@ class RequestCounter: return self.old_request(method=method, **kwargs) -def is_flaky(max_attempts: int = 5, wait_before_retry: Optional[float] = None): +def is_flaky(max_attempts: int = 5, wait_before_retry: Optional[float] = None, description: Optional[str] = None): """ To decorate flaky tests. They will be retried on failures. @@ -1699,6 +1699,9 @@ def is_flaky(max_attempts: int = 5, wait_before_retry: Optional[float] = None): The maximum number of attempts to retry the flaky test. wait_before_retry (`float`, *optional*): If provided, will wait that number of seconds before retrying the test. + description (`str`, *optional*): + A string to describe the situation (what / where / why is flaky, link to GH issue/PR comments, errors, + etc.) """ def decorator(test_func_ref):