ForcePilot/backend/test/unit/channels/sdk/test_retry.py

from __future__ import annotations

import asyncio

import pytest

from yuxi.channels.sdk.retry import RetryConfig, with_retry


class TestRetryConfig:
    def test_default_values(self):
        cfg = RetryConfig()
        assert cfg.max_retries == 3
        assert cfg.base_delay == 1.0
        assert cfg.max_delay == 60.0
        assert cfg.backoff_factor == 2.0
        assert cfg.jitter is True
        assert cfg.retryable_exceptions == (Exception,)

    def test_custom_values(self):
        cfg = RetryConfig(
            max_retries=5,
            base_delay=2.0,
            max_delay=30.0,
            backoff_factor=3.0,
            jitter=False,
            retryable_exceptions=(ValueError, TypeError),
        )
        assert cfg.max_retries == 5
        assert cfg.base_delay == 2.0
        assert cfg.max_delay == 30.0
        assert cfg.backoff_factor == 3.0
        assert cfg.jitter is False
        assert cfg.retryable_exceptions == (ValueError, TypeError)


class TestWithRetry:
    @pytest.mark.asyncio
    async def test_success_first_attempt(self):
        call_count = 0

        async def succeed():
            nonlocal call_count
            call_count += 1
            return "ok"

        result = await with_retry(succeed)
        assert result == "ok"
        assert call_count == 1

    @pytest.mark.asyncio
    async def test_retry_on_failure_then_succeed(self):
        call_count = 0

        async def fail_then_succeed():
            nonlocal call_count
            call_count += 1
            if call_count < 3:
                raise RuntimeError("transient error")
            return "recovered"

        cfg = RetryConfig(max_retries=3, jitter=False)
        result = await with_retry(fail_then_succeed, cfg)
        assert result == "recovered"
        assert call_count == 3

    @pytest.mark.asyncio
    async def test_exhausts_retries(self):
        call_count = 0

        async def always_fail():
            nonlocal call_count
            call_count += 1
            raise RuntimeError("persistent error")

        cfg = RetryConfig(max_retries=2, jitter=False)
        with pytest.raises(RuntimeError, match="persistent error"):
            await with_retry(always_fail, cfg)
        assert call_count == 3  # initial + 2 retries

    @pytest.mark.asyncio
    async def test_non_retryable_exception_does_not_retry(self):
        call_count = 0

        async def raise_value_error():
            nonlocal call_count
            call_count += 1
            raise ValueError("not retryable")

        cfg = RetryConfig(
            max_retries=3,
            jitter=False,
            retryable_exceptions=(RuntimeError,),
        )
        with pytest.raises(ValueError, match="not retryable"):
            await with_retry(raise_value_error, cfg)
        assert call_count == 1

    @pytest.mark.asyncio
    async def test_retryable_subclass_exception(self):
        call_count = 0

        class CustomRetryableError(RuntimeError):
            pass

        async def raise_custom():
            nonlocal call_count
            call_count += 1
            if call_count < 2:
                raise CustomRetryableError("custom")
            return "ok"

        cfg = RetryConfig(
            max_retries=2,
            jitter=False,
            retryable_exceptions=(RuntimeError,),
        )
        result = await with_retry(raise_custom, cfg)
        assert result == "ok"
        assert call_count == 2

    @pytest.mark.asyncio
    async def test_zero_retries(self):
        call_count = 0

        async def fail():
            nonlocal call_count
            call_count += 1
            raise RuntimeError("fail")

        cfg = RetryConfig(max_retries=0, jitter=False)
        with pytest.raises(RuntimeError, match="fail"):
            await with_retry(fail, cfg)
        assert call_count == 1

    @pytest.mark.asyncio
    async def test_backoff_delay_increases(self):
        delays: list[float] = []

        async def fail():
            raise RuntimeError("fail")

        original_sleep = asyncio.sleep

        async def mock_sleep(delay):
            delays.append(delay)

        asyncio.sleep = mock_sleep
        try:
            cfg = RetryConfig(max_retries=3, base_delay=1.0, backoff_factor=2.0, jitter=False)
            with pytest.raises(RuntimeError):
                await with_retry(fail, cfg)
        finally:
            asyncio.sleep = original_sleep

        assert len(delays) == 3
        assert delays[0] == pytest.approx(1.0)
        assert delays[1] == pytest.approx(2.0)
        assert delays[2] == pytest.approx(4.0)

    @pytest.mark.asyncio
    async def test_max_delay_cap(self):
        delays: list[float] = []

        async def fail():
            raise RuntimeError("fail")

        original_sleep = asyncio.sleep

        async def mock_sleep(delay):
            delays.append(delay)

        asyncio.sleep = mock_sleep
        try:
            cfg = RetryConfig(
                max_retries=5,
                base_delay=10.0,
                backoff_factor=2.0,
                max_delay=15.0,
                jitter=False,
            )
            with pytest.raises(RuntimeError):
                await with_retry(fail, cfg)
        finally:
            asyncio.sleep = original_sleep

        for d in delays:
            assert d <= 15.0

    @pytest.mark.asyncio
    async def test_jitter_produces_variable_delays(self):
        delays: list[float] = []

        async def fail():
            raise RuntimeError("fail")

        original_sleep = asyncio.sleep

        async def mock_sleep(delay):
            delays.append(delay)

        asyncio.sleep = mock_sleep
        try:
            cfg = RetryConfig(max_retries=10, base_delay=1.0, jitter=True)
            with pytest.raises(RuntimeError):
                await with_retry(fail, cfg)
        finally:
            asyncio.sleep = original_sleep

        assert len(set(round(d, 4) for d in delays)) > 1