From 8cb5aa8ca991f064f0587773e216a2a36d4594be Mon Sep 17 00:00:00 2001
From: Adeel Ahmed <37113831+adeelsohailahmed@users.noreply.github.com>
Date: Tue, 1 Oct 2024 20:30:41 +0500
Subject: [PATCH] test: add test to ensure dict with enum keys are encoded
 properly (#1001)

---
 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 c6c30adb..12320ae2 100644
--- a/tests/odm/conftest.py
+++ b/tests/odm/conftest.py
@@ -45,6 +45,7 @@
     DocumentWithCustomInit,
     DocumentWithDecimalField,
     DocumentWithDeprecatedHiddenField,
+    DocumentWithEnumKeysDict,
     DocumentWithExtras,
     DocumentWithHttpUrlField,
     DocumentWithIndexedObjectId,
@@ -292,6 +293,7 @@ async def init(db):
         DocumentToTestSync,
         DocumentWithLinkForNesting,
         DocumentWithBackLinkForNesting,
+        DocumentWithEnumKeysDict,
         LongSelfLink,
         BsonRegexDoc,
         NativeRegexDoc,
diff --git a/tests/odm/models.py b/tests/odm/models.py
index 34699d6e..1de5851c 100644
--- a/tests/odm/models.py
+++ b/tests/odm/models.py
@@ -1143,6 +1143,15 @@ class Settings:
         max_nesting_depth = 50
 
 
+class DictEnum(str, Enum):
+    RED = "Red"
+    BLUE = "Blue"
+
+
+class DocumentWithEnumKeysDict(Document):
+    color: Dict[DictEnum, str]
+
+
 class BsonRegexDoc(Document):
     regex: Optional[Regex] = None
 
diff --git a/tests/odm/test_encoder.py b/tests/odm/test_encoder.py
index 46e8d33e..f839a969 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
@@ -11,10 +12,12 @@
 from tests.odm.models import (
     BsonRegexDoc,
     Child,
+    DictEnum,
     DocumentForEncodingTest,
     DocumentForEncodingTestDate,
     DocumentWithComplexDictKey,
     DocumentWithDecimalField,
+    DocumentWithEnumKeysDict,
     DocumentWithHttpUrlField,
     DocumentWithKeepNullsFalse,
     DocumentWithStringField,
@@ -173,6 +176,17 @@ async def test_dict_with_complex_key():
     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
+
+
 async def test_native_regex():
     regex = re.compile(r"^1?$|^(11+?)\1+$", (re.I | re.M | re.S) ^ re.UNICODE)
     doc = await NativeRegexDoc(regex=regex).insert()