From 9afd043350e2d2a643271ff0c8faa88a7566d5b5 Mon Sep 17 00:00:00 2001
From: tarsil <tiago.arasilva@gmail.com>
Date: Tue, 16 Jan 2024 11:17:22 +0000
Subject: [PATCH] Add unique index validation

---
 mongoz/core/db/documents/metaclasses.py |  4 ++-
 tests/indexes/test_unique.py            | 41 +++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)
 create mode 100644 tests/indexes/test_unique.py

diff --git a/mongoz/core/db/documents/metaclasses.py b/mongoz/core/db/documents/metaclasses.py
index d5743e7..95398d6 100644
--- a/mongoz/core/db/documents/metaclasses.py
+++ b/mongoz/core/db/documents/metaclasses.py
@@ -340,10 +340,12 @@ def __search_for_fields(base: Type, attrs: Any) -> None:
             if not new_class.is_proxy_document:
                 # For the indexes
                 _index: Union[Index, None] = None
-                if hasattr(field, "index") and field.index and field.unique is True:
+                if hasattr(field, "index") and field.index and field.unique:
                     _index = Index(name, unique=True, sparse=field.sparse)
                 elif hasattr(field, "index") and field.index:
                     _index = Index(name, sparse=field.sparse)
+                elif hasattr(field, "unique") and field.unique:
+                    _index = Index(name, unique=True)
 
                 if _index is not None:
                     index_names = [index.name for index in meta.indexes or []]
diff --git a/tests/indexes/test_unique.py b/tests/indexes/test_unique.py
new file mode 100644
index 0000000..61fb08b
--- /dev/null
+++ b/tests/indexes/test_unique.py
@@ -0,0 +1,41 @@
+from typing import AsyncGenerator
+
+import pydantic
+import pytest
+from pymongo.errors import DuplicateKeyError
+from tests.conftest import client
+
+import mongoz
+from mongoz import Document
+
+pytestmark = pytest.mark.anyio
+pydantic_version = pydantic.__version__[:3]
+
+
+class Movie(Document):
+    name: str = mongoz.String(unique=True)
+    is_published: bool = mongoz.Boolean(default=False)
+
+    class Meta:
+        registry = client
+        database = "test_db"
+
+
+@pytest.fixture(scope="function", autouse=True)
+async def prepare_database() -> AsyncGenerator:
+    await Movie.drop_indexes(force=True)
+    await Movie.objects.delete()
+    await Movie.create_indexes()
+    yield
+    await Movie.drop_indexes(force=True)
+    await Movie.objects.delete()
+    await Movie.create_indexes()
+
+
+async def test_unique():
+    movie = await Movie.objects.create(name="Barbie")
+
+    assert movie.is_published is False
+
+    with pytest.raises(DuplicateKeyError):
+        await Movie.objects.create(name="Barbie")