About
Side Projects
Blog
2020-08-13

Mocking a Class in a Python Unit Test

I have recently been writing some Python logic to manipulate data stored in Google Cloud Storage. I needed to write some build-time unit tests for this logic and in doing so wanted to mock out the Google Cloud Storage classes that are provided for the Python language. I wanted the logic that I had written to write to work as it normally does but not to access Cloud Storage while it is under test but to access a fake Cloud Storage.

To show how this is done, here is an artificial example of some Python logic that uses a class called Transform;

class Transform:
  msg = None

  def __init__(self, msg):
    self.msg = msg

  def transform(self):
    return '>> ' + self.msg


def process():
  outputta = Transform('grok')
  return outputta.transform()


if __name__ == '__main__':
  process()

A test is written to exercise the process function below. In this test, I want to mock the class Transform so that when a Transform is created, a different class is created instead. Here is how it looks;

import unittest
from unittest.mock import patch
import example as e


class TransformErsatz:

  def __init__(self, ignore): pass

  def transform(self):
    return '>> mock'

class ExampleTests(unittest.TestCase):

  @patch("example.Transform", side_effect=TransformErsatz)
  def test_process(self, mock_transform):
    # GIVEN

    # WHEN
    result = e.process()

    # THEN
    self.assertEqual(result, '>> mock')


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

The @patch argument is provided to the test function in order to identify the class to mock and, also to provide the replacement mocked class; in this case TransformErsatz. In the actual logic under test when an instance of Transform is created, instead an instance of TransformErsatz will be created. If there is a need to mock multiple classes in this way, it is possible to add additional @patch(...) clauses and to add additional arguments carrying the mocks to the test function arguments in the signature.

In this way it was possible in my actual work to create a mock set of Google Cloud Storage classes to test the logic against.