Skip to content

Instantly share code, notes, and snippets.

@velp
Last active February 27, 2017 10:12
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save velp/a56c6877bfa61548621ebc9f2457c3e6 to your computer and use it in GitHub Desktop.
Save velp/a56c6877bfa61548621ebc9f2457c3e6 to your computer and use it in GitHub Desktop.

Example use of the 'mock' module for unit testing in Python

import time
import unittest
try:
    from unittest import mock
except ImportError: # pragma: no cover
    import mock


# Some libraries client performs a long-acting
class ClientClass: # pragma: no cover
    state = None

    def __init__(self):
        print("Call %s.__init__()" % type(self).__name__)
        self.status = None

    def __enter__(self):
        print("Call %s.__enter__()" % type(self).__name__)
        return self

    def __exit__(self, *args, **kwargs):
        print("Call %s.__exit__()" % type(self).__name__)

    def start(self):
        print("Call %s.start()" % type(self).__name__)
        self.status = "running"
        return True

    def stop(self):
        print("Call %s.stop()" % type(self).__name__)
        self.status = "stop"
        return True

    def get_status(self):
        print("Call %s.get_status()" % type(self).__name__)
        return self.status


# Tested class, used ClientClass
class ProductionClass:

    def __init__(self):
        self.client = ClientClass()

    def _start(self):
        # Prepare config
        return self.client.start()

    def _stop(self):
        # Pre-stoped operations
        return self.client.stop()

    def _wait(self):
        ticker = 0
        while True:
            status = self.client.get_status()
            if status != "running":
                return status
            elif ticker >= 10:
                return "timeout %d" % ticker
            else:
                ticker += 1
                time.sleep(1)

    def run(self):
        self._start()
        res = self._wait()
        self._stop()
        return res

    def run_context(self):
        try:
            with ClientClass() as cc:
                res = cc.start()
                cc.stop()
                return cc.state
        except Exception as err:
            return str(err)


# Tests
class TestTestedClass(unittest.TestCase):

    @mock.patch('%s.ClientClass.__init__' % __name__, return_value=None)
    def setUp(self, m):
        self.pc = ProductionClass()

    @mock.patch('%s.ClientClass.start' % __name__)
    def test_start(self, m):
        m.return_value = True
        self.assertTrue(self.pc._start())
        self.assertTrue(m.called)

    @mock.patch('%s.ClientClass.stop' % __name__)
    def test_stop(self, m):
        m.return_value = True
        self.assertTrue(self.pc._stop())
        self.assertTrue(m.called)

    @mock.patch('time.sleep', return_value=None)
    @mock.patch('%s.ClientClass.get_status' % __name__)
    def test_wait(self, m, mock_sleep):
        # test timeout
        m.return_value = "running"
        self.assertIn('timeout', self.pc._wait())
        self.assertTrue(m.called)
        # test status != "running"
        m.return_value = "stop"
        self.assertEqual(self.pc._wait(), 'stop')
        self.assertTrue(m.called)

    @mock.patch('%s.ProductionClass._start' % __name__)
    @mock.patch('%s.ProductionClass._stop' % __name__)
    @mock.patch('%s.ProductionClass._wait' % __name__)
    def test_run(self, mock_wait, mock_stop, mock_start):
        mock_wait.return_value = "test"
        self.assertEqual(self.pc.run(), "test")
        mock_start.assert_called_with()
        mock_stop.assert_called_with()
        mock_wait.assert_called_with()

    @mock.patch('%s.ProductionClass.__init__' % __name__, return_value=None)
    @mock.patch("%s.ClientClass" % __name__)
    def test_context_method1(self, mock_client, mock_init):
        class FakeDocker:
            state = "state"
            def __init__(self, *args, **kwargs):
                pass
            def __enter__(self):
                return self
            def __exit__(self, *args, **kwargs):
                pass
            def start(self):
                return "start"
            def stop(self):
                return "stop"
        mock_client.return_value = FakeDocker()
        pc = ProductionClass()
        self.assertEqual(pc.run_context(), "state")

    @mock.patch('%s.ProductionClass.__init__' % __name__, return_value=None)
    @mock.patch("%s.ClientClass.__init__" % __name__, return_value=None)
    @mock.patch("%s.ClientClass.__enter__" % __name__)
    @mock.patch("%s.ClientClass.__exit__" % __name__)
    def test_context_method2(self, mock_client_exit, mock_client_enter, mock_client_init, mock_init):
        mock_context = mock_client_enter.return_value
        mock_context.start.return_value = "start"
        mock_context.stop.return_value = "stop"
        mock_context.state = "state"
        pc = ProductionClass()
        self.assertEqual(pc.run_context(), "state")
        mock_context.start.assert_called_with()
        mock_context.stop.assert_called_with()

    @mock.patch('%s.ProductionClass.__init__' % __name__, return_value=None)
    @mock.patch("%s.ClientClass" % __name__)
    def test_context_method3(self, mock_client, mock_init):
        mock_context = mock_client.return_value.__enter__.return_value
        mock_context.start.return_value = "start"
        mock_context.stop.return_value = "stop"
        mock_context.state = "state"
        pc = ProductionClass()
        self.assertEqual(pc.run_context(), "state")
        mock_context.start.assert_called_with()
        mock_context.stop.assert_called_with()
        # Test exception
        mock_context.start.side_effect = Exception("error")
        self.assertEqual(pc.run_context(), "error")
        mock_context.start.assert_called_with()

Run

$ python -m unittest tests.test_mock
.......
----------------------------------------------------------------------
Ran 7 tests in 0.011s

OK

Real class 'ClientClass' not called

Test coverage 100% (ClientClass ignored):

$ coverage run -m unittest tests.test_mock
.......
----------------------------------------------------------------------
Ran 7 tests in 0.014s

OK

$ coverage report                         
Name                 Stmts   Miss  Cover
----------------------------------------
tests/__init__.py        0      0   100%
tests/test_mock.py     103      0   100%
----------------------------------------
TOTAL                  103      0   100%

Links

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment