Python Mock Patch Two Functions that are similar

Question

Let's say I have a function that has two similar function calls:

def foo():
  test = one_func("SELECT * from Users;")
  test1 = one_func("SELECT * from Addresses;")
  return test, test1

How do I patch each of these function contents? Here's my attempt:

@patch('one_func')
def test_foo(self, mock_one_func):
     mock_one_func.return_value = one_func("SELECT * from TestUsers;")
     mock_one_func.return_value = one_func("SELECT * from TestAddresses;")

But I think this method patches one_func as a whole per function. Which results to:

def foo():
  test = one_func("SELECT * from TestUsers;")
  test1 = one_func("SELECT * from TestUsers;")
  return test, test1

Then on the next line

def foo():
  test = one_func("SELECT * from TestAddresses;")
  test1 = one_func("SELECT * from TestAddresses;")
  return test, test1

What I want to happen in the patched function is.

def foo():
  test = one_func("SELECT * from TestUsers;")
  test1 = one_func("SELECT * from TestAddresses;")
  return test, test1

score 0 · Answer 1 · answered Jun 22 '21 at 13:14

The way to achieve what you need is using side_effect instead of return_value. side_effect can be many things. If it is an Exception class or object, it will raise that exception whenever the patched method is called. If it is a list of values, will return each value in sequence for each method call. If it is a function, will execute that function with the parameters from each mocked method call.

Here is a working example, where I show you how to use a list of values on side_effect and a side_effect function. What's nice about using a function is that you can make it return specific values depending on the arguments of the patched method call.

from mock import patch

import unittest

class MyClass(object):
    def one_func(self, query):
        return ''

    def foo(self):
        test = self.one_func("SELECT * from Users;")
        test1 = self.one_func("SELECT * from Addresses;")
        return test, test1

class Test(unittest.TestCase):
    @patch.object(MyClass, 'one_func')
    def test_foo(self, one_func_mock):
        # side_effect can be a list of responses that will be returned in
        # subsequent calls
        one_func_mock.side_effect = ['users', 'addresses']
        self.assertEqual(('users', 'addresses'), MyClass().foo())

        # side_effect can also be a method which will return different mock
        # responses depending on args:
        def side_effect(args):
            if args == "SELECT * from Users;":
                return 'other users'

            if args == "SELECT * from Addresses;":
                return 'other addresses'
        
        one_func_mock.side_effect = side_effect
        self.assertEqual(('other users', 'other addresses'), MyClass().foo())

unittest.main()

Python Mock Patch Two Functions that are similar

1 Answers1