Fixed database resetting bug and test typos

This commit is contained in:
VakarisZ 2020-04-01 12:15:31 +03:00
parent 483a3576d4
commit 8a479145c6
2 changed files with 12 additions and 6 deletions

View File

@ -7,9 +7,9 @@ class TestMitreApiInterface(TestCase):
def test_get_all_mitigations(self): def test_get_all_mitigations(self):
mitigations = MitreApiInterface.get_all_mitigations() mitigations = MitreApiInterface.get_all_mitigations()
self.assertTrue((len(mitigations.items()) >= 282)) self.assertIsNotNone((len(mitigations.items()) >= 282))
mitigation = next(iter(mitigations.values())) mitigation = next(iter(mitigations.values()))
self.assertEqual(mitigation['type'], "course-of-action") self.assertEqual(mitigation['type'], "course-of-action")
self.assertTrue(mitigation['name']) self.assertIsNotNone(mitigation['name'])
self.assertTrue(mitigation['description']) self.assertIsNotNone(mitigation['description'])
self.assertTrue(mitigation['external_references']) self.assertIsNotNone(mitigation['external_references'])

View File

@ -16,15 +16,21 @@ class Database(object):
@staticmethod @staticmethod
def reset_db(): def reset_db():
logger.info('Resetting database')
remove_PBA_files() remove_PBA_files()
# We can't drop system collections. # We can't drop system collections.
[mongo.db[x].drop() for x in mongo.db.collection_names() if not x.startswith('system.') [Database.drop_collection(x) for x in mongo.db.collection_names() if not x.startswith('system.')
and not AttackMitigations.COLLECTION_NAME] and not x == AttackMitigations.COLLECTION_NAME]
ConfigService.init_config() ConfigService.init_config()
AttackConfig.reset_config() AttackConfig.reset_config()
logger.info('DB was reset') logger.info('DB was reset')
return jsonify(status='OK') return jsonify(status='OK')
@staticmethod
def drop_collection(collection_name: str):
mongo.db[collection_name].drop()
logger.info("Dropped collection {}".format(collection_name))
@staticmethod @staticmethod
def init_db(): def init_db():
if not mongo.db.collection_names(): if not mongo.db.collection_names():