Fixed UTs - enum change wasn't 100% correct

This commit is contained in:
Shay Nehmad 2020-01-21 16:53:00 +02:00
parent 0f73567ad4
commit 5548cf2acf
2 changed files with 5 additions and 4 deletions

View File

@ -4,16 +4,16 @@ from common.data.system_info_collectors_names import ENVIRONMENT_COLLECTOR
from infection_monkey.system_info.system_info_collector import SystemInfoCollector from infection_monkey.system_info.system_info_collector import SystemInfoCollector
def get_monkey_environment() -> Environment: def get_monkey_environment() -> str:
""" """
Get the Monkey's running environment. Get the Monkey's running environment.
:return: One of the cloud providers if on cloud; otherwise, assumes "on premise". :return: One of the cloud providers if on cloud; otherwise, assumes "on premise".
""" """
for instance in get_all_cloud_instances(): for instance in get_all_cloud_instances():
if instance.is_instance(): if instance.is_instance():
return instance.get_cloud_provider_name() return instance.get_cloud_provider_name().value
return Environment.ON_PREMISE return Environment.ON_PREMISE.value
class EnvironmentCollector(SystemInfoCollector): class EnvironmentCollector(SystemInfoCollector):

View File

@ -45,7 +45,8 @@ class Monkey(Document):
command_control_channel = EmbeddedDocumentField(CommandControlChannel) command_control_channel = EmbeddedDocumentField(CommandControlChannel)
# Environment related fields # Environment related fields
environment = StringField(default=environment_names.Environment.UNKNOWN, choices=environment_names.ALL_ENVIRONMENTS_NAMES) environment = StringField(default=environment_names.Environment.UNKNOWN.value,
choices=environment_names.ALL_ENVIRONMENTS_NAMES)
aws_instance_id = StringField(required=False) # This field only exists when the monkey is running on an AWS aws_instance_id = StringField(required=False) # This field only exists when the monkey is running on an AWS
# instance. See https://github.com/guardicore/monkey/issues/426. # instance. See https://github.com/guardicore/monkey/issues/426.