From fe77c7c04d9d7f6b7b771ef6344f7e405e2c84ec Mon Sep 17 00:00:00 2001
From: Jeferson Daniel <jeferson.daniel412@gmail.com>
Date: Sat, 6 Jul 2024 17:14:50 -0300
Subject: [PATCH] fix: Support enums

---
 pydantic_mongo/abstract_repository.py |  2 +-
 test/test_repository.py               | 40 +++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/pydantic_mongo/abstract_repository.py b/pydantic_mongo/abstract_repository.py
index 112f961..02322ba 100644
--- a/pydantic_mongo/abstract_repository.py
+++ b/pydantic_mongo/abstract_repository.py
@@ -70,7 +70,7 @@ def to_document(model: T) -> dict:
         :return: dict
         """
         model_with_id = cast(ModelWithId, model)
-        data = model_with_id.model_dump()
+        data = model_with_id.model_dump(mode="json")
         data.pop("id")
         if model_with_id.id:
             data["_id"] = model_with_id.id
diff --git a/test/test_repository.py b/test/test_repository.py
index 81d51bc..3021a69 100644
--- a/test/test_repository.py
+++ b/test/test_repository.py
@@ -3,6 +3,8 @@
 import mongomock
 import pytest
 from bson import ObjectId
+from datetime import datetime, timezone
+from enum import Enum
 from pydantic import BaseModel, Field
 
 from pydantic_mongo import AbstractRepository, PydanticObjectId
@@ -30,6 +32,22 @@ class Meta:
         collection_name = "spams"
 
 
+class State(Enum):
+    Preparation = "Preparation"
+    Processing = "Processing"
+
+
+class Order(BaseModel):
+    id: Optional[PydanticObjectId] = None
+    created_at: datetime
+    state: Optional[State]
+
+
+class OrderRepository(AbstractRepository[Order]):
+    class Meta:
+        collection_name = "orders"
+
+
 @pytest.fixture
 def database():
     return mongomock.MongoClient().db
@@ -259,3 +277,25 @@ def test_paginate(self, database):
 
         with pytest.raises(PaginationError):
             spam_repository.paginate({}, limit=10, after="invalid string")
+
+    def test_store_enums(self, database):
+        order_repository = OrderRepository(database=database)
+        current_date = datetime.now(tz=timezone.utc)
+
+        order = Order(
+            state=State.Preparation,
+            created_at=current_date
+        )
+        order_repository.save(order)
+
+        mongo_document = database["orders"].find_one({
+            "_id": order.id,
+        })
+
+        assert mongo_document["state"] == "Preparation"
+        assert isinstance(mongo_document["created_at"], datetime)
+
+        find_result = order_repository.find_one_by_id(order.id)
+        assert find_result is not None
+        assert find_result.state == State.Preparation
+        assert find_result.created_at == current_date