From 5548cf2acfee76dcd120b9494d9c0d049b4a200a Mon Sep 17 00:00:00 2001
From: Shay Nehmad <shay.nehmad@guardicore.com>
Date: Tue, 21 Jan 2020 16:53:00 +0200
Subject: [PATCH] Fixed UTs - enum change wasn't 100% correct

---
 .../system_info/collectors/environment_collector.py         | 6 +++---
 monkey/monkey_island/cc/models/monkey.py                    | 3 ++-
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/monkey/infection_monkey/system_info/collectors/environment_collector.py b/monkey/infection_monkey/system_info/collectors/environment_collector.py
index 7a953fce9..100d23175 100644
--- a/monkey/infection_monkey/system_info/collectors/environment_collector.py
+++ b/monkey/infection_monkey/system_info/collectors/environment_collector.py
@@ -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):
diff --git a/monkey/monkey_island/cc/models/monkey.py b/monkey/monkey_island/cc/models/monkey.py
index da6d880b4..f658a3d06 100644
--- a/monkey/monkey_island/cc/models/monkey.py
+++ b/monkey/monkey_island/cc/models/monkey.py
@@ -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.