Unit Testing in Python: Part 1 – Requests

Welcome to a blog series about unit testing in Python. In this first post, I’d like to explore some basics of how to test REST API calls.

I’ve recently written several large Python classes to call a REST API. When I first started, the script was very small and I could easily run it to test all the situations. It quickly grew and it was obvious that I needed some more sophisticated testing. After some reading, I found I could solve this problem with two Python libraries: unittest and requests-mock.

Python has a very good unit testing framework in unittest. However, a script that makes many external requests can be very difficult to test because you don’t want to make a real call to a live server during a unit test. Luckily, I always use the requests library for my HTTP(S) calls in python and there is this great library called requests-mock that will capture all requests to a certain URL and return some JSON (or whatever) I specify. I’ll use an example to show how easy this is. (For the full source code, please see my GitHub.)

I’ll start with a simple class with one method to create a user.

import requests

class ApiWrapper(object):
    """Class that wraps some REST API"""

    def __init__(self, api_url: str):
        self._api_url = api_url

    def add_user(self, first_name: str, last_name: str) -> str:
        """Adds a new user to the system and returns its ID"""
        params = {'first_name': first_name,
                  'last_name': last_name}
        response = requests.post(self._api_url + '/users',
                                 json=params)
        if response.status_code != 201:
            raise RuntimeError

        return response.json()['id']

I need to test the logic in the add_user method but I don’t want to hit the real REST API during my unit testing. You can see below that this is quite simple.

I first add the @requests_mock.mock() annotation to the class so that each test method will be passed a mock object.

@requests_mock.mock()
class ApiWrapperTest(unittest.TestCase):

In this test, I want to test the case where the POST fails and the case where it succeeds and returns valid JSON. In the first case I set up mock to return 401 and test that it does raise a RuntimeError. In the second case I tell mock to return a valid JSON and test that add_user returns the correct ID.

def test_add_user(self, mock):
    wrapper = ApiWrapper(API_URL)

    mock.post(API_URL + '/users', status_code=401)
    with self.assertRaises(RuntimeError):
        wrapper.add_user('The', 'User')

    mock.post(API_URL + '/users', text='{"id": "1234"}',
              status_code=201)
    self.assertEqual('1234',
                     wrapper.add_user('The', 'User'))

Lastly, I add the main method so that I can easily run this from the command line by executing the test file (e.g. python apiwrapper_test.py).

if __name__ == '__main__':
    unittest.main()

Here is a look at the file in its entirety.

import unittest
import requests_mock
from apiwrapper import ApiWrapper

API_URL = 'http://example.com'

@requests_mock.mock()
class ApiWrapperTest(unittest.TestCase):
    """Tests ApiWrapper"""

    def test_add_user(self, mock):
        wrapper = ApiWrapper(API_URL)

        mock.post(API_URL + '/users', status_code=401)
        with self.assertRaises(RuntimeError):
            wrapper.add_user('The', 'User')

        mock.post(API_URL + '/users', text='{"id": "1234"}',
                  status_code=201)
        self.assertEqual('1234',
                         wrapper.add_user('The', 'User'))

if __name__ == '__main__':
    unittest.main()

I hope you enjoyed this look at Python’s unit testing capabilities. Join me next time as I explore how to see the code coverage of the unit tests.

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