Mock a function of an attribute Python

Question

I have here:

class Authentication(MethodView):

    @staticmethod
    def make_payload(identity):
        iat = datetime.utcnow()
        exp = iat + timedelta(seconds=300)
        nbf = iat + timedelta(seconds=0)
        identity = identity.key.urlsafe()
        return {'exp': exp, 'iat': iat, 'nbf': nbf, 'identity': identity}

which is based on JWT. I want to mock identity.key.urlsafe() to return 1 which is a user id:

def test_make_payload(self):
    def get_urlsafe():
        return self.user_id

    now = datetime.utcnow()
    mock_identity = MagicMock()
    mock_identity.key.return_value = MagicMock(urlsafe=get_urlsafe)
    payload = Authentication.make_payload(mock_identity)

Right now, everything works except my mock. The goal was for it to return 1:

ipdb> payload
{'identity': <MagicMock name='mock.key.urlsafe()' id='4392456656'>, 'iat': datetime.datetime(2016, 11, 22, 21, 34, 41, 605698), 'nbf': datetime.datetime(2016, 11, 22, 21, 34, 41, 605698), 'exp': datetime.datetime(2016, 11, 22, 21, 39, 41, 605698)}
ipdb> payload['identity']
<MagicMock name='mock.key.urlsafe()' id='4392456656'>

How can I mock this nested call to make urlsafe return 1 in my mock?

Answer 1

It looks to me like you want to set the side_effect of your mock to be get_urlsafe:

def test_make_payload(self):
    def get_urlsafe():
        return self.user_id

    now = datetime.utcnow()
    mock_identity = MagicMock()
    mock_identity.key.urlsafe.side_effect = get_urlsafe
    payload = Authentication.make_payload(mock_identity)

From the documentation:

side_effect: A function to be called whenever the Mock is called. See the side_effect attribute. Useful for raising exceptions or dynamically changing return values. The function is called with the same arguments as the mock, and unless it returns DEFAULT, the return value of this function is used as the return value.