side_effect to return values conditionally

Question

I'm mocking one of the functions similar to below. Is there a way to pass arguments as part of side_effect so that I can use the same function to (mock) load both order and product data?

def mock_load_data(self, name="order"):
    if name == "order":
       return 1
    elif name == "product":
       return 2

@patch('src.shop.order.load_data', side_effect=mock_load_data)
def test_order:

Answer 1

Pass different parameters to mock_load_data when execute it will do that.

Eg

main.py :

from order import load_data


def main(name):
    return load_data(name)

order.py :

def load_data():
    return 'real data'

test_main.py :

import unittest
from unittest.mock import patch
from main import main


def mock_load_data(name="order"):
    if name == "order":
        return 1
    elif name == "product":
        return 2


class TestMain(unittest.TestCase):
    @patch('main.load_data', side_effect=mock_load_data)
    def test_main(self, mock_load_data):
        rval1 = main('order')
        self.assertEqual(rval1, 1)
        rval2 = main('product')
        self.assertEqual(rval2, 2)


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

unit test results with coverage report:

.
----------------------------------------------------------------------
Ran 1 test in 0.001s

OK
Name                                      Stmts   Miss  Cover   Missing
-----------------------------------------------------------------------
src/stackoverflow/60572053/main.py            3      0   100%
src/stackoverflow/60572053/order.py           2      1    50%   2
src/stackoverflow/60572053/test_main.py      16      0   100%
-----------------------------------------------------------------------
TOTAL                                        21      1    95%