Python - create mock test for class method that has context manager

Question

I'm trying to write a unit test for a method of a class that has a context manager and many calls. I'm having a difficulty time understanding how to properly mock the function so that I can test the return value.
The class I am trying to mock is db. As you can see below I'm using a patch, but I'm not able to figure out how to get it to return the correct method call. I'm getting a generic mock function instead of the return value I expect.

db_class.py

import db

class Foo():
    def __init__(self):
        pass
    def method(self):
        with db.a() as a:
            b = a.b
            return b.fetch()

unit_db.py

 from mock import Mock, patch, MagicMock
 from db_class import Foo

 @patch('db_class.db')
 def test(db_mock):
     expected_result = [5,10]
     db_mock.return_value = Mock(__enter__ = db_mock,
                                 __exit___ = Mock(),
                                 b = Mock(fetch=expected_result))
 
     foo = Foo()
     result = foo.method()
     assert result == expected_result
     
 

Answer 1

Thanks to the commenters I have found a solution that works for me. The trick was to patch the correct class, in this case I wanted to patch db_class.db.a instead of db_class.db. After that, it is important to make sure that the fetch() call is a method (I think I'm getting that correct). The tricky part about this problem for me was patching the correct thing as well as dealing with the context manager which requires a bit of extra tinkering.

@patch('db_class.db.a')
def test(db_a):
    expected_result = [5,10]
    b_fetch = MagicMock()
    b_fetch.fetch.return_value = expected_result 
    db_a.return_value = Mock(b = b_fetch,
                         __enter__= db_a,
                         __exit__ =Mock())
    foo = Foo()
    result = foo.method()
    assert result == expected_result

if __name__ == "__main__":
    test()