Mocking the OpenAI API in Python: A Step-by-Step Guide

As developers, we often need to test our applications without making actual API calls, especially when dealing with APIs that have rate limits or costs associated with them. At Areto (www.areto.de), we experiment with OpenAI’s powerful ChatGPT models to enhance our services.

However, making numerous calls during testing can quickly eat into your API limits. To avoid this issue, you can mock the API calls in your Python code, allowing you to simulate the responses and test your application efficiently.

In this blog post, I’ll walk you through the process of mocking the OpenAI API in Python using the pytest_mock library.

All the code can be found on my GitHub: https://github.com/TobiLang/chatgpt_mocker

Install Required Libraries

First, make sure you have the OpenAI Python library installed, as well as the pytest_mock library. You can install them using pip:

pip install openai pytest_mock

Setup a minimal project

We use a minimal project with the following files and directory structure:

chatgpt_mock/
│
├── src/chatgpt/
│   └── chatgpt_handler.py
│
└── src/tests/
    ├── chatgpt/
    │   └── test_chatgpt_handler.py
    │
    ├── __init__.py
    └── conftest.py

Let us start with the chatgpt_handler.py file itself. It is pretty straightforward code, implementing the interaction with the OpenAI API:

"""Process data using ChatGPT."""
import logging
from typing import Optional

import openai
from openai.error import APIConnectionError, AuthenticationError

logger = logging.getLogger(__name__)


# pylint: disable=too-few-public-methods
class ChatGPTHandler:
    """Process data using ChatGPT."""

    MODEL = "gpt-3.5-turbo"

    def __init__(self, openai_key: str = "") -> None:
        """
        Initialize the class.
        """
        openai.api_key = openai_key

    def query_api(self, query: str) -> Optional[str]:
        """
        Query the ChatGPT API.

        Args:
            query: Query to send to the API

        Returns:
            Response message from the API
        """
        logging.info("Querying API...")

        # No need to query the API if there is no query content
        if not query:
            return None

        message = [{"role": "user", "content": query}]

        result = None
        try:
            completion = openai.ChatCompletion.create(
                model=self.MODEL, messages=message
            )  # type: ignore[no-untyped-call]
            result = completion.choices[0].message.content
        except AuthenticationError as ex:
            logger.error("Authentication error: %s", ex)
        except APIConnectionError as ex:
            logger.error("APIConnection error: %s", ex)

        return result

The ChatGPTHandler class is defined within the ChatGPTHandler.py file and has two main methods:

1. __init__(self, openai_key: str = ""): The class constructor that initializes the ChatGPTHandler instance and sets the OpenAI API key.
2. query_api(self, query: str) -> Optional[str]: The main method that queries the ChatGPT API and returns the response message. The method also includes error handling for AuthenticationError and APIConnectionError, which are both imported from the openai.error module. If an exception is caught, it logs the error message using the logger.error() method.

Testing

To test the above class, we first set up some helper classes in a conftest.py file. These classes are designed to help with testing by converting nested dictionaries into objects and mocking OpenAI API completions.

"""Configuration for ChatGPT tests."""
from typing import Any, Dict, Optional


# pylint: disable=too-few-public-methods
class NestedDictToObject:
    """
    Converts a nested dictionary into an object.
    """

    def __init__(self, dictionary: Dict[str, Any]) -> None:
        """
        Initializes an instance of NestedDictToObject.

        Args:
            dictionary: A nested dictionary to convert into an object.
        """
        for key, value in dictionary.items():
            if isinstance(value, dict):
                setattr(self, key, NestedDictToObject(value))
            else:
                setattr(self, key, value)


# pylint: disable=too-few-public-methods
class MockedCompletion:
    """
    Mock OpenAI API completion.
    """

    def __init__(self, message: Optional[Dict[str, Any]] = None) -> None:
        """
        Initialize the mocked API parameters.

        Args:
            message: Dictionary containing the message to be mocked.
                     Will be converted to an object.
        """
        if message is None:
            message = {"message": {"content": "This is a mocked message."}}

        # Convert to object
        message_obj = NestedDictToObject(message)

        self.choices = [message_obj]

    def __next__(self) -> "MockedCompletion":
        """
        Return the class instance itself.

        Unittest mock expects an iterable object, so we need to implement
        this method. Otherwise, we would get an error like this:
            TypeError: 'MockedCompletion' object is not iterable

        Returns:
            The class instance itself.
        """
        return self

Class 1: NestedDictToObject

The NestedDictToObject class converts a nested dictionary into an object, making it easier to access the dictionary’s keys as object attributes. As the current implementation of the OpenAI API returns objects rather than dictionaries, we need it to not break the implementation while testing against the mocked version.

Class 2: MockedCompletion

The MockedCompletion class will be used to mock the openai.ChatCompletion.create call. It accepts an optional dictionary argument message, in case the unit tests ask for a specific return content. The message dictionary is then converted to an object using the NestedDictToObject class, and the resulting object is added to the choices attribute as a list.

This class could be extended to handle more than just the choices parameter. Other possible candidates are (depending on the actual use case and test scenario), e.g.:

id: 'example-id'
created: 1234567890
model: 'text-davinci-002'
finish_reason: 'stop'

See https://platform.openai.com/docs/api-reference/making-requests for more details on the response object.

The __next__ method returns the class instance itself. This method is used for making the class instance iterable and allows it to be used in for loops or other iterable constructs. Unittest mock expects an iterable object, so we must implement this method. Otherwise, we would get an error like this:

TypeError: ‘MockedCompletion’ object is not iterable

And finally, the test_chatgpt_handler.py file itself:

"""Tests processing using ChatGPT."""

import pytest
from openai.error import AuthenticationError
from pytest_mock import MockerFixture

from src.chatgpt.chatgpt_handler import ChatGPTHandler
from tests.conftest import MockedCompletion


class TestChatGPTHandler:
    """Test class for chatgpt_handler.py."""

    @pytest.fixture
    def gpt_handler(self, mocker: MockerFixture) -> ChatGPTHandler:
        """
        Fixture for testing, return a mocked ChatGPTHandler instance.

        Args:
            mocker: MockerFixture

        Returns:
            ChatGPTHandler instance
        """
        # Mock the openai.ChatCompletion.create method
        mocker.patch("openai.ChatCompletion.create", side_effect=MockedCompletion())
        return ChatGPTHandler(openai_key="SECRET_API_KEY")

    def test_query_api(self, gpt_handler: ChatGPTHandler) -> None:
        """
        Test for querying the OpenAI API.

        Args:
            gpt_processor: Mocked ChatGPTHandler instance

        Returns:
            None
        """
        query = "Riddle me this ..."

        result = gpt_handler.query_api(query)

        assert isinstance(result, str)
        assert len(result) > 0

    def test_query_api_empty(self, gpt_handler: ChatGPTHandler) -> None:
        """
        Test for querying the OpenAI API with empty query.

        Args:
            gpt_processor: Mocked ChatGPTHandler instance

        Returns:
            None
        """
        result = gpt_handler.query_api("")

        assert result is None

    def test_query_api_authentication_error(
        self,
        mocker: MockerFixture,
        gpt_handler: ChatGPTHandler,
    ) -> None:
        """
        Test for querying the OpenAI API, triggering an AuthenticationError.

        Returns:
            None
        """
        side_effect = AuthenticationError("mocked error")
        mocker.patch("openai.ChatCompletion.create", side_effect=side_effect)

        query = "Riddle me this ..."

        result = gpt_handler.query_api(query)

        assert result is None

In this Python test file, a series of tests are designed to check the functionality of the ChatGPTHandler class, which interacts with the OpenAI API for processing chat inputs. The test suite uses the pytest framework and pytest-mock to create test cases with mocked objects and functions.

The test file begins with several imports, including the required classes, methods, and fixtures from different modules. The TestChatGPTHandler class is then defined to include all the tests related to the ChatGPTHandler class.

There are three test cases and one fixture in this test suite:

1. gpt_handler: This is a pytest fixture that returns a mocked ChatGPTHandler instance. It mocks the openai.ChatCompletion.create method using MockedCompletion from the tests.conftest module. This allows the test suite to avoid making actual API calls while testing. A default message is set for all tests using this fixture.
2. test_query_api: This test case verifies that the query_api method of the ChatGPTHandler class returns a non-empty string when provided with a valid query.
3. test_query_api_empty_query: This test case checks the behavior of the query_api method when given an empty string as input.
4. test_query_api_authentication_error: This test case tests the behavior of the query_api method when an AuthenticationError occurs. The test uses the mocker fixture to mock the openai.ChatCompletion.create method to throw a mocked AuthenticationError.

Conclusion

Mocking the OpenAI API in Python allows you to test your application efficiently and avoid consuming your API limits.

The classes and test files are kept simple to show the basics of testing different things while mocking the API. In a real-world scenario, consider what to test and avoid testing the mock instead of your actual code.

Happy coding!