Further improving and refactoring performance test code

envs/monkey_zoo/blackbox/analyzers/performance_analyzer.py

@@ -1,37 +1,39 @@
 import logging
 from datetime import timedelta
+from typing import Dict
 
 from envs.monkey_zoo.blackbox.analyzers.analyzer import Analyzer
-from envs.monkey_zoo.blackbox.island_client.monkey_island_client import MonkeyIslandClient
+from envs.monkey_zoo.blackbox.tests.performance.performance_test_config import PerformanceTestConfig
-
-MAX_ALLOWED_SINGLE_PAGE_TIME = timedelta(seconds=2)
-MAX_ALLOWED_TOTAL_TIME = timedelta(seconds=5)
-
-REPORT_URLS = [
-    "api/report/security",
-    "api/attack/report",
-    "api/report/zero_trust/findings",
-    "api/report/zero_trust/principles",
-    "api/report/zero_trust/pillars"
-]
-
 LOGGER = logging.getLogger(__name__)
 
 
 class PerformanceAnalyzer(Analyzer):
 
-    def __init__(self, island_client: MonkeyIslandClient, break_if_took_too_long=False):
+    def __init__(self, performance_test_config: PerformanceTestConfig, endpoint_timings: Dict[str, timedelta]):
-        self.break_if_took_too_long = break_if_took_too_long
+        self.performance_test_config = performance_test_config
-        self.island_client = island_client
+        self.endpoint_timings = endpoint_timings
 
+    def analyze_test_results(self):
+        # Calculate total time and check each page
+        single_page_time_less_then_max = True
+        total_time = timedelta()
+        for page, elapsed in self.endpoint_timings.items():
+            LOGGER.info(f"page {page} took {str(elapsed)}")
+            total_time += elapsed
+            if elapsed > self.performance_test_config.max_allowed_single_page_time:
+                single_page_time_less_then_max = False
 
+        total_time_less_then_max = total_time < self.performance_test_config.max_allowed_total_time
 
-    def get_elapsed_for_get_request(self, url):
+        LOGGER.info(f"total time is {str(total_time)}")
-        response = self.island_client.requests.get(url)
+
-        if response.ok:
+        performance_is_good_enough = total_time_less_then_max and single_page_time_less_then_max
-            LOGGER.debug(f"Got ok for {url} content peek:\n{response.content[:120].strip()}")
+
-            return response.elapsed
+        if self.performance_test_config.break_on_timeout and not performance_is_good_enough:
-        else:
+            LOGGER.warning(
-            LOGGER.error(f"Trying to get {url} but got unexpected {str(response)}")
+                "Calling breakpoint - pausing to enable investigation of island. Type 'c' to continue once you're done "
-            # instead of raising for status, mark failed responses as maxtime
+                "investigating. Type 'p timings' and 'p total_time' to see performance information."
-            return timedelta.max()
+            )
+            breakpoint()
+
+        return performance_is_good_enough

envs/monkey_zoo/blackbox/island_client/monkey_island_client.py

@@ -96,13 +96,3 @@ class MonkeyIslandClient(object):
         response = self.requests.get("api/test/clear_caches")
         response.raise_for_status()
         return response
-
-    def get_elapsed_for_get_request(self, url):
-        response = self.requests.get(url)
-        if response.ok:
-            LOGGER.debug(f"Got ok for {url} content peek:\n{response.content[:120].strip()}")
-            return response.elapsed
-        else:
-            LOGGER.error(f"Trying to get {url} but got unexpected {str(response)}")
-            # instead of raising for status, mark failed responses as maxtime
-            return timedelta.max()

envs/monkey_zoo/blackbox/tests/performance/map_generation.py

@@ -0,0 +1,47 @@
+from datetime import timedelta
+
+from envs.monkey_zoo.blackbox.tests.basic_test import BasicTest
+from envs.monkey_zoo.blackbox.tests.exploitation import ExploitationTest
+from envs.monkey_zoo.blackbox.tests.performance.performance_test_config import PerformanceTestConfig
+from envs.monkey_zoo.blackbox.tests.performance.performance_test import PerformanceTest
+
+MAX_ALLOWED_SINGLE_PAGE_TIME = timedelta(seconds=2)
+MAX_ALLOWED_TOTAL_TIME = timedelta(seconds=5)
+
+REPORT_RESOURCES = [
+    "api/report/security",
+    "api/attack/report",
+    "api/report/zero_trust/findings",
+    "api/report/zero_trust/principles",
+    "api/report/zero_trust/pillars"
+]
+
+
+class MapGenerationTest(BasicTest):
+
+    def __init__(self, name, island_client, config_parser, analyzers,
+                 timeout, log_handler, break_on_timeout=False):
+        self.name = name
+        self.island_client = island_client
+        self.config_parser = config_parser
+        self.exploitation_test = ExploitationTest(name, island_client, config_parser, analyzers, timeout, log_handler)
+        self.break_on_timeout = break_on_timeout
+
+    def run(self):
+        self.island_client.import_config(self.config_parser.config_raw)
+        self.exploitation_test.print_test_starting_info()
+        try:
+            self.island_client.run_monkey_local()
+            self.exploitation_test.test_until_timeout()
+        finally:
+            self.island_client.kill_all_monkeys()
+            self.exploitation_test.wait_until_monkeys_die()
+            self.exploitation_test.wait_for_monkey_process_to_finish()
+            performance_config = PerformanceTestConfig(max_allowed_single_page_time=MAX_ALLOWED_SINGLE_PAGE_TIME,
+                                                       max_allowed_total_time=MAX_ALLOWED_TOTAL_TIME,
+                                                       endpoints_to_test=REPORT_RESOURCES,
+                                                       break_on_timeout=self.break_on_timeout)
+            performance_test = PerformanceTest("Report generation test", performance_config, self.exploitation_test)
+            performance_test.run()
+            self.exploitation_test.parse_logs()
+            self.island_client.reset_env()

envs/monkey_zoo/blackbox/tests/performance/performance_test.py

@@ -0,0 +1,42 @@
+import logging
+from datetime import timedelta
+
+from envs.monkey_zoo.blackbox.tests.basic_test import BasicTest
+from envs.monkey_zoo.blackbox.tests.exploitation import ExploitationTest
+from envs.monkey_zoo.blackbox.tests.performance.performance_test_config import PerformanceTestConfig
+from envs.monkey_zoo.blackbox.analyzers.performance_analyzer import PerformanceAnalyzer
+
+
+LOGGER = logging.getLogger(__name__)
+
+
+class PerformanceTest(BasicTest):
+
+    def __init__(self, name, test_config: PerformanceTestConfig, exploitation_test: ExploitationTest):
+        self.name = name
+        self.test_config = test_config
+        self.exploitation_test = exploitation_test
+
+    def run(self) -> bool:
+        if not self.exploitation_test.island_client.is_all_monkeys_dead():
+            raise RuntimeError("Can't test report times since not all Monkeys have died.")
+
+        # Collect timings for all pages
+        self.exploitation_test.island_client.clear_caches()
+        endpoint_timings = {}
+        for endpoint in self.test_config.endpoints_to_test:
+            endpoint_timings[endpoint] = self.exploitation_test.island_client.get_elapsed_for_get_request(endpoint)
+
+        analyzer = PerformanceAnalyzer(self.test_config, endpoint_timings)
+
+        return analyzer.analyze_test_results()
+
+    def get_elapsed_for_get_request(self, url):
+        response = self.exploitation_test.island_client.requests.get(url)
+        if response.ok:
+            LOGGER.debug(f"Got ok for {url} content peek:\n{response.content[:120].strip()}")
+            return response.elapsed
+        else:
+            LOGGER.error(f"Trying to get {url} but got unexpected {str(response)}")
+            # instead of raising for status, mark failed responses as maxtime
+            return timedelta.max()

envs/monkey_zoo/blackbox/tests/performance/performance_test_config.py

@@ -0,0 +1,12 @@
+from datetime import timedelta
+from typing import List
+
+
+class PerformanceTestConfig:
+
+    def __init__(self, max_allowed_single_page_time: timedelta, max_allowed_total_time: timedelta,
+                 endpoints_to_test: List[str], break_on_timeout=False):
+        self.max_allowed_single_page_time = max_allowed_single_page_time
+        self.max_allowed_total_time = max_allowed_total_time
+        self.endpoints_to_test = endpoints_to_test
+        self.break_on_timeout = break_on_timeout

envs/monkey_zoo/blackbox/tests/performance/report_generation.py

@@ -1,66 +1,27 @@
-import logging
 from datetime import timedelta
 
 from envs.monkey_zoo.blackbox.tests.basic_test import BasicTest
 from envs.monkey_zoo.blackbox.tests.exploitation import ExploitationTest
+from envs.monkey_zoo.blackbox.tests.performance.performance_test_config import PerformanceTestConfig
+from envs.monkey_zoo.blackbox.tests.performance.performance_test import PerformanceTest
 
 MAX_ALLOWED_SINGLE_PAGE_TIME = timedelta(seconds=2)
 MAX_ALLOWED_TOTAL_TIME = timedelta(seconds=5)
 
-REPORT_URLS = [
+MAP_RESOURCES = [
-    "api/report/security",
+    "api/report/security"
-    "api/attack/report",
-    "api/report/zero_trust/findings",
-    "api/report/zero_trust/principles",
-    "api/report/zero_trust/pillars"
 ]
 
-LOGGER = logging.getLogger(__name__)
-
 
 class ReportGenerationTest(BasicTest):
 
     def __init__(self, name, island_client, config_parser, analyzers,
-                 timeout, log_handler, break_if_took_too_long=False):
+                 timeout, log_handler, break_on_timeout=False):
         self.name = name
         self.island_client = island_client
         self.config_parser = config_parser
         self.exploitation_test = ExploitationTest(name, island_client, config_parser, analyzers, timeout, log_handler)
-        self.break_if_took_too_long = break_if_took_too_long
+        self.break_on_timeout = break_on_timeout
-
-    def test_report_generation_performance(self) -> bool:
-        if not self.island_client.is_all_monkeys_dead():
-            raise RuntimeError("Can't test report times since not all Monkeys have died.")
-
-        # Collect timings for all pages
-        self.island_client.clear_caches()
-        report_resource_to_response_time = {}
-        for url in REPORT_URLS:
-            report_resource_to_response_time[url] = self.island_client.get_elapsed_for_get_request(url)
-
-        # Calculate total time and check each page
-        single_page_time_less_then_max = True
-        total_time = timedelta()
-        for page, elapsed in report_resource_to_response_time.items():
-            LOGGER.info(f"page {page} took {str(elapsed)}")
-            total_time += elapsed
-            if elapsed > MAX_ALLOWED_SINGLE_PAGE_TIME:
-                single_page_time_less_then_max = False
-
-        total_time_less_then_max = total_time < MAX_ALLOWED_TOTAL_TIME
-
-        LOGGER.info(f"total time is {str(total_time)}")
-
-        performance_is_good_enough = total_time_less_then_max and single_page_time_less_then_max
-
-        if self.break_if_took_too_long and not performance_is_good_enough:
-            LOGGER.warning(
-                "Calling breakpoint - pausing to enable investigation of island. Type 'c' to continue once you're done "
-                "investigating. Type 'p timings' and 'p total_time' to see performance information."
-            )
-            breakpoint()
-
-        return performance_is_good_enough
 
     def run(self):
         self.island_client.import_config(self.config_parser.config_raw)
@@ -72,10 +33,11 @@ class ReportGenerationTest(BasicTest):
             self.island_client.kill_all_monkeys()
             self.exploitation_test.wait_until_monkeys_die()
             self.exploitation_test.wait_for_monkey_process_to_finish()
-            self.test_post_exec_analyzers()
+            performance_config = PerformanceTestConfig(max_allowed_single_page_time=MAX_ALLOWED_SINGLE_PAGE_TIME,
+                                                       max_allowed_total_time=MAX_ALLOWED_TOTAL_TIME,
+                                                       endpoints_to_test=REPORT_RESOURCES,
+                                                       break_on_timeout=self.break_on_timeout)
+            performance_test = PerformanceTest("Report generation test", performance_config, self.exploitation_test)
+            performance_test.run()
             self.exploitation_test.parse_logs()
             self.island_client.reset_env()
-
-    def test_post_exec_analyzers(self):
-        post_exec_analyzers_results = [analyzer.analyze_test_results() for analyzer in self.post_exec_analyzers]
-        assert all(post_exec_analyzers_results)