title: Python - create mock test for class method that has context manager
body: <p>I'm trying to write a unit test for a method of a class function 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. </p>
  
  <p>db_class.py</p>
  
  <pre><code>import db
  
  class Foo():
      def __init__(self):
          pass
      def method(self):
          with db.a() as a:
              b = a.b
              return b.fetch()
  </code></pre>
  
  <p>unit_db.py</p>
  
  <pre><code> 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
  </code></pre>