platform/reworkd_platform/tests/test_helpers.py

import pytest
from openai.error import InvalidRequestError, ServiceUnavailableError

from reworkd_platform.schemas.agent import ModelSettings
from reworkd_platform.web.api.agent.helpers import openai_error_handler
from reworkd_platform.web.api.errors import OpenAIError


async def act(*args, settings: ModelSettings = ModelSettings(), **kwargs):
    return await openai_error_handler(*args, settings=settings, **kwargs)


@pytest.mark.asyncio
async def test_service_unavailable_error():
    async def mock_service_unavailable_error():
        raise ServiceUnavailableError("Service Unavailable")

    with pytest.raises(OpenAIError):
        await act(mock_service_unavailable_error)


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "settings,should_log",
    [
        (ModelSettings(custom_api_key="xyz"), False),
        (ModelSettings(custom_api_key=None), True),
    ],
)
async def test_should_log(settings, should_log):
    async def mock_invalid_request_error_model_access():
        raise InvalidRequestError(
            "The model: xyz does not exist or you do not have access to it.",
            param="model",
        )

    with pytest.raises(Exception) as exc_info:
        await openai_error_handler(
            mock_invalid_request_error_model_access, settings=settings
        )

    assert isinstance(exc_info.value, OpenAIError)
    error: OpenAIError = exc_info.value

    assert error.should_log == should_log