llm_handler/openai_handler_test.py

import os
from typing import List
from unittest import mock
import pytest
import openai
from openai.types.chat.chat_completion_message_param import ChatCompletionMessageParam

import llm_handler.openai_handler as openai_handler


# Clear openai.api_key before each test
@pytest.fixture(autouse=True)
def test_env_setup():
    openai.api_key = None


# This allows us to clear the OPENAI_API_KEY before any test we want
@pytest.fixture()
def mock_settings_env_vars():
    with mock.patch.dict(os.environ, clear=True):
        yield


# Define some constants for our tests
_TEST_KEY: str = "TEST_KEY"
_TEST_MESSAGE: str = "Hello how are you?"
_TEST_MESSAGE_FOR_CHAT_COMPLETION: List[ChatCompletionMessageParam] = [
    {
        "role": "system", "content": "You are serving as a en endpoint to verify a test."
    }, {
        "role": "user", "content": "Respond with something to help us verify our code is working."
    }
]
_TEXT_EMBEDDING_SMALL_3_LENGTH: int = 1536


def test_init_without_key(mock_settings_env_vars):
    # Ensure key not set as env var and openai.api_key not set
    with pytest.raises(KeyError):
        os.environ[openai_handler.OpenAIHandler._ENV_KEY_NAME]
    assert openai.api_key == None
    # Ensure proper exception raised when instantiating handler without key as param or env var
    with pytest.raises(ValueError) as excinfo:
        openai_handler.OpenAIHandler()
    assert f'{openai_handler.OpenAIHandler._ENV_KEY_NAME} not set' in str(excinfo.value)
    assert openai.api_key == None


def test_init_with_key_as_param():
    # Ensure key is set as env var, key value is unique from _TEST_KEY, and openai.api_key not set
    assert not os.environ[openai_handler.OpenAIHandler._ENV_KEY_NAME] == _TEST_KEY
    assert openai.api_key == None
    # Ensure successful instantiation and openai.api_key properly set
    handler = openai_handler.OpenAIHandler(openai_api_key=_TEST_KEY)
    assert isinstance(handler, openai_handler.OpenAIHandler)
    assert openai.api_key == _TEST_KEY


def test_init_with_key_as_env_var(mock_settings_env_vars):
    # Ensure key not set as env var and openai.api_key not set
    with pytest.raises(KeyError):
        os.environ[openai_handler.OpenAIHandler._ENV_KEY_NAME]
    assert openai.api_key == None
    # Set key as env var
    os.environ.setdefault(openai_handler.OpenAIHandler._ENV_KEY_NAME, _TEST_KEY)
    # Ensure successful instantiation and openai.api_key properly set
    handler = openai_handler.OpenAIHandler()
    assert isinstance(handler, openai_handler.OpenAIHandler)
    assert openai.api_key == _TEST_KEY