tests/test_llm.py

# tests/test_llm.py
import json
from unittest.mock import AsyncMock, MagicMock, patch

import httpx
import pytest

from app.core.llm import _log_retry, _wait_for_rate_limit, call_llm, call_llm_for_json


def _ok_response(content: str = "hello", tokens: int = 0) -> MagicMock:
    resp = MagicMock(spec=httpx.Response)
    resp.status_code = 200
    resp.headers = {}
    body: dict = {"choices": [{"message": {"content": content}}]}
    if tokens:
        body["usage"] = {"total_tokens": tokens}
    resp.json.return_value = body
    resp.request = MagicMock()
    return resp


# ── _log_retry ────────────────────────────────────────────────────────────────


def test_log_retry_with_exception() -> None:
    rs = MagicMock()
    rs.attempt_number = 2
    rs.outcome = MagicMock()
    rs.outcome.exception.return_value = ValueError("boom")
    _log_retry(rs)  # must not raise


def test_log_retry_no_outcome() -> None:
    rs = MagicMock()
    rs.attempt_number = 1
    rs.outcome = None
    _log_retry(rs)  # must not raise


# ── _wait_for_rate_limit ──────────────────────────────────────────────────────


@pytest.mark.asyncio
async def test_wait_for_rate_limit_sleeps_when_header_present() -> None:
    resp = MagicMock()
    resp.headers = {"Retry-After": "0"}
    exc = httpx.HTTPStatusError("rate limited", request=MagicMock(), response=resp)
    await _wait_for_rate_limit(exc)  # must not raise


@pytest.mark.asyncio
async def test_wait_for_rate_limit_noop_when_no_header() -> None:
    resp = MagicMock()
    resp.headers = {}
    exc = httpx.HTTPStatusError("rate limited", request=MagicMock(), response=resp)
    await _wait_for_rate_limit(exc)  # must not raise


# ── call_llm ──────────────────────────────────────────────────────────────────


@pytest.mark.asyncio
async def test_call_llm_returns_content_string() -> None:
    mock_client = AsyncMock()
    mock_client.post = AsyncMock(return_value=_ok_response("the answer"))

    with patch("app.core.llm.httpx.AsyncClient") as MockClient:
        MockClient.return_value.__aenter__ = AsyncMock(return_value=mock_client)
        MockClient.return_value.__aexit__ = AsyncMock(return_value=False)
        result = await call_llm("test prompt")

    assert result == "the answer"


@pytest.mark.asyncio
async def test_call_llm_includes_system_prompt() -> None:
    mock_client = AsyncMock()
    captured: dict = {}

    async def capture(url: str, headers: dict, json: dict) -> MagicMock:
        captured.update(json)
        return _ok_response("ok")

    mock_client.post = capture

    with patch("app.core.llm.httpx.AsyncClient") as MockClient:
        MockClient.return_value.__aenter__ = AsyncMock(return_value=mock_client)
        MockClient.return_value.__aexit__ = AsyncMock(return_value=False)
        await call_llm("user msg", system_prompt="be concise")

    assert captured["messages"][0] == {"role": "system", "content": "be concise"}
    assert captured["messages"][1]["role"] == "user"


@pytest.mark.asyncio
async def test_call_llm_records_tokens_when_present() -> None:
    mock_client = AsyncMock()
    mock_client.post = AsyncMock(return_value=_ok_response("ok", tokens=99))
    mock_metrics = MagicMock()

    with (
        patch("app.core.llm.httpx.AsyncClient") as MockClient,
        patch("app.core.llm.metrics", mock_metrics),
    ):
        MockClient.return_value.__aenter__ = AsyncMock(return_value=mock_client)
        MockClient.return_value.__aexit__ = AsyncMock(return_value=False)
        await call_llm("prompt")

    mock_metrics.record_tokens.assert_called_once_with("llm_total", 99)


@pytest.mark.asyncio
async def test_call_llm_skips_token_recording_when_zero() -> None:
    mock_client = AsyncMock()
    mock_client.post = AsyncMock(return_value=_ok_response("ok", tokens=0))
    mock_metrics = MagicMock()

    with (
        patch("app.core.llm.httpx.AsyncClient") as MockClient,
        patch("app.core.llm.metrics", mock_metrics),
    ):
        MockClient.return_value.__aenter__ = AsyncMock(return_value=mock_client)
        MockClient.return_value.__aexit__ = AsyncMock(return_value=False)
        await call_llm("prompt")

    mock_metrics.record_tokens.assert_not_called()


# ── call_llm_for_json ─────────────────────────────────────────────────────────


@pytest.mark.asyncio
async def test_call_llm_for_json_parses_clean_json() -> None:
    payload = {"result": "ok", "score": 7}
    with patch(
        "app.core.llm.call_llm_routed", new=AsyncMock(return_value=json.dumps(payload))
    ):
        result = await call_llm_for_json("prompt")
    assert result == payload


@pytest.mark.asyncio
async def test_call_llm_for_json_strips_markdown_fences() -> None:
    payload = {"a": 1}
    raw = f"```json\n{json.dumps(payload)}\n```"
    with patch("app.core.llm.call_llm_routed", new=AsyncMock(return_value=raw)):
        result = await call_llm_for_json("prompt")
    assert result == payload


@pytest.mark.asyncio
async def test_call_llm_for_json_corrective_prompt_on_bad_json() -> None:
    good = {"fixed": True}
    with (
        patch("app.core.llm._call_llm_with_model", new=AsyncMock(return_value="not json")),
        patch("app.core.llm.call_llm", new=AsyncMock(return_value=json.dumps(good))),
    ):
        result = await call_llm_for_json("prompt")
    assert result == good


@pytest.mark.asyncio
async def test_call_llm_for_json_returns_empty_dict_on_total_failure() -> None:
    with (
        patch("app.core.llm.call_llm_routed", new=AsyncMock(return_value="still not json")),
        patch("app.core.llm.call_llm", new=AsyncMock(return_value="also not json")),
    ):
        result = await call_llm_for_json("prompt")
    assert result == {}