diff --git a/monkey/monkey_island/cc/models/test_event.py b/monkey/monkey_island/cc/models/test_event.py index f36f4f562..3bc201f62 100644 --- a/monkey/monkey_island/cc/models/test_event.py +++ b/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 ) diff --git a/monkey/monkey_island/cc/models/test_finding.py b/monkey/monkey_island/cc/models/test_finding.py index 2b52553da..d111b0513 100644 --- a/monkey/monkey_island/cc/models/test_finding.py +++ b/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) diff --git a/monkey/monkey_island/cc/testing/IslandTestCase.py b/monkey/monkey_island/cc/testing/IslandTestCase.py index e894f13df..5b050684c 100644 --- a/monkey/monkey_island/cc/testing/IslandTestCase.py +++ b/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()