How to supply a mock class method for python unit test?

Question

Let’s say I have a class like this.

class SomeProductionProcess(CustomCachedSingleTon):
    
    @classmethod
    def loaddata(cls):
        """
        Uses an iterator over a large file in Production for the Data pipeline.
        """
        pass

Now at test time I want to change the logic inside the loaddata() method. It would be a simple custom logic that doesn’t process large data.

How do we supply custom implementation of loaddata() at testtime using Python Mock UnitTest framework?

Asked By: Vineel

||

Source

Answer 1

Lets say you have a module named awesome.py and in it, you had:

import time

class SomeProductionProcess(CustomCachedSingleTon):

    def loaddata(self):
        time.sleep(30) # simulating a long running process
        return 2

Then your unittest where you mock loaddata could look like this:

import unittest

import awesome # your application module


class TestSomeProductionProcess(unittest.TestCase):
    """Example of direct monkey patching"""

    def test_loaddata(self):
        some_prod_proc = awesome.SomeProductionProcess()
        some_prod_proc.loaddata = lambda x: 2 # will return 2 every time called
        output = some_prod_proc.loaddata()
        expected = 2

        self.assertEqual(output, expected)

Or it could look like this:

import unittest
from mock import patch

import awesome # your application module

class TestSomeProductionProcess(unittest.TestCase):
    """Example of using the mock.patch function"""

    @patch.object(awesome.SomeProductionProcess, 'loaddata')
    def test_loaddata(self, fake_loaddata):
        fake_loaddata.return_value = 2
        some_prod_proc = awesome.SomeProductionProcess()

        output = some_prod_proc.loaddata()
        expected = 2

        self.assertEqual(output, expected)

Now when you run your test, loaddata wont take 30 seconds for those test cases.

Answered By: willnx

Answer 2

To easily mock out a class method with a structured return_value, can use unittest.mock.Mock.

from unittest.mock import Mock

mockObject = SomeProductionProcess
mockObject.loaddata = Mock(return_value=True)

EDIT:

Since you want to mock out the method with a custom implementation, you could just create a custom mock method object and swap out the original method at testing runtime.

def custom_method(*args, **kwargs):
    # do custom implementation

SomeProductionProcess.loaddata = custom_method

Answered By: Mattew Whitt

Answer 3

Here is a simple way to do it using mock

import mock


def new_loaddata(cls, *args, **kwargs):
    # Your custom testing override
    return 1


def test_SomeProductionProcess():
    with mock.patch.object(SomeProductionProcess, 'loaddata', new=new_loaddata):
        obj = SomeProductionProcess()
        obj.loaddata()  # This will call your mock method

I’d recommend using pytest instead of the unittest module if you’re able. It makes your test code a lot cleaner and reduces a lot of the boilerplate you get with unittest.TestCase-style tests.

Answered By: Brendan Abel

How to supply a mock class method for python unit test?

Question:

Answers: