From 8a479145c6de9700dc63f70be3c95d681e71a592 Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Wed, 1 Apr 2020 12:15:31 +0300 Subject: [PATCH] Fixed database resetting bug and test typos --- .../cc/services/attack/test_mitre_api_interface.py | 8 ++++---- monkey/monkey_island/cc/services/database.py | 10 ++++++++-- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/monkey/monkey_island/cc/services/attack/test_mitre_api_interface.py b/monkey/monkey_island/cc/services/attack/test_mitre_api_interface.py index b896828cb..4866a6729 100644 --- a/monkey/monkey_island/cc/services/attack/test_mitre_api_interface.py +++ b/monkey/monkey_island/cc/services/attack/test_mitre_api_interface.py @@ -7,9 +7,9 @@ class TestMitreApiInterface(TestCase): def test_get_all_mitigations(self): mitigations = MitreApiInterface.get_all_mitigations() - self.assertTrue((len(mitigations.items()) >= 282)) + self.assertIsNotNone((len(mitigations.items()) >= 282)) mitigation = next(iter(mitigations.values())) self.assertEqual(mitigation['type'], "course-of-action") - self.assertTrue(mitigation['name']) - self.assertTrue(mitigation['description']) - self.assertTrue(mitigation['external_references']) + self.assertIsNotNone(mitigation['name']) + self.assertIsNotNone(mitigation['description']) + self.assertIsNotNone(mitigation['external_references']) diff --git a/monkey/monkey_island/cc/services/database.py b/monkey/monkey_island/cc/services/database.py index 5961fa70b..85812dd6e 100644 --- a/monkey/monkey_island/cc/services/database.py +++ b/monkey/monkey_island/cc/services/database.py @@ -16,15 +16,21 @@ class Database(object): @staticmethod def reset_db(): + logger.info('Resetting database') remove_PBA_files() # We can't drop system collections. - [mongo.db[x].drop() for x in mongo.db.collection_names() if not x.startswith('system.') - and not AttackMitigations.COLLECTION_NAME] + [Database.drop_collection(x) for x in mongo.db.collection_names() if not x.startswith('system.') + and not x == AttackMitigations.COLLECTION_NAME] ConfigService.init_config() AttackConfig.reset_config() logger.info('DB was reset') return jsonify(status='OK') + @staticmethod + def drop_collection(collection_name: str): + mongo.db[collection_name].drop() + logger.info("Dropped collection {}".format(collection_name)) + @staticmethod def init_db(): if not mongo.db.collection_names():