How to unit test nested for loop function in Python

Question

I'm trying to unit test a function with nested for loop using the Mock module in Python:

main.py:

@dataclass
Helper_class():
  Attr1: Set[some_object]
  Attr2: Set[str]

def fetch_rule_for(api, id):
    return some rule_objects based on id

def my_func(list_ids):
    api = create some api connection
    result_dict = dict()
    for id in list_ids:
      Attr1_set = set()
      Attr2_set = set()
      for rule in fetch_rule_for(api, id):
        if any(a.name == 'some_action_name' for a in rule.actions):
          Attr2_set.add(rule.attr1)
          Attr1_set.add(rule)
      result_dict[id] = Helper_class(Attr1_set, Attr2_set)

    return result_dict
      

Rule is a class object with method actions, action is another class object with attribute name.

Two questions I'm struggling with:

(1) How do I pacth the return value of function fetch_rule_for(api, id) where the return is a complicated class object?

(2) How do I deal with the for loops in Python using unit test? I have seen mock.call_count mentioned but can someone please get into more details or point me to relevant resources?

I'm new to unit testing in Python so any help is much appreicated!

Answer 1

If you want to patch the return value of a class just instantiate the class and manually set any necessary class variables then it's just a matter of using configure_mock(return_value=class_object)

Here is an example with requests.Response:

    resp = Response()
    resp.status_code = 200
    resp._content = json.dumps({'token': 'xyz123'})
    mock_post.configure_mock(return_value=resp)

In the case of the for loop you likely want to mock any returns contained within using side_effect

@mock.patch('your_class.fetch_rule_for', side_effect=iter([1,2,3]))