tests/unit/test_retry.py

"""Unit tests for orchestrator/retry.py."""

from __future__ import annotations

import time

import pytest

from orchestrator.retry import RetryPolicy, with_retry_sync


class TestRetryPolicy:
    def test_backoff_increases(self):
        p = RetryPolicy(max_attempts=3, initial_backoff_seconds=0.1, max_backoff_seconds=1.0)
        b1 = p.backoff(1)
        b2 = p.backoff(2)
        b3 = p.backoff(3)
        # Backoff should generally increase (jitter makes exact assertion impossible)
        # max backoff never exceeds cap
        assert b1 <= 0.2
        assert b2 <= 0.4
        assert b3 <= 0.8

    def test_backoff_capped(self):
        p = RetryPolicy(max_attempts=10, initial_backoff_seconds=1.0, max_backoff_seconds=2.0)
        for attempt in range(1, 10):
            assert p.backoff(attempt) <= 2.0


class TestWithRetrySync:
    def test_succeeds_first_try(self):
        calls = [0]
        def fn():
            calls[0] += 1
            return "ok"
        result = with_retry_sync(fn, RetryPolicy(max_attempts=3, initial_backoff_seconds=0.001),
                                 label="test")
        assert result == "ok"
        assert calls[0] == 1

    def test_retries_on_retriable_exception(self):
        calls = [0]
        def fn():
            calls[0] += 1
            if calls[0] < 3:
                raise ConnectionError("transient")
            return "ok"
        result = with_retry_sync(fn, RetryPolicy(max_attempts=3, initial_backoff_seconds=0.001),
                                 label="test")
        assert result == "ok"
        assert calls[0] == 3

    def test_gives_up_after_max_attempts(self):
        calls = [0]
        def fn():
            calls[0] += 1
            raise ConnectionError("persistent")
        with pytest.raises(ConnectionError):
            with_retry_sync(fn, RetryPolicy(max_attempts=3, initial_backoff_seconds=0.001),
                            label="test")
        assert calls[0] == 3

    def test_does_not_retry_non_retriable_exception(self):
        calls = [0]
        def fn():
            calls[0] += 1
            raise ValueError("not retriable")
        with pytest.raises(ValueError):
            with_retry_sync(fn, RetryPolicy(max_attempts=3, initial_backoff_seconds=0.001),
                            label="test")
        # ValueError is not in retriable_exceptions, so no retry
        assert calls[0] == 1

    def test_on_retry_callback_called(self):
        calls = [0]
        retry_calls = []
        def fn():
            calls[0] += 1
            if calls[0] < 3:
                raise ConnectionError("transient")
            return "ok"
        def on_retry(attempt, exc):
            retry_calls.append((attempt, str(exc)))
        result = with_retry_sync(
            fn,
            RetryPolicy(max_attempts=3, initial_backoff_seconds=0.001),
            label="test",
            on_retry=on_retry,
        )
        assert result == "ok"
        assert len(retry_calls) == 2
        assert retry_calls[0][0] == 1
        assert retry_calls[1][0] == 2