How do I patch a python dict such that it returns mutliple values?

1.3k views Asked by At

Python's unittest.mock allows me to patch a function such that it returns multiple values:

m = Mock(side_effect=["myName", 100, 200])

Calling m() multiple times will then return "myName", 100 and finally 200.

I can also patch a dict with patch.dict to return a mocked value, but what I am after is:

with DictMock(d, return_values=(('a',1), ('a', 2))) as d:
     assert d['a'] == 1
     assert d['a'] == 2

I have created a class that does that:

In [19]: d_
Out[19]: {'a': 1, 'b': 1}

In [20]: class MockDict:
    ...:     
    ...:     def __init__(self, in_dict, values=(), clear=False):
    ...:         self.in_dict = in_dict
    ...:         keys, values = zip(*values)
    ...:         self.keys = list(keys)
    ...:         self.values = list(values)
    ...:     
    ...:     def __enter__(self):
    ...:         self.orig_dict = self.in_dict.copy()
    ...:         return self
    ...: 
    ...:     def __exit__(self, *args):
    ...:         
    ...:         return False
    ...:     
    ...:     def __getitem__(self, key):
    ...:         try:
    ...:             idx = self.keys.index(key)
    ...:             self.keys.pop(idx)
    ...:             val = self.values.pop(idx)
    ...:             return val
    ...:         except ValueError:
    ...:             raise KeyError(key)
    ...:         

In [21]: with MockDict(d_, values=(("a",4),("a",5))) as d:
    ...:     print(d["a"])
    ...:     print(d["a"])
    ...:     
4
5

I know my class is limited in a way that is does not copy the original key value pairs ... this can be improved of course.

Can you suggest a better way to achieve this?

update: explaining my motivation

I need to test the following snippet:

    code = self.instance.state['Code']
    while code != status:
        time.sleep(3)
        self.instance.reload()
        code = self.instance.state['Code']

Hence, I have to call instance.state['Code'] twice, to go through the function code and reach self.instance.reload()

0

There are 0 answers