Unit Testing with Mock and JSON Schema

NicePython

— 3 min read

Writing tests for your code is essential—it's the foundation of ensuring your functionality is stable. This is the key principle behind TDD (Test Driven Development), where tests are written before the actual code. In this article, I’ll walk you through how to effectively use two powerful Python libraries, unittest.mock and jsonschema, which can make testing, especially API testing, much easier..

Example Setup

Let’s say you have two files:

|- get_data.py
|- tests.py

In get_data.py, you make a request to an external API to retrieve data:

import requests as rq

def get_post():
    response = rq.get('http://some-json-backend')
    if response.status_code == 200:
        return response.json()
    raise Exception

In tests.py, you want to write a simple test to check this function:

import unittest
from get_data import get_post

class TestGetData(unittest.TestCase):

    def test_get_post(self):
        response = get_post()
        self.assertIsInstance(response, dict)
        self.assertTrue(response)

unittest.main()

You can run the test by executing tests.py:

python3 tests.py

However, there’s a major drawback—this test depends on an external API. If the API is down or slow, your tests will fail or become unreliable. How can you fix this? You can use mock to replace the part of your code that interacts with the API.

Introducing Mocking with unittest.mock

Let’s refactor the code in get_data.py to separate the API call, making it easier to mock:

import requests as rq

def make_request():
    return rq.get('http://some-json-backend')

def get_post():
    response = make_request()
    if response.status_code == 200:
        return response.json()
    raise Exception

Now, in tests.py, you can mock make_request() to control the behavior and avoid depending on the external API:

import unittest
from get_data import get_post
from unittest.mock import patch, Mock

class TestGetData(unittest.TestCase):

    @patch('get_data.make_request')
    def test_get_post(self, mock_make_request):
        mock_response = Mock(status_code=200)
        test_body = "sample_body"
        mock_response.json.return_value = {
            "userId": 1,
            "id": 1,
            "title": "sample_title",
            "body": test_body
        }
        mock_make_request.return_value = mock_response

        response = get_post()
        self.assertIsInstance(response, dict)
        self.assertTrue(response)
        self.assertEqual(response["body"], test_body)

unittest.main()

Here’s what’s happening:

  • We use the @patch decorator to replace make_request() with a mock.
  • We create a mock object (mock_response) to simulate a successful API call.
  • The mock returns a predefined JSON response that mimics the actual API output.
  • This ensures our tests run smoothly, independent of any external API.

One downside of this approach is that if the API changes, you may need to update your test logic. But at least, your test won’t break due to API downtime.

Validating JSON with jsonschema

One issue with the current test is that we are only checking one field (body) from the API response. What if the entire JSON structure needs validation? This is where jsonschema comes in handy.

First, install jsonschema:

pip install jsonschema

Now, you can define a schema to validate the entire API response:

import unittest
from get_data import get_post
from unittest.mock import patch, Mock
from jsonschema import validate

class TestGetData(unittest.TestCase):

    def setUp(self):
        self.schema = {
            "type": "object",
            "properties": {
                "userId": {"type": "number"},
                "id": {"type": "number"},
                "title": {"type": "string"},
                "body": {"type": "string"}
            },
            "required": ["userId", "id", "title", "body"]
        }

    @patch('get_data.make_request')
    def test_get_post(self, mock_make_request):
        mock_response = Mock(status_code=200)
        test_body = "sample_body"
        mock_response.json.return_value = {
            "userId": 1,
            "id": 1,
            "title": "sample_title",
            "body": test_body
        }
        mock_make_request.return_value = mock_response

        response = get_post()
        self.assertIsInstance(response, dict)
        self.assertTrue(response)
        self.assertEqual(response["body"], test_body)
        validate(instance=response, schema=self.schema)

unittest.main()

With this setup, you can now validate the entire JSON structure according to the defined schema:

  • The schema defines the expected structure of the API response.
  • The validate() function ensures that the response matches the schema.
  • The required field ensures that all necessary fields are present.

If you remove a required field (e.g., userId), the schema validation will fail:

mock_response.json.return_value = {
    # "userId": 1,  # intentionally omitted
    "id": 1,
    "title": "sample_title",
    "body": "sample_body"
}

When you run the test, you’ll get an error like this:

jsonschema.exceptions.ValidationError: 'userId' is a required property

This prevents passing a test with incomplete data, making your tests more robust.

Conclusion

In this article, I showed you how to test code that interacts with external APIs using Python’s unittest.mock and jsonschema libraries. By mocking external dependencies and validating JSON responses, you can make your tests more reliable and comprehensive. I hope this helps you write better tests for your projects!

Read more