Handle limit and session in .count() method (#1040)

---------

Co-authored-by: ᑕᗩᑭTᗩIᑎᗰᗩᖇᐯᗴᒪ <CAPITAINMARVEL@users.noreply.github.com>

Author: CAPITAINMARVEL

beanie/odm/queries/find.py

@@ -71,7 +71,7 @@ class FindQuery(
     AggregationQueryType = AggregationQuery
 
     def __init__(self, document_model: Type["DocType"]):
-        self.document_model: Type["DocType"] = document_model
+        self.document_model = document_model
         self.find_expressions: List[Mapping[str, Any]] = []
         self.projection_model: Type[FindQueryResultType] = cast(
             Type[FindQueryResultType], self.document_model
@@ -147,9 +147,15 @@ async def count(self) -> int:
         Number of found documents
         :return: int
         """
+        kwargs = {}
+        if isinstance(self, FindMany):
+            if self.limit_number:
+                kwargs["limit"] = self.limit_number
+            if self.skip_number:
+                kwargs["skip"] = self.skip_number
         return (
             await self.document_model.get_motor_collection().count_documents(
-                self.get_filter_query()
+                self.get_filter_query(), session=self.session, **kwargs
             )
         )
 

tests/odm/documents/test_count.py

@@ -13,3 +13,11 @@ async def test_count_with_filter_query(documents):
     await documents(1, "cuatro", True)
     c = await DocumentTestModel.find_many({"test_str": "dos"}).count()
     assert c == 2
+
+
+async def test_count_with_limit(documents):
+    await documents(5, "five", True)
+    c = await DocumentTestModel.find_all().limit(1).count()
+    assert c == 1
+    d = await DocumentTestModel.find_all().count()
+    assert d == 5