chore: import upstream snapshot with attribution

python/ray/_common/tests/test_retry.py

@@ -0,0 +1,136 @@
+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__]))