devplacepy/tests/unit/services/openai_gateway/reliability.py

# retoor <retoor@molodetz.nl>

import asyncio
import time

from tests.conftest import run_async
from devplacepy.services.openai_gateway.reliability import (
    _retry_after_seconds,
    retry_send,
)


class FakeHeaders_reliability(dict):
    def get(self, key, default=None):
        return super().get(key.lower(), default)


class FakeResp_reliability:
    def __init__(self, status_code, retry_after=None):
        self.status_code = status_code
        headers = FakeHeaders_reliability()
        if retry_after is not None:
            headers["retry-after"] = retry_after
        self.headers = headers

    async def aclose(self):
        pass


def test_retry_after_seconds_parses_numeric_header():
    assert _retry_after_seconds(FakeResp_reliability(429, "2")) == 2.0


def test_retry_after_seconds_absent_returns_none():
    assert _retry_after_seconds(FakeResp_reliability(429)) is None


def test_retry_after_seconds_is_capped():
    assert _retry_after_seconds(FakeResp_reliability(429, "99999")) == 30.0


def test_retry_after_seconds_ignores_garbage():
    assert _retry_after_seconds(FakeResp_reliability(429, "not-a-number-or-date")) is None


def test_retry_send_retries_429_and_honors_retry_after():
    responses = [FakeResp_reliability(429, "0.01"), FakeResp_reliability(200)]
    calls = []

    async def do_call():
        calls.append(time.monotonic())
        return responses.pop(0)

    sem = asyncio.Semaphore(1)
    resp, exc, attempts, queue_wait_ms = run_async(
        retry_send(do_call, sem, max_retries=2, backoff_ms=5000)
    )
    assert exc is None
    assert resp.status_code == 200
    assert attempts == 2
    # The 429 branch waits on the short Retry-After (0.01s), never the
    # much larger fixed 5000ms*attempt linear backoff it would otherwise use.
    assert calls[1] - calls[0] < 1.0


def test_retry_send_gives_up_after_max_retries_on_429():
    async def do_call():
        return FakeResp_reliability(429, "0.001")

    sem = asyncio.Semaphore(1)
    resp, exc, attempts, queue_wait_ms = run_async(
        retry_send(do_call, sem, max_retries=1, backoff_ms=1)
    )
    assert exc is None
    assert resp.status_code == 429
    assert attempts == 2