def _reinit():
"""Restart the class attributes to initial values."""
InfrastructureList.infrastructure_list = {}
InfrastructureList._lock = threading.Lock()
db = DataBase(Config.DATA_DB)
if db.connect():
if db.db_type == DataBase.MONGO:
db.delete("inf_list", {})
else:
db.execute("delete from inf_list")
db.close()
def test_mongo_db(self, mongo):
client = MagicMock()
mongo.return_value = client
database = MagicMock()
client.__getitem__.return_value = database
database.client = client
table = MagicMock()
database.__getitem__.return_value = table
table.database = database
db_url = "mongodb://*****:*****@server/db_name"
db = DataBase(db_url)
self.assertTrue(db.connect())
database.collection_names.return_value = ['table1']
res = db.table_exists("test")
self.assertFalse(res)
res = db.table_exists("table1")
self.assertTrue(res)
res = db.replace('table', {}, {'id': 1, 'data': 'test1'})
self.assertTrue(res)
self.assertEqual(table.replace_one.call_args_list[0][0], ({}, {'data': 'test1', 'id': 1}, True))
table.find.return_value = [{'id': 2, 'data': 'test2', '_id': 2}]
res = db.find('table', {'id': 2}, {'data': True})
self.assertEqual(len(res), 1)
self.assertEqual(table.find.call_args_list[0][0], ({'id': 2}, {'_id': False, 'data': True}))
del_res = MagicMock()
del_res.deleted_count = 1
table.delete_many.return_value = del_res
res = db.delete('table', {'id': 1})
self.assertEqual(res, 1)
self.assertEqual(table.delete_many.call_args_list[0][0], ({'id': 1},))
db.close()
