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
def get_monkey_environment() -> Environment:
def get_monkey_environment() -> str:
"""
Get the Monkey's running environment.
:return: One of the cloud providers if on cloud; otherwise, assumes "on premise".
"""
for instance in get_all_cloud_instances():
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):

View File

@ -45,7 +45,8 @@ class Monkey(Document):
command_control_channel = EmbeddedDocumentField(CommandControlChannel)
# 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
# instance. See https://github.com/guardicore/monkey/issues/426.