Search code examples
pythonmockingpytestpython-unittestpython-unittest.mock

Use methods on Mock object


I have an object that is used for fetching information from another service which is very simple. Since the object is simple and the initialization method could be easily patched I thought I would try to write my code to be super reusable and extendable. But alas, I cannot figure out how to make it work. The code below is pretty well sudo code and is super simplified but it should get the point across.

class SimpleClient:
    def __init__(self):
        pass
    def read(self, key, path='some/path'):
        return value_from_get_on_another_service

I then have a request handler object that initializes a client via get_client() (seen below)

def get_client():
    return SimpleClient()

Then a method on the request handler uses the client.read() method a few times with different parameters (2nd dependent upon the 1st).

For my tests, I thought I could "patch" the get_client method to return my own simple object that could then be used "regularly" and eliminate the dependence on the third party service and actually use the values retrieved from the method execution. I was disappointed to find it was not that easy and clean. The test pattern is seen below.

class MockClient:
    def __init__(self, addr='someAddr', token='someToken'):
        pass

    def read(self, value, prefix):
        data = {}
        if prefix == 'path/1':
            data = self.p1_lookup(value)
        elif prefix == 'path/2':
            data = self.p2_lookup(value)

        return self.response_wrapper(data)

    def p2_lookup(self, key):
        data = {
        'key1': {
            'sub_key': {"55B3FE7D-9F43-4DD4-9090-9D89330C918A": "Dev2",
                        "7A1C2F4B-E91C-4659-A33E-1B18B0BEE2B3": "Dev"}
        }
    }

    return data.get(key, {})


@mock.patch('a.module.get_client')
def test_authorize_valid_request_no_body(mock_get_client):
    request = RequestMock()
    request.body = None
    handler = RequestHandler(Application(), request=request, logging_level='INFO')
    mock_get_client.return_value = MockClient()
    handler.authorize_request()
    assert handler.verified_headers is None
    assert handler.verified_body is None
    assert handler.user_authenticated is False

I have seen where I can mock the responses for the actual client.read() to return multiple values with a list. But this just seems like I will be doing lots of copy and paste and have to do the same thing over and over for each little test. Forgive me if this is simple, sadly I am just learning the art of testing. Is there a way to accomplish what I am trying to do? Maybe there is something super simple I am missing. Or maybe I am just totally on the wrong track for no good reason. Help?!


Solution

  • After a sleep, with fresh eyes I was able to figure this out relatively quickly thanks to a couple other similar questions/answers that I had not found before. Primarily this one, Python Mock Object with Method called Multiple Times.

    Rather than needing to rebuild the module object completely I need to let mock do that for me and then override the specific method on it with the side_effect attribute. So below is what sanitized version of the code looks like.

    def read_override(value, prefix):
        lookup_data1 = {"lookup1": {'key1': 'value1'}}
        lookup_data2 = {'some_id': {'akey': {'12345678': 'DEV'}}
    
        data = {}
        if prefix == 'path1/1a':
            data = lookup_data1.get(value, {})
        elif prefix == 'path2/2a':
            data = lookup_data2.get(value, {})
    
        return {'data': data}
    
    # Create a true Mock of the entire LookupClient Object
    VAULT_MOCK = mock.Mock(spec=LookupClient)
    
    # make the read method work the way I want it to with an "override" of sorts
    VAULT_MOCK.read.side_effect = vault_read_override
    

    Then the test simply looked like this...

    @mock.patch('a.module.get_client')
    def test_authorize_valid_request_no_body(get_client):
        get_client.return_value = VAULT_MOCK
        request = RequestMock()
        request.body = None
        handler = RequestHandler(Application(), request=request, logging_level='INFO')
        handler.authorize_request()
        assert handler.verified_headers is None
        assert handler.verified_body is None
        assert handler.user_authenticated is False