Unit testing - how to set values for mock object Python

Question

I'm a beginner to unit test in Python - I'm trying to test a function which takes in a class object and returns a set of strings. Input is a class object with some_attr Ps, p is also a class object with attribute 'name'.

def my_func(input):
    results = set(map(lambda p: p.name, input.Ps))
    return results

I'm trying to use Mock module by setting the input as a mock object.

@dataclass
class TestClass:
    name: str

@patch('some related module')
def Test_my_func(TestCase, mock_input):
    import my_func
    mock_input = MagicMock()
    p1 = TestClass('name1')
    p2 = TestClass('name2')
    mock_input.Ps.return_value = tuple([p1, p2])
    expectedResult = set(['name1', 'name2'])

    self.assertEqual(my_func(mock_input), expectedResult)

The issue I'm having is that my_func(mock_input) returns set() and did not pick up the mock_input.Ps.return_value = tuple([p1, p2]) as I expected. Can someone please shed some light on this? Much appreciated.

Error message: AssertionError: Items in the second set but not the first

Answer 1

You can just set mock_input.Ps to [p1, p2] since Ps is an attribute, here is the code snippet that works on my environment with PyTest,

from dataclasses import dataclass

def my_func(input):
    results = set(map(lambda p: p.name, input.Ps))
    return results

@dataclass
class TestClass:
    name: str

# test --------------------
from unittest.mock import MagicMock

def test_my_func():
    # import my_func
    mock_input = MagicMock()
    p1 = TestClass('name1')
    p2 = TestClass('name2')
    mock_input.Ps = [p1, p2]
    expectedResult = set(['name1', 'name2'])

    assert my_func(mock_input) == expectedResult

Results:

tests/unit/test_foo.py::test_my_func PASSED