Python Mock.Patch a function used at several places

Question

This works:

I got my program:

# module/core_functions.py

def new_input(question):
    print(question)
    value = input(">").strip()
    return value


def main_function():
    #do things
    new_input("question1")
    #do other things
    new_input("question2")
    ...

I wrote an unittest:

import unittest
from unittest.mock import patch, Mock
from module.core_functions import main_function

class MytestClass(unittest.TestCase):

    @patch('module.core_functions.new_input')
    def test_multiple_answer(self, mock_input):
        mock_input.return_value = Mock()
        mock_input.side_effect = ['Answer1', 'Answer2']
        result = main_function()
        self.assertIn('ExpectedResult', result)

This works perfectly fine (I use nose2 to run all my tests).

Now this is not working:

As my code is getting bigger and bigger, I would like to involve other people in my project and I need to modularise my functions to make modification easier and cleaner.

So I put the new_input function in a new file module/io.py and I got a new subfunctions:

# module/subfunctions.py
from module.io import new_input

def subfunction():
    # do things
    new_input("question")
    # do things
    return things

And my core program evolved to:

# module/core_functions.py
from module.io import new_input
from module.subfunctions import subfunction

def main_function():
    #do things
    new_input("question1")
    #do other things
    subfuntion()
    ...

So the problem is that the function to mock is in several places: in the main function and in some subfunctions. I can't find a way to have my unittest working (I can't predict if AnswerX will be needed in the main function or in one subfunction).

Do you have any idea how can I fix my test so it works? Thank you (I hope I've been rather clear).

Edit:

I tried something like:

@patch('module.core_functions.new_input')
@patch('module.subfunctions.new_input')
def test_multiple_answer(self, mock_input):
    mock_input.return_value = Mock()
    mock_input.side_effect = ['Answer1', 'Answer2']
    result = main_function()
    self.assertIn('ExpectedResult', result)

But I got the error: TypeError: test_multiple_answer() takes 2 positional arguments but 3 were given.

Answer 1

I just asked myself the same question, and while I was pretty certain of the answer, I wrote a stupidly small but understandable proof of concept.

Like @whats-done-is pointed out there are two ways to go about this; either specify the mock objects to inline using mock.patch()'s new argument, or add a parameter for each @patch annotation.

utils.py

def expensive_function():
    return 12345

module1.py

from utils import expensive_function
from module2 import method2

def method1():
    return {'method1': expensive_function(), 'method2': method2()}

module2.py

from utils import expensive_function

def method2():
    return expensive_function()

main.py

import unittest
from unittest.mock import patch
from module1 import method1

def mock_expensive_function():
    return 123

class TestCases(unittest.TestCase):
    def test_unpatched(self):
        self.assertEqual({'method1': 12345, 'method2': 12345}, method1())

    @patch('module1.expensive_function', new=mock_expensive_function)
    def test_method1_patched(self):
        self.assertEqual({'method1': 123, 'method2': 12345}, method1())

    @patch('module2.expensive_function', new=mock_expensive_function)
    def test_method2_patched(self):
        self.assertEqual({'method1': 12345, 'method2': 123}, method1())

    @patch('module1.expensive_function', new=mock_expensive_function)
    @patch('module2.expensive_function', new=mock_expensive_function)
    def test_both_patched_inline(self):
        self.assertEqual({'method1': 123, 'method2': 123}, method1())

    @patch('module1.expensive_function')
    @patch('module2.expensive_function')
    def test_both_patched_magicmock(self, mock_in_module2, mock_in_module1):
        mock_in_module1.return_value = mock_expensive_function()
        mock_in_module2.return_value = mock_expensive_function()
        self.assertEqual({'method1': 123, 'method2': 123}, method1())

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