From 736e779f4cb91ab6c08164e39c33a3280faec67e Mon Sep 17 00:00:00 2001
From: Ilija Lazoroski <ilija.la@live.com>
Date: Wed, 10 Aug 2022 17:32:39 +0200
Subject: [PATCH] BB: Gather enabled tests and select GCP machines needed for
 the specific tests

---
 envs/monkey_zoo/blackbox/conftest.py      | 31 +++++++++++++++++++++--
 envs/monkey_zoo/blackbox/test_blackbox.py | 16 +++---------
 2 files changed, 32 insertions(+), 15 deletions(-)

diff --git a/envs/monkey_zoo/blackbox/conftest.py b/envs/monkey_zoo/blackbox/conftest.py
index d844bc403..09dce3c27 100644
--- a/envs/monkey_zoo/blackbox/conftest.py
+++ b/envs/monkey_zoo/blackbox/conftest.py
@@ -1,5 +1,10 @@
 import pytest
 
+from envs.monkey_zoo.blackbox.gcp_test_machine_list import (
+    GCP_SINGLE_TEST_LIST,
+    GCP_TEST_MACHINE_LIST,
+)
+
 
 def pytest_addoption(parser):
     parser.addoption(
@@ -33,8 +38,30 @@ def no_gcp(request):
 
 
 @pytest.fixture(scope="session")
-def machines_to_start(request):
-    return request.config.getoption("-k")
+def list_machines(request):
+    enabled_tests = [test.name for test in request.node.items]
+
+    if len(enabled_tests) == len(GCP_SINGLE_TEST_LIST.keys()):
+        return GCP_TEST_MACHINE_LIST
+
+    try:
+        list_machines_to_start = [GCP_SINGLE_TEST_LIST[test] for test in enabled_tests]
+
+        if len(list_machines_to_start) == 1:
+            return list_machines_to_start[0]
+
+        single_machine_list = {}
+
+        for machine_dict in list_machines_to_start:
+            for zone, machines in machine_dict.items():
+                for machine in machines:
+                    if machine not in single_machine_list[zone]:
+                        single_machine_list[zone].append(machine)
+
+    except KeyError:
+        return GCP_TEST_MACHINE_LIST
+
+    return single_machine_list
 
 
 def pytest_runtest_setup(item):
diff --git a/envs/monkey_zoo/blackbox/test_blackbox.py b/envs/monkey_zoo/blackbox/test_blackbox.py
index 4731a70ae..52e9e10dd 100644
--- a/envs/monkey_zoo/blackbox/test_blackbox.py
+++ b/envs/monkey_zoo/blackbox/test_blackbox.py
@@ -6,10 +6,6 @@ import pytest
 
 from envs.monkey_zoo.blackbox.analyzers.communication_analyzer import CommunicationAnalyzer
 from envs.monkey_zoo.blackbox.analyzers.zerologon_analyzer import ZerologonAnalyzer
-from envs.monkey_zoo.blackbox.gcp_test_machine_list import (
-    GCP_SINGLE_TEST_LIST,
-    GCP_TEST_MACHINE_LIST,
-)
 from envs.monkey_zoo.blackbox.island_client.monkey_island_client import MonkeyIslandClient
 from envs.monkey_zoo.blackbox.island_client.test_configuration_parser import get_target_ips
 from envs.monkey_zoo.blackbox.log_handlers.test_logs_handler import TestLogsHandler
@@ -38,16 +34,10 @@ LOGGER = logging.getLogger(__name__)
 
 
 @pytest.fixture(autouse=True, scope="session")
-def GCPHandler(request, no_gcp, machines_to_start):
+def GCPHandler(request, no_gcp, list_machines):
     if not no_gcp:
-        list_machines = GCP_TEST_MACHINE_LIST
-        if machines_to_start:
-            try:
-                list_machines = GCP_SINGLE_TEST_LIST[machines_to_start]
-            except KeyError as err:
-                LOGGER.warning(
-                    f"Partial or wrong test name provided. " f"Starting all GCP machines:{err}"
-                )
+        LOGGER.info(f"MACHINES TO START: {list_machines}")
+
         try:
             initialize_gcp_client()
             start_machines(list_machines)