Search code examples
pythonpython-3.xunit-testingmockingpydantic

How to mock pydantic BaseModel that expects a Response object?


I'm writing tests for my API client. I need to mock the get function so that it won't make any requests. So instead of returning a Response object I want to return a MagicMock. But then pydantic raises ValidationError because it is going to the model.

I have the following pydantic model:

class Meta(BaseModel):
    raw: Optional[str]
    response: Optional[Response]

    class Config:
        arbitrary_types_allowed = True

which raises:

>   ???
E   pydantic.error_wrappers.ValidationError: 1 validation error for OneCallResponse
E   meta -> response
E     instance of Response expected (type=type_error.arbitrary_type; expected_arbitrary_type=Response)

The one solution would be to add Union with MagicMock but I really don't want to change the code for tests. That is not the way.

class Meta(BaseModel):
    raw: Optional[str]
    response: Optional[Union[Response, MagicMock]]

    class Config:
        arbitrary_types_allowed = True

Any ideas how to patch/mock it?


Solution

  • Instead of using a MagicMock/Mock, you can create a subclass of Response for tests, then patch requests.get to return an instance of that subclass.

    This lets you:

    • Maintain the type of your mock as Response (making pydantic happy)
    • Control most of the expected response behavior for tests
    • Avoid polluting app code with test code (Yes, the "one solution would be to add Union with MagicMock" is definitely not the way.)

    (I'm going to assume the Response is from the requests library. If it isn't, then appropriately adjust the attributes and methods to be mocked. The idea is the same.)

    # TEST CODE
    
    import json
    from requests import Response
    from requests.models import CaseInsensitiveDict
    
    class MockResponse(Response):
        def __init__(self, mock_response_data: dict, status_code: int) -> None:
            super().__init__()
    
            # Mock attributes or methods depending on the use-case.
            # Here, mock to make .text, .content, and .json work.
    
            self._content = json.dumps(mock_response_data).encode()
            self.encoding = "utf-8"
            self.status_code = status_code
            self.headers = CaseInsensitiveDict(
                [
                    ("content-length", str(len(self._content))),
                ]
            )
    

    Then, in tests, you just need to instantiate a MockResponse and tell patch to return that:

    # APP CODE
    
    import requests
    from pydantic import BaseModel
    from typing import Optional
    
    class Meta(BaseModel):
        raw: Optional[str]
        response: Optional[Response]
    
        class Config:
            arbitrary_types_allowed = True
    
    def get_meta(url: str) -> Meta:
        resp = requests.get(url)
        meta = Meta(raw=resp.json()["status"], response=resp)
        return meta
    
    # TEST CODE
    
    from unittest.mock import patch
    
    def test_get_meta():
        mocked_response_data = {"status": "OK"}
        mocked_response = MockResponse(mocked_response_data, 200)
    
        with patch("requests.get", return_value=mocked_response) as mocked_get:
            meta = get_meta("http://test/url")
    
        mocked_get.call_count == 1
        assert meta.raw == "OK"
        assert meta.response == mocked_response
        assert isinstance(meta.response, Response)