pythonunit-testingmockingautomated-testspytest

How to Unit Test a Python Class Which Needs to Make an API Call to an External Service?


The class is meant to wrap a remotely hosted large language model and has to make an API call to the service to fetch the results. This is an example.


class ModelWrapper(AbstractLLMInterface):
    """The Claude 3 Sonnet model wrapper following the interface."""

    def __init__(
        self,
        region: str = CONFIGS["GCP"]["REGION"],
        project: str = get_gcp_project_id(),
        model: str = CONFIGS["MODELS"]["SONNET_ID"],
    ) -> None:
        """Set up the claude client using the region and project."""
        self.client: AnthropicVertex = AnthropicVertex(region=region, project_id=project)
        self.model_name: str = model
        self.role_key: str = "role"
        self.content_key: str = "content"
        self.user_key: str = "user"
        llm_logger.debug(msg=f"Initialised sonnet client for {region}, {project} and {self.model_name}.")

    def get_completion(self, user_prompt: str, system_prompt: str, history: List[Correspondence]) -> Iterator[str]:
        """
        Fetch a response from the model.This requires an egress request to GCP and the
        service for Anthropic model must be enabled in the VertexAI console.
        """

        # This is where the API call to GCP service happens
        return self.client.messages.stream(user_prompt, system_prompt, history)

I am aware of MagicMocks for python unit testing which are objects that can be configured to return anything I want. But in this instance, all the parameters of the constructor are simple strings, and I the client is built inside the class. So there is no scope for a mock client, right?

Does it mean any unit test must make an API call? Or is the class designed incorrectly? Any help will be appreciated.

The class should be testable without having to call the API.

Note

I know the changing the class constructor to accept an AnthropicVertex object make the mocking trivially simple. But the purpose of the question is to test this class as it exists, not to change the class.

However, if the current design of the class violates any basic design principle (such as SOLID), that is something I would love to know and understand. So any reference will be appreciated.


Solution

  • import importlib
    from unittest import mock
    
    import anthropic 
    
    import your_module
    from your_module import ModelWrapper
    
    with mock.patch.object(target=anthropic, attribute="AnthropicVertex"):
        importlib.reload(module=your_module)
        assert isinstance(ModelWrapper().client, mock.MagicMock)
    

    If you instantiate your class ModelWrapper under the with block, AnthropicVertex will remain patched.