def tearDown(self):
'''
清除测试数据
:return:
'''
with DBSession(self.database.Session) as session:
results = session.query(ItemTable).filter(
ItemTable.stock_market == 'insert').all()
map(session.delete, results)
def setUp(self):
'''
连接测试数据库
:return:
'''
self.database = InitDB('test_database.yaml')
InitDB.Base.metadata.create_all(self.database.engine)
with DBSession(self.database.Session) as session:
results = session.query(ItemTable).filter(
ItemTable.stock_market == 'insert').all()
map(session.delete, results)
def update(self):
'''
用当前的对象更新数据
:return:
'''
with DBSession(self.database.Session) as session:
mapped_values = {}
for field_name, field_type in ItemTable.__dict__.iteritems():
is_column = isinstance(field_type, InstrumentedAttribute)
if is_column:
mapped_values[field_name] = getattr(self, field_name)
session.query(ItemTable).filter(ItemTable.stock_market == self.
stock_market).update(mapped_values)
def test_insert(self):
value_dict = {
'stock_market': 'insert',
'stock_market_link': 'test1',
'scrape_time': '2016-10-18 00:00:00'
}
obj = ItemTable(database=self.database, **value_dict)
obj.insert()
# 测试1：插入之后再取出数据，验证数据项的个数和内容是否一致
with DBSession(self.database.Session) as session:
results = session.query(ItemTable).filter(
ItemTable.stock_market == 'insert').all()
assert_equal(len(results), 1)
result = results.pop()
for key, value in value_dict.iteritems():
real_value = getattr(result, key)
if key == 'scrape_time':
real_value = real_value.strftime("%Y-%m-%d %H:%M:%S")
assert_equal(real_value, value)
# 测试2：测试update主键，验证数据是否被更新
value_dict = {
'stock_market': 'insert&update',
'stock_market_link': 'test1',
'scrape_time': time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
}
obj = ItemTable(database=self.database, **value_dict)
obj.insert()
with DBSession(self.database.Session) as session:
results = session.query(ItemTable).filter(
ItemTable.stock_market == 'insert&update').all()
assert_equal(len(results), 1)
result = results.pop()
for key, value in value_dict.iteritems():
real_value = getattr(result, key)
if key == 'scrape_time':
real_value = real_value.strftime("%Y-%m-%d %H:%M:%S")
assert_equal(real_value, value)
def test_update(self):
value_dict = {
'stock_market': 'update',
'stock_market_link': 'test1',
'scrape_time': time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
}
obj = ItemTable(database=self.database, **value_dict)
obj.update()
with DBSession(self.database.Session) as session:
results = session.query(ItemTable).filter(
ItemTable.stock_market == 'update').all()
assert_equal(len(results), 1)
result = results.pop()
for key, value in value_dict.iteritems():
real_value = getattr(result, key)
if key == 'scrape_time':
real_value = real_value.strftime("%Y-%m-%d %H:%M:%S")
assert_equal(real_value, value)
def insert(self):
'''
插入数据，如果数据存在就更新
:return:
'''
with DBSession(self.database.Session) as session:
try:
session.add(self)
session.flush()
except IntegrityError, e:
logger.info(e.message)
# 如果该主键值已经存在，就更新该数据
if 'Duplicate' in e.message:
session.rollback()
logger.info('Records exists')
self.update()
# 如果是其他的问题，就把异常继续向上抛
else:
raise e
def test_DBSession():
# 建立mock对象
Session = Mock()
mock_session = Mock()
mock_session.commit.return_value = None
mock_session.rollback.return_value = None
mock_session.close.return_value = None
Session.return_value = mock_session
# 测试1：正常调用DBSession
# 检查commit、close的调用状态
Session.assert_not_called()
mock_session.commit.assert_not_called()
mock_session.close.assert_not_called()
from stock_price_crawler.database import DBSession
with DBSession(Session) as session:
assert_equal(mock_session, session)
Session.assert_called()
mock_session.commit.assert_called()
mock_session.close.assert_called()
# close的状态清空
mock_session.close.reset_mock()
# 测试2：不正常调用DBSession
# 检查rollback、close的调用状态
mock_session.rollback.assert_not_called()
mock_session.close.assert_not_called()
def _nested_func(Session):
with DBSession(Session) as session:
# with语句抛出异常
raise Exception('Boom!')
assert_raises(Exception, _nested_func, Session)
mock_session.rollback.assert_called()
mock_session.close.assert_called()
def _nested_func(Session):
with DBSession(Session) as session:
# with语句抛出异常
raise Exception('Boom!')