Extracted function in EnvCollector for reuse in other parts of the Monkey

monkey/infection_monkey/system_info/collectors/environment_collector.py

@@ -4,18 +4,21 @@ from common.cloud.environment_names import ON_PREMISE, AZURE, AWS
 from infection_monkey.system_info.system_info_collector import SystemInfoCollector
 
 
+def get_monkey_environment():
+    # Check if on any cloud env. Default is on prem.
+    if AwsInstance().is_aws_instance():
+        env = AWS
+    elif AzureInstance().is_azure_instance():
+        env = AZURE
+    # TODO: elif GcpInstance().is_gcp_instance():
+    else:
+        env = ON_PREMISE
+    return env
+
+
 class EnvironmentCollector(SystemInfoCollector):
     def __init__(self):
         super(EnvironmentCollector, self).__init__(name="EnvironmentCollector")
 
     def collect(self) -> dict:
-        # Check if on any cloud env. Default is on prem.
-        if AwsInstance().is_aws_instance():
-            env = AWS
-        elif AzureInstance().is_azure_instance():
-            env = AZURE
-        # TODO: elif GcpInstance().is_gcp_instance():
-        else:
-            env = ON_PREMISE
-
-        return {"environment": env}
+        return {"environment": get_monkey_environment()}