From 96f80565637f7cd0fe18ccb5d9a661c0c03d060b Mon Sep 17 00:00:00 2001 From: Harrison Zhao Date: Sun, 31 May 2026 00:28:04 -0700 Subject: [PATCH 1/2] add support for get_by_ids and count --- README.md | 20 +++ firebase.json | 2 +- firedantic/_async/model.py | 134 ++++++++++++++++++--- firedantic/_sync/model.py | 134 ++++++++++++++++++--- firedantic/tests/tests_async/test_model.py | 46 +++++++ firedantic/tests/tests_sync/test_model.py | 42 +++++++ setup.cfg | 2 +- 7 files changed, 340 insertions(+), 40 deletions(-) diff --git a/README.md b/README.md index e80d025..fd4b975 100644 --- a/README.md +++ b/README.md @@ -92,11 +92,21 @@ Product.find({"stock": {">=": 3}}) Product.find({"stock": {op.GTE: 3}}) Product.find({"stock": {">=": 1}}, order_by=[('unit_value', Query.ASCENDING)], limit=25, offset=50) Product.find(order_by=[('unit_value', Query.ASCENDING), ('stock', Query.DESCENDING)], limit=2) + +# Count documents, with or without filters +Product.count() +Product.count({"stock": {op.GTE: 3}}) + +# Fetch multiple documents by Firestore document ID in one batch request +Product.get_by_ids(["firestore-doc-id-1", "firestore-doc-id-2"]) ``` The query operators are found at [https://firebase.google.com/docs/firestore/query-data/queries#query_operators](https://firebase.google.com/docs/firestore/query-data/queries#query_operators). +If your model customizes the document ID field by extending `BareModel`, use +`get_by_doc_ids()` instead of `get_by_ids()`. + ### Async usage Firedantic can also be used in an async way, like this: @@ -144,6 +154,12 @@ async def main(): assert bob.id == found_bob.id print(f"Found Bob: {found_bob.id}") + people = await Person.get_by_ids([alice.id, bob.id]) + print(f"Found {len(people)} people") + + total_people = await Person.count() + print(f"Total people: {total_people}") + await alice.delete() print("Deleted Alice") await bob.delete() @@ -327,13 +343,17 @@ Firedantic has basic support for The following methods can be used in a transaction: - `Model.delete(transaction=transaction)` +- `Model.count(transaction=transaction)` - `Model.find_one(transaction=transaction)` - `Model.find(transaction=transaction)` - `Model.get_by_doc_id(transaction=transaction)` +- `Model.get_by_doc_ids(transaction=transaction)` - `Model.get_by_id(transaction=transaction)` +- `Model.get_by_ids(transaction=transaction)` - `Model.reload(transaction=transaction)` - `Model.save(transaction=transaction)` - `SubModel.get_by_id(transaction=transaction)` +- `SubModel.get_by_ids(transaction=transaction)` When using transactions, note that read operations must come before write operations. diff --git a/firebase.json b/firebase.json index 6e53d0a..95556b1 100644 --- a/firebase.json +++ b/firebase.json @@ -5,7 +5,7 @@ }, "ui": { "enabled": true, - "port": 4000 + "port": 4343 } } } diff --git a/firedantic/_async/model.py b/firedantic/_async/model.py index 84251df..afa2da6 100644 --- a/firedantic/_async/model.py +++ b/firedantic/_async/model.py @@ -181,28 +181,59 @@ async def find( # pylint: disable=too-many-arguments if offset is not None: query = query.offset(offset) # type: ignore - def _cls(doc_id: str, data: Dict[str, Any]) -> TAsyncBareModel: - if cls.__document_id__ in data: - logger.warning( - "%s document ID %s contains conflicting %s in data with value %s", - cls.__name__, - doc_id, - cls.__document_id__, - data[cls.__document_id__], - ) - data[cls.__document_id__] = doc_id - model = cls(**data) - setattr(model, cls.__document_id__, doc_id) - return model - return [ - _cls(doc_id, doc_dict) + cls._create_from_doc_data(doc_id, doc_dict) async for doc_id, doc_dict in ( (doc.id, doc.to_dict()) async for doc in query.stream(transaction=transaction) # type: ignore ) if doc_dict is not None ] + @classmethod + async def count( + cls: Type[TAsyncBareModel], + filter_: Optional[Dict[str, Union[str, dict]]] = None, + transaction: Optional[AsyncTransaction] = None, + ) -> int: + """ + Returns the number of models matching the given filter. + + Example: `Company.count()`. + Example: `Company.count({"owner.first_name": "John"})`. + + :param filter_: The filter criteria. + :param transaction: Optional transaction to use. + :return: Number of matching models. + """ + query: Union[AsyncQuery, AsyncCollectionReference] = cls._get_col_ref() + if filter_: + for key, value in filter_.items(): + query = cls._add_filter(query, key, value) + + aggregation_results = await query.count().get(transaction=transaction) # type: ignore + for aggregation_result in aggregation_results: + if len(aggregation_result) == 0: + continue + return int(aggregation_result[0].value) + return 0 + + @classmethod + def _create_from_doc_data( + cls: Type[TAsyncBareModel], doc_id: str, data: Dict[str, Any] + ) -> TAsyncBareModel: + if cls.__document_id__ in data: + logger.warning( + "%s document ID %s contains conflicting %s in data with value %s", + cls.__name__, + doc_id, + cls.__document_id__, + data[cls.__document_id__], + ) + data[cls.__document_id__] = doc_id + model = cls(**data) + setattr(model, cls.__document_id__, doc_id) + return model + @classmethod def _add_filter( cls, query: Union[AsyncQuery, AsyncCollectionReference], field: str, value: Any @@ -279,10 +310,47 @@ async def get_by_doc_id( raise ModelNotFoundError( f"No '{cls.__name__}' found with {cls.__document_id__} '{doc_id}'" ) - data[cls.__document_id__] = doc_id - model = cls(**data) - setattr(model, cls.__document_id__, doc_id) - return model + return cls._create_from_doc_data(doc_id, data) + + @classmethod + async def get_by_doc_ids( + cls: Type[TAsyncBareModel], + doc_ids: Iterable[str], + transaction: Optional[AsyncTransaction] = None, + ) -> List[TAsyncBareModel]: + """ + Returns models matching the given document IDs. + + Missing documents are omitted from the returned list. + + :param doc_ids: Document IDs to fetch. + :param transaction: Optional transaction to use. + :return: List of found models. + :raise ModelNotFoundError: Raised if a document ID is invalid. + """ + doc_refs = [] + for doc_id in doc_ids: + try: + cls._validate_document_id(doc_id) + except InvalidDocumentID as e: + raise ModelNotFoundError( + f"No '{cls.__name__}' found with {cls.__document_id__} '{doc_id}'" + ) from e + doc_refs.append(cls._get_col_ref().document(doc_id)) + if len(doc_refs) == 0: + return [] + + return [ + cls._create_from_doc_data(doc_id, doc_dict) + async for doc_id, doc_dict in ( + (doc.id, doc.to_dict()) + async for doc in CONFIGURATIONS["db"].get_all( + doc_refs, + transaction=transaction, + ) + ) + if doc_dict is not None + ] @classmethod async def truncate_collection(cls, batch_size: int = 128) -> int: @@ -372,6 +440,20 @@ async def get_by_id( """ return await cls.get_by_doc_id(id_, transaction=transaction) + @classmethod + async def get_by_ids( + cls: Type[TAsyncBareModel], + ids: Iterable[str], + transaction: Optional[AsyncTransaction] = None, + ) -> List[TAsyncBareModel]: + """ + Get models by document IDs. + + :param ids: Document IDs. + :param transaction: Optional transaction to use. + """ + return await cls.get_by_doc_ids(ids, transaction=transaction) + class AsyncBareSubCollection(ABC): __collection_tpl__: Optional[str] = None @@ -445,6 +527,20 @@ async def get_by_id( """ return await cls.get_by_doc_id(id_, transaction=transaction) + @classmethod + async def get_by_ids( + cls: Type[TAsyncBareModel], + ids: Iterable[str], + transaction: Optional[AsyncTransaction] = None, + ) -> List[TAsyncBareModel]: + """ + Get items by document IDs. + + :param ids: Document IDs. + :param transaction: Optional transaction to use. + """ + return await cls.get_by_doc_ids(ids, transaction=transaction) + class AsyncSubCollection(AsyncBareSubCollection, ABC): __document_id__ = "id" diff --git a/firedantic/_sync/model.py b/firedantic/_sync/model.py index 78ec372..3a9f98c 100644 --- a/firedantic/_sync/model.py +++ b/firedantic/_sync/model.py @@ -181,28 +181,59 @@ def find( # pylint: disable=too-many-arguments if offset is not None: query = query.offset(offset) # type: ignore - def _cls(doc_id: str, data: Dict[str, Any]) -> TBareModel: - if cls.__document_id__ in data: - logger.warning( - "%s document ID %s contains conflicting %s in data with value %s", - cls.__name__, - doc_id, - cls.__document_id__, - data[cls.__document_id__], - ) - data[cls.__document_id__] = doc_id - model = cls(**data) - setattr(model, cls.__document_id__, doc_id) - return model - return [ - _cls(doc_id, doc_dict) + cls._create_from_doc_data(doc_id, doc_dict) for doc_id, doc_dict in ( (doc.id, doc.to_dict()) for doc in query.stream(transaction=transaction) # type: ignore ) if doc_dict is not None ] + @classmethod + def count( + cls: Type[TBareModel], + filter_: Optional[Dict[str, Union[str, dict]]] = None, + transaction: Optional[Transaction] = None, + ) -> int: + """ + Returns the number of models matching the given filter. + + Example: `Company.count()`. + Example: `Company.count({"owner.first_name": "John"})`. + + :param filter_: The filter criteria. + :param transaction: Optional transaction to use. + :return: Number of matching models. + """ + query: Union[BaseQuery, CollectionReference] = cls._get_col_ref() + if filter_: + for key, value in filter_.items(): + query = cls._add_filter(query, key, value) + + aggregation_results = query.count().get(transaction=transaction) # type: ignore + for aggregation_result in aggregation_results: + if len(aggregation_result) == 0: + continue + return int(aggregation_result[0].value) + return 0 + + @classmethod + def _create_from_doc_data( + cls: Type[TBareModel], doc_id: str, data: Dict[str, Any] + ) -> TBareModel: + if cls.__document_id__ in data: + logger.warning( + "%s document ID %s contains conflicting %s in data with value %s", + cls.__name__, + doc_id, + cls.__document_id__, + data[cls.__document_id__], + ) + data[cls.__document_id__] = doc_id + model = cls(**data) + setattr(model, cls.__document_id__, doc_id) + return model + @classmethod def _add_filter( cls, query: Union[BaseQuery, CollectionReference], field: str, value: Any @@ -277,10 +308,47 @@ def get_by_doc_id( raise ModelNotFoundError( f"No '{cls.__name__}' found with {cls.__document_id__} '{doc_id}'" ) - data[cls.__document_id__] = doc_id - model = cls(**data) - setattr(model, cls.__document_id__, doc_id) - return model + return cls._create_from_doc_data(doc_id, data) + + @classmethod + def get_by_doc_ids( + cls: Type[TBareModel], + doc_ids: Iterable[str], + transaction: Optional[Transaction] = None, + ) -> List[TBareModel]: + """ + Returns models matching the given document IDs. + + Missing documents are omitted from the returned list. + + :param doc_ids: Document IDs to fetch. + :param transaction: Optional transaction to use. + :return: List of found models. + :raise ModelNotFoundError: Raised if a document ID is invalid. + """ + doc_refs = [] + for doc_id in doc_ids: + try: + cls._validate_document_id(doc_id) + except InvalidDocumentID as e: + raise ModelNotFoundError( + f"No '{cls.__name__}' found with {cls.__document_id__} '{doc_id}'" + ) from e + doc_refs.append(cls._get_col_ref().document(doc_id)) + if len(doc_refs) == 0: + return [] + + return [ + cls._create_from_doc_data(doc_id, doc_dict) + for doc_id, doc_dict in ( + (doc.id, doc.to_dict()) + for doc in CONFIGURATIONS["db"].get_all( + doc_refs, + transaction=transaction, + ) + ) + if doc_dict is not None + ] @classmethod def truncate_collection(cls, batch_size: int = 128) -> int: @@ -370,6 +438,20 @@ def get_by_id( """ return cls.get_by_doc_id(id_, transaction=transaction) + @classmethod + def get_by_ids( + cls: Type[TBareModel], + ids: Iterable[str], + transaction: Optional[Transaction] = None, + ) -> List[TBareModel]: + """ + Get models by document IDs. + + :param ids: Document IDs. + :param transaction: Optional transaction to use. + """ + return cls.get_by_doc_ids(ids, transaction=transaction) + class BareSubCollection(ABC): __collection_tpl__: Optional[str] = None @@ -443,6 +525,20 @@ def get_by_id( """ return cls.get_by_doc_id(id_, transaction=transaction) + @classmethod + def get_by_ids( + cls: Type[TBareModel], + ids: Iterable[str], + transaction: Optional[Transaction] = None, + ) -> List[TBareModel]: + """ + Get items by document IDs. + + :param ids: Document IDs. + :param transaction: Optional transaction to use. + """ + return cls.get_by_doc_ids(ids, transaction=transaction) + class SubCollection(BareSubCollection, ABC): __document_id__ = "id" diff --git a/firedantic/tests/tests_async/test_model.py b/firedantic/tests/tests_async/test_model.py index ed27c0d..b29d010 100644 --- a/firedantic/tests/tests_async/test_model.py +++ b/firedantic/tests/tests_async/test_model.py @@ -121,6 +121,39 @@ async def test_find(configure_db, create_company, create_product) -> None: await Product.find({"product_id": {"<>": "a"}}) +@pytest.mark.asyncio +async def test_count(configure_db, create_product) -> None: + for p in TEST_PRODUCTS: + await create_product(**p) + + assert await Product.count() == 4 + assert await Product.count({"stock": {op.GTE: 1}}) == 3 + assert await Product.count({"stock": {op.GTE: 2, op.LT: 4}}) == 2 + assert await Product.count({"product_id": "missing"}) == 0 + + +@pytest.mark.asyncio +async def test_get_by_ids(configure_db, create_company) -> None: + company_a: Company = await create_company(company_id="1234555-1") + company_b: Company = await create_company(company_id="1231231-2") + + assert company_a.id + assert company_b.id + found = await Company.get_by_ids([company_a.id, "missing", company_b.id]) + + assert {company.id for company in found} == {company_a.id, company_b.id} + assert {company.company_id for company in found} == { + company_a.company_id, + company_b.company_id, + } + + +@pytest.mark.asyncio +async def test_get_by_ids_with_invalid_id(configure_db) -> None: + with pytest.raises(ModelNotFoundError): + await Company.get_by_ids([""]) + + @pytest.mark.asyncio async def test_find_not_in(configure_db, create_company) -> None: ids = ["1234555-1", "1234567-8", "2131232-4", "4124432-4"] @@ -385,6 +418,19 @@ async def test_custom_id_model(configure_db) -> None: assert m.bar == "bar" +@pytest.mark.asyncio +async def test_custom_id_model_get_by_doc_ids(configure_db) -> None: + c = CustomIDModel(bar="bar") # type: ignore + await c.save() + assert c.foo + + models = await CustomIDModel.get_by_doc_ids([c.foo, "missing"]) + + assert len(models) == 1 + assert models[0].foo == c.foo + assert models[0].bar == "bar" + + @pytest.mark.asyncio async def test_custom_id_conflict(configure_db) -> None: await CustomIDConflictModel(foo="foo", bar="bar").save() diff --git a/firedantic/tests/tests_sync/test_model.py b/firedantic/tests/tests_sync/test_model.py index 37e27a4..1ecd844 100644 --- a/firedantic/tests/tests_sync/test_model.py +++ b/firedantic/tests/tests_sync/test_model.py @@ -115,6 +115,36 @@ def test_find(configure_db, create_company, create_product) -> None: Product.find({"product_id": {"<>": "a"}}) +def test_count(configure_db, create_product) -> None: + for p in TEST_PRODUCTS: + create_product(**p) + + assert Product.count() == 4 + assert Product.count({"stock": {op.GTE: 1}}) == 3 + assert Product.count({"stock": {op.GTE: 2, op.LT: 4}}) == 2 + assert Product.count({"product_id": "missing"}) == 0 + + +def test_get_by_ids(configure_db, create_company) -> None: + company_a: Company = create_company(company_id="1234555-1") + company_b: Company = create_company(company_id="1231231-2") + + assert company_a.id + assert company_b.id + found = Company.get_by_ids([company_a.id, "missing", company_b.id]) + + assert {company.id for company in found} == {company_a.id, company_b.id} + assert {company.company_id for company in found} == { + company_a.company_id, + company_b.company_id, + } + + +def test_get_by_ids_with_invalid_id(configure_db) -> None: + with pytest.raises(ModelNotFoundError): + Company.get_by_ids([""]) + + def test_find_not_in(configure_db, create_company) -> None: ids = ["1234555-1", "1234567-8", "2131232-4", "4124432-4"] for company_id in ids: @@ -361,6 +391,18 @@ def test_custom_id_model(configure_db) -> None: assert m.bar == "bar" +def test_custom_id_model_get_by_doc_ids(configure_db) -> None: + c = CustomIDModel(bar="bar") # type: ignore + c.save() + assert c.foo + + models = CustomIDModel.get_by_doc_ids([c.foo, "missing"]) + + assert len(models) == 1 + assert models[0].foo == c.foo + assert models[0].bar == "bar" + + def test_custom_id_conflict(configure_db) -> None: CustomIDConflictModel(foo="foo", bar="bar").save() diff --git a/setup.cfg b/setup.cfg index 4e7cb69..b9c06b8 100644 --- a/setup.cfg +++ b/setup.cfg @@ -4,7 +4,7 @@ filterwarnings = ignore::DeprecationWarning # https://github.com/googleapis/python-firestore/issues/804 ignore::UserWarning:.*google.cloud.firestore_v1.base_collection.*: -asyncio_default_fixture_loop_scope = "function" +asyncio_default_fixture_loop_scope = function [flake8] # Black takes care of line formatting so flake8 should not mess with them From b965886f6ea09ee250c86eaf069f881e3da1d04f Mon Sep 17 00:00:00 2001 From: Harrison Zhao Date: Sun, 31 May 2026 03:12:38 -0700 Subject: [PATCH 2/2] implement increment --- README.md | 34 ++++++++++++++++ firedantic/_async/model.py | 45 +++++++++++++++++++++ firedantic/_sync/model.py | 45 +++++++++++++++++++++ firedantic/tests/tests_async/test_model.py | 31 ++++++++++++++ firedantic/tests/tests_sync/test_indexes.py | 4 +- firedantic/tests/tests_sync/test_model.py | 29 +++++++++++++ 6 files changed, 186 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index fd4b975..dd6e1ba 100644 --- a/README.md +++ b/README.md @@ -107,6 +107,34 @@ The query operators are found at If your model customizes the document ID field by extending `BareModel`, use `get_by_doc_ids()` instead of `get_by_ids()`. +### Atomic increments + +Numeric fields can be incremented atomically without first reloading and saving the +whole model. Pass a negative amount to decrement the value. + +```python +from firedantic import Model + + +class Product(Model): + __collection__ = "products" + product_id: str + stock: int + + +product = Product(product_id="abc-123", stock=10) +product.save() + +product.increment("stock", 5) +assert product.stock == 15 + +product.increment("stock", -3) +assert product.stock == 12 +``` + +This uses Firestore's atomic increment operation. If the field does not exist or the +current field value is not numeric, Firestore sets the field to the given amount. + ### Async usage Firedantic can also be used in an async way, like this: @@ -136,12 +164,16 @@ configure(client, prefix="firedantic-test-") class Person(AsyncModel): __collection__ = "persons" name: str + login_count: int = 0 async def main(): alice = Person(name="Alice") await alice.save() print(f"Saved Alice as {alice.id}") + await alice.increment("login_count") + print(f"Alice has logged in {alice.login_count} time(s)") + bob = Person(name="Bob") await bob.save() print(f"Saved Bob as {bob.id}") @@ -350,10 +382,12 @@ The following methods can be used in a transaction: - `Model.get_by_doc_ids(transaction=transaction)` - `Model.get_by_id(transaction=transaction)` - `Model.get_by_ids(transaction=transaction)` +- `Model.increment("field", amount, transaction=transaction)` - `Model.reload(transaction=transaction)` - `Model.save(transaction=transaction)` - `SubModel.get_by_id(transaction=transaction)` - `SubModel.get_by_ids(transaction=transaction)` +- `SubModel.increment("field", amount, transaction=transaction)` When using transactions, note that read operations must come before write operations. diff --git a/firedantic/_async/model.py b/firedantic/_async/model.py index afa2da6..5600b59 100644 --- a/firedantic/_async/model.py +++ b/firedantic/_async/model.py @@ -8,6 +8,7 @@ AsyncDocumentReference, DocumentSnapshot, FieldFilter, + Increment, ) from google.cloud.firestore_v1.async_query import AsyncQuery from google.cloud.firestore_v1.async_transaction import AsyncTransaction @@ -24,6 +25,7 @@ TAsyncBareModel = TypeVar("TAsyncBareModel", bound="AsyncBareModel") TAsyncBareSubModel = TypeVar("TAsyncBareSubModel", bound="AsyncBareSubModel") +Number = Union[int, float] logger = getLogger("firedantic") # https://firebase.google.com/docs/firestore/query-data/queries#query_operators @@ -100,6 +102,49 @@ async def save( await doc_ref.set(data) setattr(self, self.__document_id__, doc_ref.id) + async def increment( + self, + field: str, + amount: Number = 1, + transaction: Optional[AsyncTransaction] = None, + ) -> None: + """ + Atomically increments or decrements a numeric field in the database. + + :param field: Firestore field path to update. + :param amount: Amount to increment by. Use a negative value to decrement. + :param transaction: Optional transaction to use. + """ + doc_ref = self._get_doc_ref() + data = {field: Increment(amount)} + if transaction is not None: + transaction.update(doc_ref, data) + else: + await doc_ref.update(data) + self._increment_field_locally(field, amount) + + def _increment_field_locally(self, field: str, amount: Number) -> None: + if "." in field: + return + + model_field = field + if model_field not in self.__class__.model_fields: + for field_name, field_info in self.__class__.model_fields.items(): + if field_info.alias == field: + model_field = field_name + break + else: + return + + current_value = getattr(self, model_field, None) + if isinstance(current_value, (int, float)) and not isinstance( + current_value, bool + ): + value = current_value + amount + else: + value = amount + setattr(self, model_field, value) + async def delete(self, transaction: Optional[AsyncTransaction] = None) -> None: """ Deletes this model from the database. diff --git a/firedantic/_sync/model.py b/firedantic/_sync/model.py index 3a9f98c..f40adc1 100644 --- a/firedantic/_sync/model.py +++ b/firedantic/_sync/model.py @@ -8,6 +8,7 @@ DocumentReference, DocumentSnapshot, FieldFilter, + Increment, ) from google.cloud.firestore_v1.base_query import BaseQuery from google.cloud.firestore_v1.transaction import Transaction @@ -24,6 +25,7 @@ TBareModel = TypeVar("TBareModel", bound="BareModel") TBareSubModel = TypeVar("TBareSubModel", bound="BareSubModel") +Number = Union[int, float] logger = getLogger("firedantic") # https://firebase.google.com/docs/firestore/query-data/queries#query_operators @@ -100,6 +102,49 @@ def save( doc_ref.set(data) setattr(self, self.__document_id__, doc_ref.id) + def increment( + self, + field: str, + amount: Number = 1, + transaction: Optional[Transaction] = None, + ) -> None: + """ + Atomically increments or decrements a numeric field in the database. + + :param field: Firestore field path to update. + :param amount: Amount to increment by. Use a negative value to decrement. + :param transaction: Optional transaction to use. + """ + doc_ref = self._get_doc_ref() + data = {field: Increment(amount)} + if transaction is not None: + transaction.update(doc_ref, data) + else: + doc_ref.update(data) + self._increment_field_locally(field, amount) + + def _increment_field_locally(self, field: str, amount: Number) -> None: + if "." in field: + return + + model_field = field + if model_field not in self.__class__.model_fields: + for field_name, field_info in self.__class__.model_fields.items(): + if field_info.alias == field: + model_field = field_name + break + else: + return + + current_value = getattr(self, model_field, None) + if isinstance(current_value, (int, float)) and not isinstance( + current_value, bool + ): + value = current_value + amount + else: + value = amount + setattr(self, model_field, value) + def delete(self, transaction: Optional[Transaction] = None) -> None: """ Deletes this model from the database. diff --git a/firedantic/tests/tests_async/test_model.py b/firedantic/tests/tests_async/test_model.py index b29d010..f21877f 100644 --- a/firedantic/tests/tests_async/test_model.py +++ b/firedantic/tests/tests_async/test_model.py @@ -1,4 +1,5 @@ from operator import attrgetter +from typing import Optional from uuid import uuid4 import pytest @@ -132,6 +133,36 @@ async def test_count(configure_db, create_product) -> None: assert await Product.count({"product_id": "missing"}) == 0 +@pytest.mark.asyncio +async def test_increment_model_field(configure_db, create_product) -> None: + product = await create_product(stock=3) + assert product.id + + await product.increment("stock", 5) + assert product.stock == 8 + assert (await Product.get_by_id(product.id)).stock == 8 + + await product.increment("stock", -3) + await product.reload() + assert product.stock == 5 + + +@pytest.mark.asyncio +async def test_increment_missing_model_field(configure_db) -> None: + class Counter(AsyncModel): + __collection__ = "counters" + + visits: Optional[int] = None + + counter = Counter() + await counter.save(exclude_none=True) + assert counter.id + + await counter.increment("visits", 2) + assert counter.visits == 2 + assert (await Counter.get_by_id(counter.id)).visits == 2 + + @pytest.mark.asyncio async def test_get_by_ids(configure_db, create_company) -> None: company_a: Company = await create_company(company_id="1234555-1") diff --git a/firedantic/tests/tests_sync/test_indexes.py b/firedantic/tests/tests_sync/test_indexes.py index c4fb399..926eea8 100644 --- a/firedantic/tests/tests_sync/test_indexes.py +++ b/firedantic/tests/tests_sync/test_indexes.py @@ -7,10 +7,10 @@ from firedantic import ( CONFIGURATIONS, Model, - collection_group_index, - collection_index, set_up_composite_indexes, set_up_composite_indexes_and_ttl_policies, + collection_group_index, + collection_index, ) from firedantic.common import IndexField from firedantic.tests.tests_sync.conftest import MockListIndexOperation diff --git a/firedantic/tests/tests_sync/test_model.py b/firedantic/tests/tests_sync/test_model.py index 1ecd844..db0164b 100644 --- a/firedantic/tests/tests_sync/test_model.py +++ b/firedantic/tests/tests_sync/test_model.py @@ -1,4 +1,5 @@ from operator import attrgetter +from typing import Optional from uuid import uuid4 import pytest @@ -125,6 +126,34 @@ def test_count(configure_db, create_product) -> None: assert Product.count({"product_id": "missing"}) == 0 +def test_increment_model_field(configure_db, create_product) -> None: + product = create_product(stock=3) + assert product.id + + product.increment("stock", 5) + assert product.stock == 8 + assert (Product.get_by_id(product.id)).stock == 8 + + product.increment("stock", -3) + product.reload() + assert product.stock == 5 + + +def test_increment_missing_model_field(configure_db) -> None: + class Counter(Model): + __collection__ = "counters" + + visits: Optional[int] = None + + counter = Counter() + counter.save(exclude_none=True) + assert counter.id + + counter.increment("visits", 2) + assert counter.visits == 2 + assert (Counter.get_by_id(counter.id)).visits == 2 + + def test_get_by_ids(configure_db, create_company) -> None: company_a: Company = create_company(company_id="1234555-1") company_b: Company = create_company(company_id="1231231-2")