How to override __getitem__ on a MagicMock subclass

Question

We're mocking a 3rd party module that has a large api with many nested instances. Its various objects are often access via attribute, get(), or dict indexing []. Having a flexible mock that treat all 3 as simple attribute access would be very handy.

class FlexiMock(MagicMock):
    def get(self, x, default=None): 
        return getattr(self, x, default)

    def __getitem__(self, x):
        return getattr(self, x)


t = FlexiMock(status=1)
assert t.status == 1
assert t.get('status') == 1
assert t['status'] == 1, f"{t['status']} vs 1"

Results in:

assert t['status'] == 1, f"{t['status']} vs 1"
AssertionError: <FlexiMock name='mock.__getitem__()' id='140549581528368'> vs 1

It seems rather straightforward so I'm not sure what I'm doing wrong.

Answer 1

You can get around it by setting the mock for __getitem__ while initializing the FlexiMock like so:

from unittest.mock import MagicMock


class FlexiMock(MagicMock):
    def __init__(self, *args, **kwargs):
        super(FlexiMock, self).__init__(*args, **kwargs)
        self.__getitem__ = lambda obj, item: getattr(obj, item) # set the mock for `[...]`

    def get(self, x, default=None):
        return getattr(self, x, default)


t = FlexiMock(status=1)
assert t.status == 1
assert t.get('status') == 1
assert t['status'] == 1, f"{t['status']} vs 1"

Which runs as expected.

Answer 2

The classes in unittest.mock are not designed to play nicely with user-defined subclasses. In particular, unittest.mock.MagicMock does something kind of weird. When a new instance is created, it gets its own dynamically-generated subclass that overrides almost every single magic method there is . Your __getitem__ is overridden by the __getitem__ in the dynamically-generated subclass:

In [1]: from unittest.mock import MagicMock

In [2]: class FlexiMock(MagicMock):
   ...:     def get(self, x, default=None):
   ...:         return getattr(self, x, default)
   ...: 
   ...:     def __getitem__(self, x):
   ...:         return getattr(self, x)
   ...: 

In [3]: t = FlexiMock(status=1)

In [4]: type(t).__mro__
Out[4]: 
(unittest.mock.FlexiMock,
 __main__.FlexiMock,
 unittest.mock.MagicMock,
 unittest.mock.MagicMixin,
 unittest.mock.Mock,
 unittest.mock.CallableMixin,
 unittest.mock.NonCallableMock,
 unittest.mock.Base,
 object)

In [5]: type(t).__dict__                                                                                                                                                                                                         
Out[5]: 
mappingproxy({'__doc__': None,
              '__module__': 'unittest.mock',
              '__pos__': <unittest.mock.MagicProxy at 0x7fcd4506cb80>,
              '__rsub__': <unittest.mock.MagicProxy at 0x7fcd3ffe4d60>,
              '__int__': <unittest.mock.MagicProxy at 0x7fcd3ff40130>,
              '__ifloordiv__': <unittest.mock.MagicProxy at 0x7fcd44084fa0>,
              '__trunc__': <unittest.mock.MagicProxy at 0x7fcd440840a0>,
              '__complex__': <unittest.mock.MagicProxy at 0x7fcd3ffdf160>,
              '__rlshift__': <unittest.mock.MagicProxy at 0x7fcd3ffdf8e0>,
              '__sub__': <unittest.mock.MagicProxy at 0x7fcd3ffdf550>,
              '__delitem__': <unittest.mock.MagicProxy at 0x7fcd3ffdf520>,
              '__matmul__': <unittest.mock.MagicProxy at 0x7fcd3ffd71f0>,
              '__imatmul__': <unittest.mock.MagicProxy at 0x7fcd3ffd95e0>,
              '__rshift__': <unittest.mock.MagicProxy at 0x7fcd3ffd9610>,
              '__gt__': <unittest.mock.MagicProxy at 0x7fcd3ffd9d60>,
              '__round__': <unittest.mock.MagicProxy at 0x7fcd3ffd9ee0>,
              '__idiv__': <unittest.mock.MagicProxy at 0x7fcd3ffd91f0>,
              '__neg__': <unittest.mock.MagicProxy at 0x7fcd3ffc9730>,
              '__fspath__': <unittest.mock.MagicProxy at 0x7fcd3ffc92b0>,
              '__imod__': <unittest.mock.MagicProxy at 0x7fcd4404a5b0>,
              '__rdiv__': <unittest.mock.MagicProxy at 0x7fcd4404a6a0>,
              '__add__': <unittest.mock.MagicProxy at 0x7fcd4404a250>,
              '__aiter__': <unittest.mock.MagicProxy at 0x7fcd4404af70>,
              '__floordiv__': <unittest.mock.MagicProxy at 0x7fcd4404a400>,
              '__xor__': <unittest.mock.MagicProxy at 0x7fcd4404aaf0>,
              '__next__': <unittest.mock.MagicProxy at 0x7fcd4404a850>,
              '__rrshift__': <unittest.mock.MagicProxy at 0x7fcd4404a8b0>,
              '__truediv__': <unittest.mock.MagicProxy at 0x7fcd4404afa0>,
              '__imul__': <unittest.mock.MagicProxy at 0x7fcd4404a070>,
              '__eq__': <unittest.mock.MagicProxy at 0x7fcd3ff8bc70>,
              '__ior__': <unittest.mock.MagicProxy at 0x7fcd3ff8b580>,
              '__enter__': <unittest.mock.MagicProxy at 0x7fcd3ff8b280>,
              '__abs__': <unittest.mock.MagicProxy at 0x7fcd3ff8bdf0>,
              '__le__': <unittest.mock.MagicProxy at 0x7fcd3ff8b4c0>,
              '__rdivmod__': <unittest.mock.MagicProxy at 0x7fcd3ff8ba30>,
              '__getitem__': <unittest.mock.MagicProxy at 0x7fcd3ff8bc10>,
              '__sizeof__': <unittest.mock.MagicProxy at 0x7fcd3ff8b0a0>,
              '__bool__': <unittest.mock.MagicProxy at 0x7fcd3ff8b1f0>,
              '__or__': <unittest.mock.MagicProxy at 0x7fcd3ff8b4f0>,
              '__radd__': <unittest.mock.MagicProxy at 0x7fcd3ff8b730>,
              '__lt__': <unittest.mock.MagicProxy at 0x7fcd3ff8b070>,
              '__rtruediv__': <unittest.mock.MagicProxy at 0x7fcd3ff8b3d0>,
              '__rmatmul__': <unittest.mock.MagicProxy at 0x7fcd3ff8bb20>,
              '__mul__': <unittest.mock.MagicProxy at 0x7fcd3ff8bb80>,
              '__divmod__': <unittest.mock.MagicProxy at 0x7fcd3ff8ba00>,
              '__float__': <unittest.mock.MagicProxy at 0x7fcd3ff8b6a0>,
              '__and__': <unittest.mock.MagicProxy at 0x7fcd3ff8bd00>,
              '__len__': <unittest.mock.MagicProxy at 0x7fcd3ff8bca0>,
              '__ixor__': <unittest.mock.MagicProxy at 0x7fcd3ff8b5e0>,
              '__index__': <unittest.mock.MagicProxy at 0x7fcd3ff8b670>,
              '__aexit__': <unittest.mock.MagicProxy at 0x7fcd3ff8bdc0>,
              '__lshift__': <unittest.mock.MagicProxy at 0x7fcd45062340>,
              '__rfloordiv__': <unittest.mock.MagicProxy at 0x7fcd450628b0>,
              '__ge__': <unittest.mock.MagicProxy at 0x7fcd450622b0>,
              '__irshift__': <unittest.mock.MagicProxy at 0x7fcd45062a90>,
              '__iadd__': <unittest.mock.MagicProxy at 0x7fcd440615b0>,
              '__iter__': <unittest.mock.MagicProxy at 0x7fcd44061790>,
              '__ne__': <unittest.mock.MagicProxy at 0x7fcd440613a0>,
              '__ilshift__': <unittest.mock.MagicProxy at 0x7fcd44061910>,
              '__rpow__': <unittest.mock.MagicProxy at 0x7fcd440614f0>,
              '__ror__': <unittest.mock.MagicProxy at 0x7fcd44061880>,
              '__div__': <unittest.mock.MagicProxy at 0x7fcd440618e0>,
              '__isub__': <unittest.mock.MagicProxy at 0x7fcd44061670>,
              '__hash__': <unittest.mock.MagicProxy at 0x7fcd440617c0>,
              '__rand__': <unittest.mock.MagicProxy at 0x7fcd44061640>,
              '__setitem__': <unittest.mock.MagicProxy at 0x7fcd44061610>,
              '__iand__': <unittest.mock.MagicProxy at 0x7fcd440618b0>,
              '__invert__': <unittest.mock.MagicProxy at 0x7fcd44061760>,
              '__ipow__': <unittest.mock.MagicProxy at 0x7fcd440619d0>,
              '__str__': <unittest.mock.MagicProxy at 0x7fcd44061a30>,
              '__contains__': <unittest.mock.MagicProxy at 0x7fcd44061a90>,
              '__ceil__': <unittest.mock.MagicProxy at 0x7fcd44061040>,
              '__anext__': <unittest.mock.MagicProxy at 0x7fcd44061100>,
              '__exit__': <unittest.mock.MagicProxy at 0x7fcd440610d0>,
              '__floor__': <unittest.mock.MagicProxy at 0x7fcd44061490>,
              '__rmul__': <unittest.mock.MagicProxy at 0x7fcd440613d0>,
              '__mod__': <unittest.mock.MagicProxy at 0x7fcd44061730>,
              '__aenter__': <unittest.mock.MagicProxy at 0x7fcd44061af0>,
              '__pow__': <unittest.mock.MagicProxy at 0x7fcd44061b20>,
              '__rmod__': <unittest.mock.MagicProxy at 0x7fcd44061b80>,
              '__rxor__': <unittest.mock.MagicProxy at 0x7fcd44061be0>,
              '__itruediv__': <unittest.mock.MagicProxy at 0x7fcd44061c40>})

Here, unittest.mock.FlexiMock is the dynamically-generated subclass, with a dict full of magic method overrides.

Instead of trying to subclass unittest.mock classes, the supported API for overriding methods is to just make a regular mock and set the methods on the instance. Note that due to additional unittest.mock weirdness, magic methods should take self , but other methods should not take self :

def make_fleximock(*args, **kwargs):
    mock = MagicMock(*args, **kwargs)
    mock.get = lambda x, default=None: getattr(mock, x, default)
    mock.__getitem__ = lambda self, x: getattr(self, x)
    return mock