test: add test to ensure dict with enum keys are encoded properly (#1001

)

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,

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
 

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()