vibe/tests/core/test_retry.py

from __future__ import annotations

from collections.abc import AsyncGenerator

import httpx
import pytest

from vibe.core.utils.retry import (
    _is_retryable_http_error,
    async_generator_retry,
    async_retry,
)


def _make_http_status_error(status_code: int) -> httpx.HTTPStatusError:
    response = httpx.Response(
        status_code=status_code, request=httpx.Request("GET", "https://example.com")
    )
    return httpx.HTTPStatusError(
        message=f"Error {status_code}", request=response.request, response=response
    )


def _make_request(url: str = "https://example.com") -> httpx.Request:
    return httpx.Request("POST", url)


class TestIsRetryableHttpError:
    @pytest.mark.parametrize("code", [408, 409, 425, 429, 500, 502, 503, 504, 529])
    def test_retryable_codes(self, code: int) -> None:
        assert _is_retryable_http_error(_make_http_status_error(code)) is True

    @pytest.mark.parametrize("code", [400, 401, 403, 404, 422])
    def test_non_retryable_codes(self, code: int) -> None:
        assert _is_retryable_http_error(_make_http_status_error(code)) is False

    @pytest.mark.parametrize(
        "exc",
        [
            httpx.ConnectTimeout("connect timed out", request=_make_request()),
            httpx.ReadTimeout("read timed out", request=_make_request()),
            httpx.WriteTimeout("write timed out", request=_make_request()),
            httpx.PoolTimeout("pool timed out", request=_make_request()),
            httpx.ConnectError("connection refused", request=_make_request()),
            httpx.ReadError("read failed", request=_make_request()),
            httpx.WriteError("write failed", request=_make_request()),
            httpx.RemoteProtocolError("server disconnected", request=_make_request()),
        ],
    )
    def test_retryable_network_errors(self, exc: Exception) -> None:
        assert _is_retryable_http_error(exc) is True

    def test_non_retryable_request_error(self) -> None:
        assert _is_retryable_http_error(httpx.InvalidURL("bad url")) is False

    def test_non_http_error_returns_false(self) -> None:
        assert _is_retryable_http_error(ValueError("not http")) is False

    def test_generic_exception_returns_false(self) -> None:
        assert _is_retryable_http_error(RuntimeError("boom")) is False


class TestAsyncRetry:
    @pytest.mark.asyncio
    async def test_retries_network_error_then_succeeds(self) -> None:
        attempts = 0

        @async_retry(tries=3, delay_seconds=0.0, backoff_factor=1.0)
        async def call() -> str:
            nonlocal attempts
            attempts += 1
            if attempts < 2:
                raise httpx.ConnectTimeout("timeout", request=_make_request())
            return "ok"

        result = await call()
        assert result == "ok"
        assert attempts == 2

    @pytest.mark.asyncio
    async def test_does_not_retry_non_retryable(self) -> None:
        attempts = 0

        @async_retry(tries=3, delay_seconds=0.0, backoff_factor=1.0)
        async def call() -> str:
            nonlocal attempts
            attempts += 1
            raise ValueError("nope")

        with pytest.raises(ValueError):
            await call()
        assert attempts == 1

    @pytest.mark.asyncio
    async def test_exhausts_retries(self) -> None:
        attempts = 0

        @async_retry(tries=3, delay_seconds=0.0, backoff_factor=1.0)
        async def call() -> str:
            nonlocal attempts
            attempts += 1
            raise httpx.ReadTimeout("timeout", request=_make_request())

        with pytest.raises(httpx.ReadTimeout):
            await call()
        assert attempts == 3


class TestAsyncGeneratorRetry:
    @pytest.mark.asyncio
    async def test_retries_before_first_yield(self) -> None:
        attempts = 0

        @async_generator_retry(tries=3, delay_seconds=0.0, backoff_factor=1.0)
        async def gen() -> AsyncGenerator[int]:
            nonlocal attempts
            attempts += 1
            if attempts < 2:
                raise httpx.ConnectError("connect failed", request=_make_request())
            yield 1
            yield 2

        items = [item async for item in gen()]
        assert items == [1, 2]
        assert attempts == 2

    @pytest.mark.asyncio
    async def test_does_not_retry_after_first_yield(self) -> None:
        attempts = 0

        @async_generator_retry(tries=3, delay_seconds=0.0, backoff_factor=1.0)
        async def gen() -> AsyncGenerator[int]:
            nonlocal attempts
            attempts += 1
            yield 1
            raise httpx.ReadError("midstream", request=_make_request())

        items: list[int] = []
        with pytest.raises(httpx.ReadError):
            async for item in gen():
                items.append(item)

        assert items == [1]
        assert attempts == 1

    @pytest.mark.asyncio
    async def test_does_not_retry_non_retryable_before_yield(self) -> None:
        attempts = 0

        @async_generator_retry(tries=3, delay_seconds=0.0, backoff_factor=1.0)
        async def gen() -> AsyncGenerator[int]:
            nonlocal attempts
            attempts += 1
            raise ValueError("nope")
            yield 0  # pragma: no cover

        with pytest.raises(ValueError):
            async for _ in gen():
                pass
        assert attempts == 1