Unit Testing ConfigurableResource Classes

[rover-wagnar]

Hi,

I'm looking for help with a better workaround or solution for handling mocking with classes that inherit from ConfigurableResource. Here's the issue I'm running into with an example.

Say I have this class, a very simple API client class that's a ConfigurableResource and handles setting a bearer token value for paginated API calls:

class RestApiClient(ConfigurableResource):
    _api_token: Optional[ApiToken] = None

    @property
    def api_token(self) -> ApiToken:
        if not self._api_token or self._api_token.has_expired():
            self._api_token = self._refresh_access_token()
        return self._api_token

     def refresh_access_token(self) -> ApiToken:
        # Code for handling a POST call for an API token value

    def get_paginated(self, api_path: str) -> List[Dict]:
        headers = {}
        headers["Authorization"] = "Bearer {}".format(self.api_token.value)
        
        response = requests.Session().get(request_url, headers=headers)
        response.raise_for_status()
        # code for handling response data and paginated calls, etc.

If I want to test that calling get_paginated means refresh_access_token is called, I would write a unit test like this, where I patch the object to mock a method so I can make assertions on call behavior:

@pytest.fixture
def rest_api_client():
    return RestApiClient(...)

def test_calls_refresh_access_token_if_no_access_token(rest_api_client, mock_response):
    # After initial construction, no api_token should be present so we expect a call to refresh_access_token
    with (
        patch('code_location.resources.rest_api.requests.Session.get') as mock_get,
        patch.object(rest_api_client, 'refresh_access_token') as mock_refresh_access_token
    ):
        mock_get.return_value = mock_response
        
        rest_api_client.get_paginated(FAKE_API_ENDPOINT)
        
    mock_refresh_access_token.assert_called_once()

The problem is, this test will fail. It fails because the way unittest.mock.patch apparently handles mocking is by using assignment to make a method a Mock object (which records interactions so we can call methods like .assert_called_once()). But ConfigurableResource objects are Pydantic BaseModels with Frozen=True which don't allow re-assignment like this, so the mock fails and the test fails with a Pydantic ValidationError.

One workaround for this is to autospec a MagicMock object with the ApiClient class like MagicMock(spec=RestApiClient). But this object does not have the same method definitions so the mocked version of get_paginated does not call refresh_access_token, and so I can't assert the method was called (unless I redefine mock methods for the MagicMock which defeats the purpose of the test).

Currently, I just don't write those kinds of unit tests for ConfigurableResource classes, but I feel like this limitation around mocking for ConfigurableResources limits test coverage. Am I missing a better way to unit test ConfigurableResources or a better workaround?

[theholy7]

Can anyone help answer this?

[dehume]

Playing around with this. If you patch the method on the class before instantiating the configurable resource, that avoids the Pydantic reassignment issue.

Slightly cleaning up the original code so it will run:

from typing import Optional
from pydantic import Field
import dagster as dg
from pydantic import BaseModel
import requests
from typing import List, Dict


class ApiToken(BaseModel):
    value: str = Field(description="The API token value")
    has_expired: bool = Field(description="Whether the token has expired")


class RestApiClient(dg.ConfigurableResource):
    _api_token: Optional[ApiToken] = None

    @property
    def api_token(self) -> ApiToken:
        if not self._api_token or self._api_token.has_expired():
            self._api_token = self.refresh_access_token()
        return self._api_token

    def refresh_access_token(self) -> ApiToken:
        # Code for handling a POST call for an API token value
        return "my_key"

    def get_paginated(self, api_path: str) -> List[Dict]:
        headers = {}
        request_url = f"https://api.example.com/{api_path}"
        headers["Authorization"] = "Bearer {}".format(self.api_token.value)

        response = requests.Session().get(request_url, headers=headers)
        response.raise_for_status()

You will still need to mock a number of things but you can still follow everything that is happening:

from example.defs.resources import RestApiClient, ApiToken

from unittest.mock import patch, Mock


@patch("example.defs.resources.RestApiClient.refresh_access_token")
def test_calls_refresh_access_token_if_no_access_token(mock_refresh_access_token):
    # Mock Api Token
    mock_token = Mock(ApiToken)
    mock_token.value = "test_key"
    mock_token.has_expired = False

    mock_refresh_access_token.return_value = mock_token

    # Construct after patch
    client = RestApiClient()
    mock_response = Mock()

    with patch("requests.Session.get") as mock_get:
        mock_get.return_value = mock_response
        client.get_paginated("test")

    mock_refresh_access_token.assert_called_once()
    mock_get.assert_called_once_with(
        "https://api.example.com/test", headers={"Authorization": "Bearer test_key"}
    )