From 2bbc1b093d800641f846a02194fd3adb7d3bb31d Mon Sep 17 00:00:00 2001
From: bohan-amplitude <96497806+bohan-amplitude@users.noreply.github.com>
Date: Wed, 22 Jun 2022 17:25:16 -0700
Subject: [PATCH] fix: get_event_body for enum value from ampli (#29)

---
 src/amplitude/event.py |  8 +++++++-
 src/test/test_event.py | 10 ++++++++--
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/src/amplitude/event.py b/src/amplitude/event.py
index 4b8f8a2..96f44ff 100644
--- a/src/amplitude/event.py
+++ b/src/amplitude/event.py
@@ -12,6 +12,7 @@
 """
 
 import copy
+import enum
 import json
 import logging
 from typing import Callable, Optional, Union
@@ -303,6 +304,11 @@ def get_event_body(self) -> dict:
                 event_body[value[0]] = self[key]
         if "plan" in event_body:
             event_body["plan"] = event_body["plan"].get_plan_body()
+        for properties in ["user_properties", "event_properties", "group_properties"]:
+            if properties in event_body:
+                for key, value in event_body[properties].items():
+                    if isinstance(value, enum.Enum):
+                        event_body[properties][key] = value.value
         return utils.truncate(event_body)
 
     def _verify_property(self, key, value) -> bool:
@@ -1187,7 +1193,7 @@ def is_validate_properties(key, value):
         return result
     if isinstance(value, dict):
         return is_validate_object(value)
-    if not isinstance(value, (bool, float, int, str)):
+    if not isinstance(value, (bool, float, int, str, enum.Enum)):
         return False
     return True
 
diff --git a/src/test/test_event.py b/src/test/test_event.py
index 6ce989b..a54a835 100644
--- a/src/test/test_event.py
+++ b/src/test/test_event.py
@@ -1,3 +1,4 @@
+import enum
 import unittest
 from unittest.mock import MagicMock
 
@@ -74,10 +75,15 @@ def test_callback_without_callback_function_success_callback(self):
         callback_func.assert_not_called()
 
     def test_base_event_get_event_body_success(self):
-        event = BaseEvent(event_type="test_event", user_id="test_user", user_properties={"email": "test@test"})
+        class TestEnum(enum.Enum):
+            ENUM1 = 'test'
+            ENUM2 = 'test2'
+        event = BaseEvent(event_type="test_event", user_id="test_user", user_properties={"email": "test@test"},
+                          event_properties={'enum_properties': TestEnum.ENUM1})
         expect_dict = {"event_type": "test_event",
                        "user_id": "test_user",
-                       "user_properties": {"email": "test@test"}}
+                       "user_properties": {"email": "test@test"},
+                       "event_properties": {"enum_properties": 'test'}}
         self.assertEqual(expect_dict, event.get_event_body())
 
     def test_base_event_set_dict_event_attributes_with_invalid_value_failed(self):