diff --git a/beanie/odm/queries/find.py b/beanie/odm/queries/find.py index d9a981c1..c7d426f5 100644 --- a/beanie/odm/queries/find.py +++ b/beanie/odm/queries/find.py @@ -689,14 +689,7 @@ async def count(self) -> int: if self.fetch_links: aggregation_pipeline: List[ Dict[str, Any] - ] = construct_lookup_queries(self.document_model) - - aggregation_pipeline.append({"$match": self.get_filter_query()}) - - if self.skip_number != 0: - aggregation_pipeline.append({"$skip": self.skip_number}) - if self.limit_number != 0: - aggregation_pipeline.append({"$limit": self.limit_number}) + ] = self.build_aggregation_pipeline() aggregation_pipeline.append({"$count": "count"}) diff --git a/tests/odm/test_relations.py b/tests/odm/test_relations.py index e3946afd..d3995ff7 100644 --- a/tests/odm/test_relations.py +++ b/tests/odm/test_relations.py @@ -557,33 +557,43 @@ async def test_with_chaining_aggregation(self): assert addresses_count[0] == {"count": 10} async def test_with_chaining_aggregation_and_text_search(self): + # ARRANGE + NUM_DOCS = 10 + NUM_WITH_LOWER = 5 linked_document = LinkDocumentForTextSeacrh(i=1) await linked_document.insert() - for i in range(10): + for i in range(NUM_DOCS): await DocumentWithTextIndexAndLink( - s="lower" if i < 5 else "UPPER", link=linked_document + s="lower" if i < NUM_WITH_LOWER else "UPPER", + link=linked_document, ).insert() linked_document_2 = LinkDocumentForTextSeacrh(i=2) await linked_document_2.insert() - for i in range(10): + for i in range(NUM_DOCS): await DocumentWithTextIndexAndLink( - s="lower" if i < 5 else "UPPER", link=linked_document_2 + s="lower" if i < NUM_WITH_LOWER else "UPPER", + link=linked_document_2, ).insert() - document_count = ( - await DocumentWithTextIndexAndLink.find( - {"$text": {"$search": "lower"}}, - DocumentWithTextIndexAndLink.link.i == 1, - fetch_links=True, - ) - .aggregate([{"$count": "count"}]) - .to_list() + # ACT + query = DocumentWithTextIndexAndLink.find( + {"$text": {"$search": "lower"}}, + DocumentWithTextIndexAndLink.link.i == 1, + fetch_links=True, ) - assert document_count[0] == {"count": 5} + # Test both aggregation and count methods + document_count_aggregation = await query.aggregate( + [{"$count": "count"}] + ).to_list() + document_count = await query.count() + + # ASSERT + assert document_count_aggregation[0] == {"count": NUM_WITH_LOWER} + assert document_count == NUM_WITH_LOWER async def test_with_extra_allow(self, houses): res = await House.find(fetch_links=True).to_list()