tests/test_llm_interface.py

import pytest
from unittest.mock import AsyncMock, MagicMock
from ankigen.llm_interface import OpenAIClientManager, OpenAIRateLimiter
from openai import OpenAIError

# --- OpenAIClientManager Tests ---


@pytest.mark.anyio
async def test_initialize_client_valid_key(mocker):
    manager = OpenAIClientManager()
    # Mock AsyncOpenAI class
    mock_openai_class = mocker.patch("ankigen.llm_interface.AsyncOpenAI", autospec=True)

    await manager.initialize_client("sk-valid-key")

    assert manager._api_key == "sk-valid-key"
    assert manager._client is not None
    mock_openai_class.assert_called_once_with(api_key="sk-valid-key")


@pytest.mark.anyio
async def test_initialize_client_invalid_key():
    manager = OpenAIClientManager()
    # Key not starting with sk-
    with pytest.raises(ValueError, match="Invalid OpenAI API key format."):
        await manager.initialize_client("invalid-key")
    assert manager._client is None


@pytest.mark.anyio
async def test_initialize_client_empty_key():
    manager = OpenAIClientManager()
    with pytest.raises(ValueError, match="Invalid OpenAI API key format."):
        await manager.initialize_client("")
    assert manager._client is None


@pytest.mark.anyio
async def test_initialize_client_openai_error(mocker):
    manager = OpenAIClientManager()
    mocker.patch(
        "ankigen.llm_interface.AsyncOpenAI", side_effect=OpenAIError("API Error")
    )

    with pytest.raises(OpenAIError):
        await manager.initialize_client("sk-error")

    assert manager._client is None


@pytest.mark.anyio
async def test_initialize_client_unexpected_error(mocker):
    manager = OpenAIClientManager()
    mocker.patch(
        "ankigen.llm_interface.AsyncOpenAI", side_effect=Exception("Unexpected")
    )

    with pytest.raises(
        RuntimeError, match="Unexpected error initializing AsyncOpenAI client."
    ):
        await manager.initialize_client("sk-unexpected")

    assert manager._client is None


def test_get_client_not_initialized():
    manager = OpenAIClientManager()
    with pytest.raises(RuntimeError, match="AsyncOpenAI client is not initialized"):
        manager.get_client()


@pytest.mark.anyio
async def test_get_client_initialized(mocker):
    manager = OpenAIClientManager()
    mocker.patch("ankigen.llm_interface.AsyncOpenAI")
    await manager.initialize_client("sk-valid")

    client = manager.get_client()
    assert client is not None


def test_context_manager_sync(mocker):
    manager = OpenAIClientManager()
    mock_close = mocker.patch.object(manager, "close")

    with manager as m:
        assert m is manager

    mock_close.assert_called_once()


@pytest.mark.anyio
async def test_context_manager_async(mocker):
    manager = OpenAIClientManager()
    mock_aclose = mocker.patch.object(manager, "aclose", new_callable=AsyncMock)

    async with manager as m:
        assert m is manager

    mock_aclose.assert_called_once()


def test_close_method(mocker):
    manager = OpenAIClientManager()
    mock_client = MagicMock()
    # Ensure it has 'close' method
    mock_client.close = MagicMock()
    manager._client = mock_client

    manager.close()

    mock_client.close.assert_called_once()
    assert manager._client is None


@pytest.mark.anyio
async def test_aclose_method(mocker):
    manager = OpenAIClientManager()
    mock_client = AsyncMock()
    # Ensure it has 'aclose' method
    mock_client.aclose = AsyncMock()
    manager._client = mock_client

    await manager.aclose()

    mock_client.aclose.assert_called_once()
    assert manager._client is None


@pytest.mark.anyio
async def test_aclose_method_fallback_to_sync(mocker):
    manager = OpenAIClientManager()
    # Create a client mock that only exposes 'close', so 'aclose' is truly absent
    mock_client = MagicMock(spec_set=["close"])
    mock_client.close = MagicMock()
    manager._client = mock_client

    await manager.aclose()

    mock_client.close.assert_called_once()
    assert manager._client is None


# --- OpenAIRateLimiter Tests ---


@pytest.mark.anyio
async def test_rate_limiter_within_limits():
    limiter = OpenAIRateLimiter(tokens_per_minute=100)
    # Should not wait
    await limiter.wait_if_needed(50)
    assert limiter.tokens_used_current_window == 50

    await limiter.wait_if_needed(40)
    assert limiter.tokens_used_current_window == 90


@pytest.mark.anyio
async def test_rate_limiter_window_reset(mocker):
    # Mock monotonic to simulate time passing
    mock_monotonic = mocker.patch("time.monotonic")
    mock_monotonic.return_value = 0.0

    limiter = OpenAIRateLimiter(tokens_per_minute=100)
    await limiter.wait_if_needed(80)
    assert limiter.tokens_used_current_window == 80

    # Advance time by 61 seconds
    mock_monotonic.return_value = 61.0

    await limiter.wait_if_needed(10)

    # Should have reset
    assert limiter.tokens_used_current_window == 10
    assert limiter.current_window_start_time == 61.0


@pytest.mark.anyio
async def test_rate_limiter_wait_needed(mocker):
    mock_sleep = mocker.patch("asyncio.sleep", new_callable=AsyncMock)
    mock_monotonic = mocker.patch("time.monotonic")

    # Start at time 0
    mock_monotonic.return_value = 0.0
    limiter = OpenAIRateLimiter(tokens_per_minute=100)

    # Use 90 tokens
    await limiter.wait_if_needed(90)
    assert limiter.tokens_used_current_window == 90

    # Next token request of 20 would exceed 100 (90 + 20 = 110)
    # Current time is 0.0. window_start is 0.0.
    # time_to_wait = (0.0 + 60.0) - 0.0 = 60.0

    # Instead of side_effect list, we use a more robust approach
    # asyncio loop also calls time.monotonic(), so return_value is safer
    # we update it when sleep is called.
    async def side_effect_sleep(delay):
        mock_monotonic.return_value = 60.0  # Fast forward time

    mock_sleep.side_effect = side_effect_sleep

    await limiter.wait_if_needed(20)

    mock_sleep.assert_called_once_with(60.0)
    assert limiter.tokens_used_current_window == 20
    assert limiter.current_window_start_time == 60.0