def _save_data_to_db(db_url, inf_list, inf_id=None):
db = DataBase(db_url)
if db.connect():
infs_to_save = inf_list
if inf_id:
infs_to_save = {inf_id: inf_list[inf_id]}
for inf in infs_to_save.values():
data = inf.serialize()
if db.db_type == DataBase.MONGO:
res = db.replace("inf_list", {"id": inf.id}, {
"id": inf.id,
"deleted": int(inf.deleted),
"data": data,
"date": time.time()
})
else:
res = db.execute(
"replace into inf_list (id, deleted, data, date) values (%s, %s, %s, now())",
(inf.id, int(inf.deleted), data))
db.close()
return res
else:
InfrastructureList.logger.error(
"ERROR connecting with the database!.")
return None
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()