How do I mock a class in a Python unit test?

Question

I have a class:

class A:
    __init__(self):
        self.b = B()

   def is_authorized(self)
      name = self.b.get_name()

      if name == 'admin':
          return True
      else:
          return False

I'd like to write a unit test to test the is_authorized method. The problem is that it requires an instance of the B class which is very complicated to construct in isolation as it requires network connections and other stuff. How can I mock this class and provide something that just has the get_name method. So that I can create the A class and test the method.

Answer 1

By use mock library you can patch B class and replace it by a MagicMock() object. mock library was designed exactly to do these kind of works and to break hard dependencies from problematic object or real resources.

In your simple example the complete test will be:

module_a.py

class B():
    def __init__(self):
        print("The awful B class!!!")

    def get_name(self):
        print("The awful B.get_name() method!!!")


class A():
    def __init__(self):
        self.b = B()

    def is_authorized(self):
        name = self.b.get_name()
        if name == 'admin':
            return True
        else:
            return False

module_a_test.py

import unittest
from unittest.mock import patch
from module_a import A


class MyTestCase(unittest.TestCase):

    # patch B class in a_module by a MagicMock instance
    # mock_b_constructor passed to test method
    @patch("module_a.B")
    def test_a(self, mock_b_constructor):
        # B() return value will be the B() instance assigned to a.b property
        mock_b = mock_b_constructor.return_value
        # Now start test:
        a = A()
        # Ok! b is our mock...
        self.assertIs(a.b, mock_b)
        # Not authorized
        self.assertFalse(a.is_authorized())
        mock_b.get_name.return_value = 'admin'
        # Yeah!!! we are admin
        self.assertTrue(a.is_authorized())
        # Sanity check
        mock_b.get_name.return_value = 'guest'
        self.assertFalse(a.is_authorized())

Patch will live just for your test method context. That is a simple and straightforward example of how to use mocks and patches from mock but the real cases can be little more complicated.