Mock the return value of a function derived from an input parameter

Question

I want to mock the return value of a function that derives from an input parameter (a).
Here is the my code looks like.

 def load_data(a, b, c):
     data_source = a.get_function(b, c)
     ...
     return r_1, r_2, r_3

This is what I tried and didn't work. I was unable to find any sources that mock the return of this type of function.

@classmethod
def setUpClass(cls):
    cls.a = A.create()
    cls.data_needed_return = read_file()

def test_function(self):
    mocked = mock.Mock()
    self.a.get_function.return_value = self.data_needed_return

    import python_file
    data_source = python_file.load_data(None, None, None)

Can anyone help on this?

Answer 1

Not sure if I understood your question and problem correctly, but are you looking for something like the following?

from unittest.mock import Mock

def load_data(a, b, c):
    data_source = a.get_function(b, c)
    return data_source

a = Mock()
a.get_function.return_value = "Some data"

x = load_data(a, 2, 3)

print(x)