How can I autospec mock attributes that are None by default in python 3?

Question

Consider this code:

import unittest
from unittest.mock import patch

class Foo(object):
    def __init__(self, bar=None):
        self.bar = bar

    def methodA(self):
        print("In methodA")

    def methodB(self):
        print("In methodB")


def my_func(bar):
    foo = Foo(bar)
    if foo.bar:
        foo.methodA()
    foo.methodB()


class MyTestCase(unittest.TestCase):
    def test_my_func(self):
        bar = None

        with patch("__main__.Foo", autospec=True) as foo:
            my_func(bar)
            foo.methodB.assert_called_once_with()


if __name__ == '__main__':
    unittest.main()

The idea is fairly simple. I have a function whose behavior switches on the presence or absence of an instance attribute. I'm trying to write a unit test that verifies only certain Foo methods are executed, depending on the attribute.

Based on the mock library's patch and autospeccing docs, I thought setting autospec=True in the patch() context manager would suffice. It doesn't. The generated Mock() correctly includes mocks of methodA and methodB, but the test fails with this error:

======================================================================
ERROR: test_my_func (__main__.MyTestCase)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "so.py", line 28, in test_my_func
    my_func(bar)
  File "trash.py", line 18, in my_func
    if foo.bar:
  File "/.../python3.3/unittest/mock.py", line 549, in __getattr__
    raise AttributeError("Mock object has no attribute %r" % name)
AttributeError: Mock object has no attribute 'bar'

I'm sure I'm missing something obvious, but I can't seem to figure out what. How can I unittest my_func()?

Answer 1

Part of my problem was a misunderstanding of patch's behavior. The context manager I set up is returning an instance of the __main__.Foo mock, not the same instance used in my_func(). To put it another way, even when I was able to mock Foo properly, without autospec, I couldn't execute assert_called_once_with() over any of its methods: It wasn't the same object.

One solution is to mock the method itself. This works:

def test_my_func(self):
    bar = None
    with patch('__main__.Foo.methodB') as mock_methodB:
        my_func(bar)
        mock_methodB.assert_called_once_with()

Another method would be modifying my_func() to return foo:

def my_func(bar):
    foo = Foo(bar)
    if foo.bar:
        foo.methodA()
    foo.methodB()
    return foo

Because the function returns the mock under test, the following should work:

def test_my_func(self):
    bar = None
    with patch('__main__.Foo', spec=True, bar=None):
        foo = my_func(bar)
        assert foo.methodB.called
        assert not foo.methodA.called