Testing OAuth with patch/mock

I finally did the right thing and added unit testing to my Flask web application. However, I ran into extreme difficulty with testing the authentication feature. I’m using OAuth for authentication. I understand that the appropriate thing to do when testing a feature that makes calls to an external API is to “mock” or “stub out” the external calls, because we want to test only our own code and not that of a third party provider.

Patch and Mock

For testing, I’m using python’s unittest framework. From this helpful blog post I learned that the way to fake external calls in unittest is by using objects called Mocks (or a subclass called MagicMock) as well as the @patch decorator (also see the official documentation).

The @patch decorator gets placed before before the test function in the test module and takes as a parameter the name of a class or function that will get swapped out when the test is executed. The test function then gets an additional parameter, which by default is a MagicMock object:

import unittest
from unittest.mock import patch, MagicMock

@patch('some_module.external_function')
def test_my_function(self, mock_external_function):
    ...

To give our mock_external_function the output we expect when it’s called, we set its return value in the body of our test:

@patch('some_module.external_function')
def test_my_function(self, mock_external_function):
    mock_external_function.return_value = "expected return value"
    ...

Alternatively, if we were to mock a class instead of a function, instead of setting a return value we could just assign it whatever attributes or methods we desire.

Seems simple enough, but I found it not so straightforward to apply to my own code. Let me first show you the code responsible for authentication in my app.

Authentication code

There are two views responsible for authentication:

# views.py

@app.route('/authorize/<provider>')
def oauth_authorize(provider):
    if not current_user.is_anonymous:
	return redirect(url_for('index'))
    oauth = OAuthSignIn.get_provider(provider)
    return oauth.authorize()

@app.route('/callback/<provider>')
def oauth_callback(provider):
    if not current_user.is_anonymous:
	return redirect(url_for('index'))
    oauth = OAuthSignIn.get_provider(provider)
    social_id, username = oauth.callback()
    if social_id is None:
	flash('Authentication failed.')
	return redirect(url_for('index'))
    user = User.query.filter_by(social_id=social_id).first()
    if not user:
	user = User(social_id=social_id, nickname=username)
	db.session.add(user)
    else:
	# update user nickname in case it's changed
	user.nickname = username
    db.session.commit()
    login_user(user, True)
    return redirect(url_for('index'))

The /authorize endpoint redirects the user to the appropriate endpoint in the OAuth provider’s API. After completing authentication, the OAuth provider then calls our /callback endpoint, which extracts the desired information about the user, adds or updates the user entry to the database, and logs in the user to the app. The OAuthSignIn class is one that I define in a separate file oauth.py. It has a subclass for each OAuth provider we allow. The OAuthSignIn.get_provider(provider) method simply creates an instance of the appropriate subclass. Using the example of Facebook, the subclass looks as follows:

# oauth.py
from rauth import OAuth2Service
...
class FacebookSignIn(OAuthSignIn):
    def __init__(self):
        super(FacebookSignIn, self).__init__('facebook')
        self.service = OAuth2Service(
            name='facebook',
            client_id=self.consumer_id,
            client_secret=self.consumer_secret,
            authorize_url='https://graph.facebook.com/oauth/authorize',
            access_token_url='https://graph.facebook.com/oauth/access_token',
            base_url='https://graph.facebook.com/'
        )

    def authorize(self):
        return redirect(self.service.get_authorize_url(
            scope='public_profile',
            response_type='code',
            redirect_uri=self.get_callback_url())
        )

    def callback(self):
        # request here is coming from Facebook's call to our /callback endpoint
        if 'code' not in request.args:
            return None, None
        oauth_session = self.service.get_auth_session(
            data={'code': request.args['code'],
                  'grant_type': 'authorization_code',
                  'redirect_uri': self.get_callback_url()}
        )
        me = oauth_session.get('me?fields=id,first_name').json()
        return (
            'facebook$' + me['id'],
            me.get('first_name')  # use first name as user nickname
        )

Note that we use the rauth library, and in particular the OAuth2Service class, to handle interactions between our app and Facebook. The authorize() function is invoked by our /authorize endpoint and handles the redirect to the OAuth provider. The more interesting method is callback(), which makes two external calls when invoked by the OAuth provider:

  1. In the call self.service.get_auth_session(), the  OAuth2Service object in self.service talks to the provider and exchanges the given authentication code for an authentication session, creating an OAuth2Session object oauth_session.
  2. Using the session object, we make a GET request to the provider for the information about the user that we want in the call oauth_session.get('me?fields=id,first_name')

Testing

I decided to break my authentication testing into two tests: one for the authorization redirect and one for the callback. The first one was easy: we simply test that the /authorize/facebook endpoint redirects to the appropriate endpoint in Facebook’s API:

def test_facebook_oath_authorize(self):
    rv = self.app.get('/authorize/facebook')
    # response should be a redirect
    assert rv.status_code == 302
    # should redirect to facebook's authorization endpoint
    assert rv.location.split('?')[0] == 'https://graph.facebook.com/oauth/authorize'

Testing the callback is much more difficult, because this is the part where our code makes external calls that we need to mock out.

Attempt 1

In my first working test for the callback, I did the laziest possible thing. Because my method oauth.FacebookSignIn.callback() is the one that makes the two external calls, I mocked out that entire function, with the mock returning a fake social_id and nickname:

# callback test, attempt 1
@patch('oauth.FacebookSignIn.callback')
def test_facebook_oath_callback(self, mock_callback):
    mock_callback.return_value = 'facebook$3617923766551', 'Andrea'
    rv = self.app.get('/callback/facebook')
    # when response is new user, db entry created
    assert db.session.query(User).count() == 1
    # when response is existing user, no entry added
    mock_callback.return_value = 'facebook$3617923766551', 'Andy'
    self.app = app.test_client()
    rv = self.app.get('/callback/facebook')
    assert User.query.filter_by(social_id='facebook$3617923766551').count() == 1
    # if user's nickname has changed, db entry is updated
    user = User.query.filter_by(social_id='facebook$3617923766551').first()
    assert user.nickname == 'Andy'

This works: the test passes, and no external calls are made. However, although I haven’t had much testing experience and am not familiar with best practices, this approach strikes me as bad. Mocking out one’s own function when testing does not seem like a good idea: we only want to replace the direct calls being made to an external provider. That way our code runs in test the same way it would in production, except that it’s talking to a fake authentication server instead of a real one.

Attempt 2

Ok, now let’s try to do things the right way. Recall that two external calls we make are in oauth_session = self.service.get_auth_session() and oauth_session.get().json(), both in oauth.FacebookSignIn.callback(). Since the second call is made on an object returned by the first call, we need to set up our mocks in the following way:

  1. Patch the function get_auth_session() and have it return a mock session, which will be a MagicMock() object.
  2. The mock session should have a method called get(), which returns a mock response to a get request. The mock response is also a MagicMock() object.
  3. The mock response should have a mock method json(), which returns a python dictionary with 'id' and 'first_name' keys.

So, in total we will be needing three mocks. Here is the final test:

@patch('oauth.OAuth2Service.get_auth_session')
def test_facebook_oath_callback2(self, mock_get_auth_session):
    mock_session = MagicMock()
    mock_get_response = MagicMock(status_code=200, json=MagicMock(return_value={'first_name': 'Andrea', 'id': '3617923766551'}))
    mock_session.get.return_value = mock_get_response
    mock_get_auth_session.return_value = mock_session
    rv = self.app.get('/callback/facebook?code=some_code')
    # when response is new user, db entry created
    assert db.session.query(User).count() == 1
    # when response is existing user, no entry added
    mock_get_response.json.return_value = {'first_name': 'Andy', 'id': '3617923766551'}
    self.app = app.test_client() # need to reset test_client when making another request
    rv = self.app.get('/callback/facebook?code=some_code')
    assert User.query.filter_by(social_id='facebook$3617923766551').count() == 1
    # if user's nickname has changed, db entry is updated
    user = User.query.filter_by(social_id='facebook$3617923766551').first()
    assert user.nickname == 'Andy'

This test also has the desired functionality, but also the additional benefit of not mocking more than is necessary. Whew!

Thanks for reading!

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Google photo

You are commenting using your Google account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s