Added finding db cleaning for the UTs in IslandTestCase

monkey/monkey_island/cc/models/test_event.py

@@ -8,7 +8,7 @@ from monkey_island.cc.testing.IslandTestCase import IslandTestCase
 class TestEvent(IslandTestCase):
     def test_create_event(self):
         self.fail_if_not_testing_env()
-        self.clean_monkey_db()
+        self.clean_finding_db()
 
         with self.assertRaises(ValidationError):
             Event.create_event(
@@ -27,5 +27,5 @@ class TestEvent(IslandTestCase):
         _ = Event.create_event(
             title="skjs",
             message="bla bla",
-            event_type=EVENT_TYPE_ISLAND  # Unknown event type
+            event_type=EVENT_TYPE_ISLAND
         )

monkey/monkey_island/cc/models/test_finding.py

@@ -1,11 +1,8 @@
-from datetime import datetime
-
 from mongoengine import ValidationError
 
-from common.data.zero_trust_consts import TEST_SEGMENTATION, STATUS_CONCLUSIVE, NETWORKS, EVENT_TYPE_MONKEY_NETWORK
-from finding import Finding, UnknownTest
+from common.data.zero_trust_consts import *
+from finding import Finding
 from monkey_island.cc.models.event import Event
-
 from monkey_island.cc.testing.IslandTestCase import IslandTestCase
 
 
@@ -19,17 +16,17 @@ class TestFinding(IslandTestCase):
     """
     def test_save_finding_validation(self):
         self.fail_if_not_testing_env()
-        self.clean_monkey_db()
+        self.clean_finding_db()
 
         with self.assertRaises(ValidationError):
-            _ = Finding.save_finding(test="bla bla", status="Conclusive", events=[])
+            _ = Finding.save_finding(test="bla bla", status=STATUS_CONCLUSIVE, events=[])
 
         with self.assertRaises(ValidationError):
             _ = Finding.save_finding(test=TEST_SEGMENTATION, status="bla bla", events=[])
 
     def test_save_finding_sanity(self):
         self.fail_if_not_testing_env()
-        self.clean_monkey_db()
+        self.clean_finding_db()
 
         self.assertEquals(len(Finding.objects(test=TEST_SEGMENTATION)), 0)
 

monkey/monkey_island/cc/testing/IslandTestCase.py

@@ -1,6 +1,7 @@
 import unittest
 from monkey_island.cc.environment.environment import env
 from monkey_island.cc.models import Monkey
+from monkey_island.cc.models.finding import Finding
 
 
 class IslandTestCase(unittest.TestCase):
@@ -10,3 +11,7 @@ class IslandTestCase(unittest.TestCase):
     @staticmethod
     def clean_monkey_db():
         Monkey.objects().delete()
+
+    @staticmethod
+    def clean_finding_db():
+        Finding.objects().delete()