From d4dc42adb55f32b74ee2ec27ddbee9d8a94945b4 Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Tue, 19 Jan 2021 15:51:18 +0200 Subject: [PATCH] Removed the need to change server_config.json just to run tests. --- .travis.yml | 6 ----- .../cc/environment/environment_singleton.py | 4 +--- .../cc/environment/set_server_config.py | 2 +- monkey/monkey_island/cc/models/__init__.py | 11 +++------ monkey/monkey_island/cc/models/test_monkey.py | 21 ---------------- .../cc/models/zero_trust/test_event.py | 3 --- .../cc/models/zero_trust/test_finding.py | 7 ------ .../cc/services/edge/test_displayed_edge.py | 1 - .../cc/services/edge/test_edge.py | 3 --- .../cc/services/reporting/test_pth_report.py | 6 ----- .../test_environment.py | 3 --- .../test_system_info_telemetry_dispatcher.py | 3 --- .../zero_trust_checks/test_segmentation.py | 3 --- .../cc/testing/IslandTestCase.py | 24 ++++++++++++++++--- 14 files changed, 26 insertions(+), 71 deletions(-) diff --git a/.travis.yml b/.travis.yml index 5d638de04..8c32482f7 100644 --- a/.travis.yml +++ b/.travis.yml @@ -53,12 +53,6 @@ install: # print hugo version (useful for debugging documentation build errors) - hugo version -before_script: -# Set the server config to `testing`. This is required for for the UTs to pass. -- pushd /home/travis/build/guardicore/monkey/monkey -- python monkey_island/cc/environment/set_server_config.py testing -- popd - script: # Check Python code ## Check syntax errors and fail the build if any are found. diff --git a/monkey/monkey_island/cc/environment/environment_singleton.py b/monkey/monkey_island/cc/environment/environment_singleton.py index 121c7001d..6b98d0b7c 100644 --- a/monkey/monkey_island/cc/environment/environment_singleton.py +++ b/monkey/monkey_island/cc/environment/environment_singleton.py @@ -10,13 +10,11 @@ logger = logging.getLogger(__name__) AWS = 'aws' STANDARD = 'standard' PASSWORD = 'password' -TESTING = 'testing' ENV_DICT = { STANDARD: standard.StandardEnvironment, AWS: aws.AwsEnvironment, - PASSWORD: password.PasswordEnvironment, - TESTING: testing.TestingEnvironment + PASSWORD: password.PasswordEnvironment } env = None diff --git a/monkey/monkey_island/cc/environment/set_server_config.py b/monkey/monkey_island/cc/environment/set_server_config.py index d31dd6916..9d7c036b2 100644 --- a/monkey/monkey_island/cc/environment/set_server_config.py +++ b/monkey/monkey_island/cc/environment/set_server_config.py @@ -42,7 +42,7 @@ def main(): def parse_args(): parser = argparse.ArgumentParser() - parser.add_argument("server_config", choices=["standard", "testing", "password"]) + parser.add_argument("server_config", choices=["standard", "password"]) args = parser.parse_args() return args diff --git a/monkey/monkey_island/cc/models/__init__.py b/monkey/monkey_island/cc/models/__init__.py index 9d0114b78..87626c448 100644 --- a/monkey/monkey_island/cc/models/__init__.py +++ b/monkey/monkey_island/cc/models/__init__.py @@ -10,11 +10,6 @@ from .monkey import Monkey # noqa: F401 from .monkey_ttl import MonkeyTtl # noqa: F401 from .pba_results import PbaResults # noqa: F401 -# This section sets up the DB connection according to the environment. -# If testing, use mongomock which only emulates mongo. for more information, see -# http://docs.mongoengine.org/guide/mongomock.html . -# Otherwise, use an actual mongod instance with connection parameters supplied by env. -if env_singleton.env.testing: # See monkey_island.cc.environment.testing - connect('mongoenginetest', host='mongomock://localhost') -else: - connect(db=env_singleton.env.mongo_db_name, host=env_singleton.env.mongo_db_host, port=env_singleton.env.mongo_db_port) +connect(db=env_singleton.env.mongo_db_name, + host=env_singleton.env.mongo_db_host, + port=env_singleton.env.mongo_db_port) diff --git a/monkey/monkey_island/cc/models/test_monkey.py b/monkey/monkey_island/cc/models/test_monkey.py index b2bba9aa0..5acf8220f 100644 --- a/monkey/monkey_island/cc/models/test_monkey.py +++ b/monkey/monkey_island/cc/models/test_monkey.py @@ -22,9 +22,6 @@ class TestMonkey(IslandTestCase): """ def test_is_dead(self): - self.fail_if_not_testing_env() - self.clean_monkey_db() - # Arrange alive_monkey_ttl = MonkeyTtl.create_ttl_expire_in(30) alive_monkey_ttl.save() @@ -52,9 +49,6 @@ class TestMonkey(IslandTestCase): self.assertFalse(alive_monkey.is_dead()) def test_ttl_renewal(self): - self.fail_if_not_testing_env() - self.clean_monkey_db() - # Arrange monkey = Monkey(guid=str(uuid.uuid4())) monkey.save() @@ -65,9 +59,6 @@ class TestMonkey(IslandTestCase): self.assertIsNotNone(monkey.ttl_ref) def test_get_single_monkey_by_id(self): - self.fail_if_not_testing_env() - self.clean_monkey_db() - # Arrange a_monkey = Monkey(guid=str(uuid.uuid4())) a_monkey.save() @@ -81,9 +72,6 @@ class TestMonkey(IslandTestCase): _ = Monkey.get_single_monkey_by_id("abcdefabcdefabcdefabcdef") def test_get_os(self): - self.fail_if_not_testing_env() - self.clean_monkey_db() - linux_monkey = Monkey(guid=str(uuid.uuid4()), description="Linux shay-Virtual-Machine 4.15.0-50-generic #54-Ubuntu") windows_monkey = Monkey(guid=str(uuid.uuid4()), @@ -99,9 +87,6 @@ class TestMonkey(IslandTestCase): self.assertEqual(1, len([m for m in Monkey.objects() if m.get_os() == "unknown"])) def test_get_tunneled_monkeys(self): - self.fail_if_not_testing_env() - self.clean_monkey_db() - linux_monkey = Monkey(guid=str(uuid.uuid4()), description="Linux shay-Virtual-Machine") windows_monkey = Monkey(guid=str(uuid.uuid4()), @@ -121,9 +106,6 @@ class TestMonkey(IslandTestCase): self.assertTrue(test, "Tunneling test") def test_get_label_by_id(self): - self.fail_if_not_testing_env() - self.clean_monkey_db() - hostname_example = "a_hostname" ip_example = "1.1.1.1" linux_monkey = Monkey(guid=str(uuid.uuid4()), @@ -169,9 +151,6 @@ class TestMonkey(IslandTestCase): self.assertEqual(cache_info_after_query_3.misses, 2) def test_is_monkey(self): - self.fail_if_not_testing_env() - self.clean_monkey_db() - a_monkey = Monkey(guid=str(uuid.uuid4())) a_monkey.save() diff --git a/monkey/monkey_island/cc/models/zero_trust/test_event.py b/monkey/monkey_island/cc/models/zero_trust/test_event.py index 4699dd829..a84268e40 100644 --- a/monkey/monkey_island/cc/models/zero_trust/test_event.py +++ b/monkey/monkey_island/cc/models/zero_trust/test_event.py @@ -7,9 +7,6 @@ from monkey_island.cc.testing.IslandTestCase import IslandTestCase class TestEvent(IslandTestCase): def test_create_event(self): - self.fail_if_not_testing_env() - self.clean_finding_db() - with self.assertRaises(ValidationError): _ = Event.create_event( title=None, # title required diff --git a/monkey/monkey_island/cc/models/zero_trust/test_finding.py b/monkey/monkey_island/cc/models/zero_trust/test_finding.py index bc51d83d8..dc3a274f2 100644 --- a/monkey/monkey_island/cc/models/zero_trust/test_finding.py +++ b/monkey/monkey_island/cc/models/zero_trust/test_finding.py @@ -16,11 +16,7 @@ class TestFinding(IslandTestCase): server.json file is found and loaded. """ - @pytest.mark.skip(reason="Broken during ScoutSuite refactoring, need to be fixed") def test_save_finding_validation(self): - self.fail_if_not_testing_env() - self.clean_finding_db() - with self.assertRaises(ValidationError): _ = Finding.save_finding(test="bla bla", status=zero_trust_consts.STATUS_FAILED, events=[]) @@ -29,9 +25,6 @@ class TestFinding(IslandTestCase): @pytest.mark.skip(reason="Broken during ScoutSuite refactoring, need to be fixed") def test_save_finding_sanity(self): - self.fail_if_not_testing_env() - self.clean_finding_db() - self.assertEqual(len(Finding.objects(test=zero_trust_consts.TEST_SEGMENTATION)), 0) event_example = Event.create_event( diff --git a/monkey/monkey_island/cc/services/edge/test_displayed_edge.py b/monkey/monkey_island/cc/services/edge/test_displayed_edge.py index d2a4e1f58..cd8519327 100644 --- a/monkey/monkey_island/cc/services/edge/test_displayed_edge.py +++ b/monkey/monkey_island/cc/services/edge/test_displayed_edge.py @@ -47,7 +47,6 @@ EXPLOIT_DATA_MOCK = [{ class TestDisplayedEdgeService(IslandTestCase): def test_get_displayed_edges_by_to(self): - self.clean_edge_db() dst_id = ObjectId() diff --git a/monkey/monkey_island/cc/services/edge/test_edge.py b/monkey/monkey_island/cc/services/edge/test_edge.py index 8ebf45343..bc2c40e80 100644 --- a/monkey/monkey_island/cc/services/edge/test_edge.py +++ b/monkey/monkey_island/cc/services/edge/test_edge.py @@ -19,9 +19,6 @@ class TestEdgeService(IslandTestCase): """ def test_get_or_create_edge(self): - self.fail_if_not_testing_env() - self.clean_edge_db() - src_id = ObjectId() dst_id = ObjectId() diff --git a/monkey/monkey_island/cc/services/reporting/test_pth_report.py b/monkey/monkey_island/cc/services/reporting/test_pth_report.py index b5a628fb1..fc1e6dd27 100644 --- a/monkey/monkey_island/cc/services/reporting/test_pth_report.py +++ b/monkey/monkey_island/cc/services/reporting/test_pth_report.py @@ -7,9 +7,6 @@ from monkey_island.cc.testing.IslandTestCase import IslandTestCase class TestPTHReportServiceGenerateMapNodes(IslandTestCase): def test_generate_map_nodes(self): - self.fail_if_not_testing_env() - self.clean_monkey_db() - self.assertEqual(PTHReportService.generate_map_nodes(), []) windows_monkey_with_services = Monkey( @@ -43,9 +40,6 @@ class TestPTHReportServiceGenerateMapNodes(IslandTestCase): self.assertEqual(2, len(map_nodes)) def test_generate_map_nodes_parsing(self): - self.fail_if_not_testing_env() - self.clean_monkey_db() - monkey_id = str(uuid.uuid4()) hostname = "A_Windows_PC_1" windows_monkey_with_services = Monkey( diff --git a/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/test_environment.py b/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/test_environment.py index f85b2b88c..8752a9e9d 100644 --- a/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/test_environment.py +++ b/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/test_environment.py @@ -8,9 +8,6 @@ from monkey_island.cc.testing.IslandTestCase import IslandTestCase class TestEnvironmentTelemetryProcessing(IslandTestCase): def test_process_environment_telemetry(self): - self.fail_if_not_testing_env() - self.clean_monkey_db() - # Arrange monkey_guid = str(uuid.uuid4()) a_monkey = Monkey(guid=monkey_guid) diff --git a/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/test_system_info_telemetry_dispatcher.py b/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/test_system_info_telemetry_dispatcher.py index 2af2d5970..4c31d466e 100644 --- a/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/test_system_info_telemetry_dispatcher.py +++ b/monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/test_system_info_telemetry_dispatcher.py @@ -36,9 +36,6 @@ class SystemInfoTelemetryDispatcherTest(IslandTestCase): dispatcher.dispatch_collector_results_to_relevant_processors(good_telem_empty_collectors) def test_dispatch_to_relevant_collector(self): - self.fail_if_not_testing_env() - self.clean_monkey_db() - a_monkey = Monkey(guid=str(uuid.uuid4())) a_monkey.save() diff --git a/monkey/monkey_island/cc/services/telemetry/zero_trust_checks/test_segmentation.py b/monkey/monkey_island/cc/services/telemetry/zero_trust_checks/test_segmentation.py index e4aa49dc2..6db918b52 100644 --- a/monkey/monkey_island/cc/services/telemetry/zero_trust_checks/test_segmentation.py +++ b/monkey/monkey_island/cc/services/telemetry/zero_trust_checks/test_segmentation.py @@ -16,9 +16,6 @@ THIRD_SUBNET = "3.3.3.3-3.3.3.200" class TestSegmentationChecks(IslandTestCase): def test_create_findings_for_all_done_pairs(self): - self.fail_if_not_testing_env() - self.clean_finding_db() - all_subnets = [FIRST_SUBNET, SECOND_SUBNET, THIRD_SUBNET] monkey = Monkey( diff --git a/monkey/monkey_island/cc/testing/IslandTestCase.py b/monkey/monkey_island/cc/testing/IslandTestCase.py index b260f62c9..d6231c01e 100644 --- a/monkey/monkey_island/cc/testing/IslandTestCase.py +++ b/monkey/monkey_island/cc/testing/IslandTestCase.py @@ -1,5 +1,7 @@ import unittest +import mongoengine + import monkey_island.cc.environment.environment_singleton as env_singleton from monkey_island.cc.models import Monkey from monkey_island.cc.models.edge import Edge @@ -7,17 +9,33 @@ from monkey_island.cc.models.zero_trust.finding import Finding class IslandTestCase(unittest.TestCase): + + def __init__(self, methodName): + # Make sure test is working with mongomock + if mongoengine.connection.get_connection().server_info()['sysInfo'] != 'Mock': + mongoengine.disconnect() + mongoengine.connect('mongoenginetest', host='mongomock://localhost') + else: + IslandTestCase.clean_db() + super().__init__(methodName) + def fail_if_not_testing_env(self): self.assertFalse(not env_singleton.env.testing, "Change server_config.json to testing environment.") @staticmethod - def clean_monkey_db(): + def clean_db(): + IslandTestCase._clean_edge_db() + IslandTestCase._clean_monkey_db() + IslandTestCase._clean_finding_db() + + @staticmethod + def _clean_monkey_db(): Monkey.objects().delete() @staticmethod - def clean_edge_db(): + def _clean_edge_db(): Edge.objects().delete() @staticmethod - def clean_finding_db(): + def _clean_finding_db(): Finding.objects().delete()