lundi 25 mai 2015

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

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.

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 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

Aucun commentaire:

Enregistrer un commentaire