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