我在类中编写了一个方法,试图从数据库中读取一些数据。我还用装饰器装饰了这个方法,它修复了与数据库的连接是否打开,如果没有,则重新启动它。
class HiveConnection(object):
def __init__(self, host, user, password):
"""Instantiate a HiveConnector object."""
self.host = host
self.user = user
self.password = password
self.port = 10000
self.auth_mechanism = 'PLAIN'
self._connection = self._connect()
def _connect(self):
"""Start the connection to database."""
try:
return connect(host=self.host, port=self.port,
user=self.user, password=self.password,
auth_mechanism=self.auth_mechanism)
except TTransportException as error:
print('Failed attempt to connect')
self._connect()
def _disconnect(self):
"""Close connection to database."""
self._connection.close()
def hadoop_connection_handler(function):
"""Start a database connection if not already open."""
@wraps(function)
def wrapper(inst, *args, **kwargs):
if not inst._connection:
inst._connect()
return function(inst, *args, **kwargs)
return wrapper
@hadoop_connection_handler
def read(self, query):
"""Execute a query to pull the data.
Args:
query: [str] Query to pull the data.
Returns:
A list of namedtuple (`Row`).
"""
columns = self._columns(query)
cursor = self._connection.cursor()
cursor.execute(query)
Record = namedtuple("Record", columns)
data = map(Record._make, cursor.fetchall())
cursor.close()
return data
现在我想写一个单元测试,以确保它确实有效。
from unittest.mock import patch, MagicMock
from nose.tools import assert_equal, raises
from services.db_connections import HiveConnection
class TestHiveConnection:
"""Integration test suite for HiveConnection class."""
def setUp(self):
self.hive = HiveConnection(user='username', password='password', host='myhost.net')
def test_reconnect(self):
"""If the connection drops, the object should be able to establish a
new connection.
"""
query = 'SELECT * FROM database.table 1'
self.hive._connect = MagicMock()
self.hive._disconnect()
self.hive.read(query)
assert_equal(self.hive._connect.called, True)
以上测试总是失败。 self.hive._connect.called
实际上与False相等。我认为这是因为在装饰器中调用了connect()
方法。我应该如何更改测试以解决这个问题?