From d125ba669d999411b46f3c9355223da23f0db9a4 Mon Sep 17 00:00:00 2001
From: adeelsohailahmed <37113831+adeelsohailahmed@users.noreply.github.com>
Date: Thu, 15 Aug 2024 21:48:27 +0500
Subject: [PATCH] test: add test to ensure dict with enum keys are encoded
 properly

---
 tests/odm/conftest.py     |  2 ++
 tests/odm/models.py       |  9 +++++++++
 tests/odm/test_encoder.py | 14 ++++++++++++++
 3 files changed, 25 insertions(+)

diff --git a/tests/odm/conftest.py b/tests/odm/conftest.py
index afb582e7f..bc1e8c115 100644
--- a/tests/odm/conftest.py
+++ b/tests/odm/conftest.py
@@ -44,6 +44,7 @@
     DocumentWithCustomInit,
     DocumentWithDecimalField,
     DocumentWithDeprecatedHiddenField,
+    DocumentWithEnumKeysDict,
     DocumentWithExtras,
     DocumentWithHttpUrlField,
     DocumentWithIndexedObjectId,
@@ -290,6 +291,7 @@ async def init(db):
         DocumentToTestSync,
         DocumentWithLinkForNesting,
         DocumentWithBackLinkForNesting,
+        DocumentWithEnumKeysDict,
         LongSelfLink,
     ]
     await init_beanie(
diff --git a/tests/odm/models.py b/tests/odm/models.py
index 760f8863b..3e6f2c112 100644
--- a/tests/odm/models.py
+++ b/tests/odm/models.py
@@ -1134,3 +1134,12 @@ class LongSelfLink(Document):
 
     class Settings:
         max_nesting_depth = 50
+
+
+class DictEnum(str, Enum):
+    RED = "Red"
+    BLUE = "Blue"
+
+
+class DocumentWithEnumKeysDict(Document):
+    color: Dict[DictEnum, str]
diff --git a/tests/odm/test_encoder.py b/tests/odm/test_encoder.py
index 9e7421846..2b94c3f76 100644
--- a/tests/odm/test_encoder.py
+++ b/tests/odm/test_encoder.py
@@ -1,5 +1,6 @@
 import re
 from datetime import date, datetime
+from enum import Enum
 from uuid import uuid4
 
 import pytest
@@ -10,10 +11,12 @@
 from beanie.odm.utils.pydantic import IS_PYDANTIC_V2
 from tests.odm.models import (
     Child,
+    DictEnum,
     DocumentForEncodingTest,
     DocumentForEncodingTestDate,
     DocumentWithComplexDictKey,
     DocumentWithDecimalField,
+    DocumentWithEnumKeysDict,
     DocumentWithHttpUrlField,
     DocumentWithKeepNullsFalse,
     DocumentWithStringField,
@@ -169,3 +172,14 @@ async def test_dict_with_complex_key():
 
     assert isinstance(new_doc.dict_field, dict)
     assert new_doc.dict_field.get(uuid) == dt
+
+
+async def test_dict_with_enum_keys():
+    doc = DocumentWithEnumKeysDict(color={DictEnum.RED: "favorite"})
+    await doc.save()
+
+    assert isinstance(doc.color, dict)
+
+    for key in doc.color:
+        assert isinstance(key, Enum)
+        assert key == DictEnum.RED