backend/tests/test_openai_rate_limit.py

"""OpenAI client concurrency gate and 429 retry behaviour."""

from __future__ import annotations

import asyncio
from unittest.mock import MagicMock, patch

import pytest

from backend.config import settings
from backend.llm import openai_client


class _RateLimitError(Exception):
    status_code = 429


def test_max_concurrent_llm_calls_respects_settings(monkeypatch):
    monkeypatch.setattr(settings, "max_concurrent_llm_calls", 7)
    openai_client.reset_client()
    assert openai_client.max_concurrent_llm_calls() == 7


def test_chat_text_retries_on_429(monkeypatch):
    monkeypatch.setattr(settings, "openai_rate_limit_max_retries", 2)
    monkeypatch.setattr(settings, "openai_rate_limit_backoff_base_seconds", 0.01)
    monkeypatch.setattr(settings, "openai_rate_limit_backoff_max_seconds", 0.05)
    openai_client.reset_client()

    mock_client = MagicMock()
    mock_client.chat.completions.create.side_effect = [
        _RateLimitError("rate limited"),
        MagicMock(choices=[MagicMock(message=MagicMock(content="ok"))]),
    ]

    sleeps: list[float] = []
    monkeypatch.setattr(openai_client.time, "sleep", lambda s: sleeps.append(s))

    with patch.object(openai_client, "_get_client", return_value=mock_client), patch.object(
        openai_client, "_is_rate_limit_error", return_value=True
    ):
        out = openai_client.chat_text([{"role": "user", "content": "hi"}])

    assert out == "ok"
    assert mock_client.chat.completions.create.call_count == 2
    assert sleeps


def test_chat_text_raises_after_retry_budget_exhausted(monkeypatch):
    monkeypatch.setattr(settings, "openai_rate_limit_max_retries", 1)
    monkeypatch.setattr(settings, "openai_rate_limit_backoff_base_seconds", 0.01)
    openai_client.reset_client()

    mock_client = MagicMock()
    mock_client.chat.completions.create.side_effect = _RateLimitError("rate limited")
    monkeypatch.setattr(openai_client.time, "sleep", lambda _s: None)

    with patch.object(openai_client, "_get_client", return_value=mock_client), patch.object(
        openai_client, "_is_rate_limit_error", return_value=True
    ):
        with pytest.raises(_RateLimitError):
            openai_client.chat_text([{"role": "user", "content": "hi"}])

    assert mock_client.chat.completions.create.call_count == 2


def test_thread_semaphore_capacity(monkeypatch):
    monkeypatch.setattr(settings, "max_concurrent_llm_calls", 2)
    openai_client.reset_client()

    sem = openai_client._get_thread_semaphore()
    assert sem.acquire(blocking=False) is True
    assert sem.acquire(blocking=False) is True
    assert sem.acquire(blocking=False) is False
    sem.release()
    sem.release()


@pytest.mark.asyncio
async def test_chat_json_async_sets_gate_held(monkeypatch):
    monkeypatch.setattr(settings, "max_concurrent_llm_calls", 3)
    openai_client.reset_client()

    held: list[bool] = []

    def fake_chat_json(*args, **kwargs):
        held.append(bool(kwargs.get("_gate_held")))
        return {"passed": True}

    monkeypatch.setattr(openai_client, "chat_json", fake_chat_json)
    out = await openai_client.chat_json_async([{"role": "user", "content": "audit"}])
    assert out == {"passed": True}
    assert held == [True]