wehub-resource-sync
137 lines
3.3 KiB
Python
137 lines
3.3 KiB
Python
import sys
|
|
|
|
import pytest
|
|
|
|
from ray._common.retry import (
|
|
call_with_retry,
|
|
retry,
|
|
)
|
|
|
|
|
|
def test_call_with_retry_immediate_success_with_args():
|
|
def func(a, b):
|
|
return [a, b]
|
|
|
|
assert call_with_retry(func, "func", [], 1, 0, "a", "b") == ["a", "b"]
|
|
|
|
|
|
def test_retry_immediate_success_with_object_args():
|
|
class MyClass:
|
|
@retry("func", [], 1, 0)
|
|
def func(self, a, b):
|
|
return [a, b]
|
|
|
|
assert MyClass().func("a", "b") == ["a", "b"]
|
|
|
|
|
|
@pytest.mark.parametrize("use_decorator", [True, False])
|
|
def test_retry_last_attempt_successful_with_appropriate_wait_time(
|
|
monkeypatch, use_decorator
|
|
):
|
|
sleep_total = 0
|
|
|
|
def sleep(x):
|
|
nonlocal sleep_total
|
|
sleep_total += x
|
|
|
|
monkeypatch.setattr("time.sleep", sleep)
|
|
monkeypatch.setattr("random.uniform", lambda a, b: 1)
|
|
|
|
pattern = "have not reached 4th attempt"
|
|
call_count = 0
|
|
|
|
def func():
|
|
nonlocal call_count
|
|
call_count += 1
|
|
if call_count == 4:
|
|
return "success"
|
|
raise ValueError(pattern)
|
|
|
|
args = ["func", [pattern], 4, 3]
|
|
if use_decorator:
|
|
assert retry(*args)(func)() == "success"
|
|
else:
|
|
assert call_with_retry(func, *args) == "success"
|
|
assert sleep_total == 6 # 1 + 2 + 3
|
|
|
|
|
|
@pytest.mark.parametrize("use_decorator", [True, False])
|
|
def test_retry_unretryable_error(use_decorator):
|
|
call_count = 0
|
|
|
|
def func():
|
|
nonlocal call_count
|
|
call_count += 1
|
|
raise ValueError("unretryable error")
|
|
|
|
args = ["func", ["only retryable error"], 10, 0]
|
|
with pytest.raises(ValueError, match="unretryable error"):
|
|
if use_decorator:
|
|
retry(*args)(func)()
|
|
else:
|
|
call_with_retry(func, *args)
|
|
assert call_count == 1
|
|
|
|
|
|
@pytest.mark.parametrize("use_decorator", [True, False])
|
|
def test_retry_fail_all_attempts_retry_all_errors(use_decorator):
|
|
call_count = 0
|
|
|
|
def func():
|
|
nonlocal call_count
|
|
call_count += 1
|
|
raise ValueError(str(call_count))
|
|
|
|
args = ["func", None, 3, 0]
|
|
with pytest.raises(ValueError):
|
|
if use_decorator:
|
|
retry(*args)(func)()
|
|
else:
|
|
call_with_retry(func, *args)
|
|
assert call_count == 3
|
|
|
|
|
|
def test_call_with_retry_matches_class_name():
|
|
"""Patterns can match the exception class name (e.g., 'RateLimit')."""
|
|
|
|
class RateLimitError(Exception):
|
|
pass
|
|
|
|
call_count = 0
|
|
|
|
def func():
|
|
nonlocal call_count
|
|
call_count += 1
|
|
raise RateLimitError("Error code: 429")
|
|
|
|
with pytest.raises(RateLimitError):
|
|
call_with_retry(func, "func", ["RateLimit"], 3, 0)
|
|
assert call_count == 3
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"pattern,should_retry",
|
|
[
|
|
# Valid regex that is not a literal substring, that matches via regex search
|
|
(r"\d{3}", True),
|
|
# Invalid regex, re.error is handled by returning False and the error is not retried.
|
|
(r"[unclosed", False),
|
|
],
|
|
)
|
|
def test_call_with_retry_regex_matching(pattern, should_retry):
|
|
call_count = 0
|
|
|
|
def func():
|
|
nonlocal call_count
|
|
call_count += 1
|
|
raise ValueError("Error code: 429")
|
|
|
|
with pytest.raises(ValueError):
|
|
call_with_retry(func, "func", [pattern], 3, 0)
|
|
|
|
assert call_count == (3 if should_retry else 1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|