Skip to content

Commit

Permalink
misc: accept list as pattern for retry_command
Browse files Browse the repository at this point in the history
  • Loading branch information
pbrezina committed Sep 20, 2024
1 parent a5cda6d commit 58b5740
Showing 1 changed file with 16 additions and 6 deletions.
22 changes: 16 additions & 6 deletions sssd_test_framework/misc/ssh.py
Original file line number Diff line number Diff line change
Expand Up @@ -98,8 +98,8 @@ def __exit__(self, exception_type, exception_value, traceback) -> None:
def retry_command(
max_retries: int = 5,
delay: float = 1,
match_stdout: str | None = None,
match_stderr: str | None = None,
match_stdout: str | list[str] | None = None,
match_stderr: str | list[str] | None = None,
) -> Callable[[Callable[Param, ProcessResult]], Callable[Param, ProcessResult]]:
"""
Decorated function will be retried if its return code is non zero.
Expand All @@ -109,15 +109,25 @@ def retry_command(
:param delay: Delay in seconds between each retry, defaults to 1
:type delay: float, optional
:param match_stdout: If set, retry only of string is found in stdout, defaults to None
:type match_stdout: str | None, optional
:type match_stdout: str | list[str] | None, optional
:param match_stderr: If set, retry only of string is found in stderr, defaults to None
:type match_stderr: str | None, optional
:type match_stderr: str | list[str] | None, optional
:return: Decorated function.
:rtype: Callable
"""

def decorator(func: Callable[Param, ProcessResult]) -> Callable[Param, ProcessResult]:
def _match(pattern: str | list[str], where: str) -> bool:
if isinstance(pattern, str):
pattern = [pattern]

for p in pattern:
if p in where:
return True

return False

@wraps(func)
def wrapper(*args, **kwargs) -> ProcessResult:
error: ProcessError | None = None
Expand All @@ -144,10 +154,10 @@ def wrapper(*args, **kwargs) -> ProcessResult:
if rc == 0:
break

if match_stdout is not None and match_stdout not in stdout:
if match_stdout is not None and not _match(match_stdout, stdout):
break

if match_stderr is not None and match_stderr not in stderr:
if match_stderr is not None and not _match(match_stderr, stderr):
break

retry += 1
Expand Down

0 comments on commit 58b5740

Please sign in to comment.