Merge pull request #2445 from bagerard/fix_count_mongomock

Use Queryset._query instead of Cursor.__spec for count()
MongoEngine · Dec 13, 2020 · 2e284b9 · 2e284b9
2 parents aa3ff39 + c28bb51
commit 2e284b9
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 45 deletions.
diff --git a/docs/changelog.rst b/docs/changelog.rst
@@ -7,6 +7,7 @@ Development
 ===========
 - (Fill this out as you fix issues and develop your features).
 - Fix LazyReferenceField dereferencing in embedded documents #2426
+- Fix regarding the recent use of Cursor.__spec in .count() that was interfering with mongomock #2425
 
 Changes in 0.21.0
 =================

diff --git a/mongoengine/queryset/base.py b/mongoengine/queryset/base.py
@@ -420,7 +420,7 @@ def count(self, with_limit_and_skip=False):
 
         count = count_documents(
             collection=self._cursor.collection,
-            filter=self._cursor._Cursor__spec,
+            filter=self._query,
             **kwargs,
         )
 

diff --git a/tests/document/test_instance.py b/tests/document/test_instance.py
@@ -2900,50 +2900,32 @@ def __str__(self):
         # Checks
         assert ",".join([str(b) for b in Book.objects.all()]) == "1,2,3,4,5,6,7,8,9"
         # bob related books
-        assert (
-            ",".join(
-                [
-                    str(b)
-                    for b in Book.objects.filter(
-                        Q(extra__a=bob) | Q(author=bob) | Q(extra__b=bob)
-                    )
-                ]
-            )
-            == "1,2,3,4"
+        bob_books_qs = Book.objects.filter(
+            Q(extra__a=bob) | Q(author=bob) | Q(extra__b=bob)
         )
+        assert [str(b) for b in bob_books_qs] == ["1", "2", "3", "4"]
+        assert bob_books_qs.count() == 4
 
         # Susan & Karl related books
-        assert (
-            ",".join(
-                [
-                    str(b)
-                    for b in Book.objects.filter(
-                        Q(extra__a__all=[karl, susan])
-                        | Q(author__all=[karl, susan])
-                        | Q(extra__b__all=[karl.to_dbref(), susan.to_dbref()])
-                    )
-                ]
-            )
-            == "1"
+        susan_karl_books_qs = Book.objects.filter(
+            Q(extra__a__all=[karl, susan])
+            | Q(author__all=[karl, susan])
+            | Q(extra__b__all=[karl.to_dbref(), susan.to_dbref()])
         )
+        assert [str(b) for b in susan_karl_books_qs] == ["1"]
+        assert susan_karl_books_qs.count() == 1
 
         # $Where
-        assert (
-            ",".join(
-                [
-                    str(b)
-                    for b in Book.objects.filter(
-                        __raw__={
-                            "$where": """
+        custom_qs = Book.objects.filter(
+            __raw__={
+                "$where": """
                                             function(){
                                                 return this.name == '1' ||
                                                        this.name == '2';}"""
-                        }
-                    )
-                ]
-            )
-            == "1,2"
+            }
         )
+        assert [str(b) for b in custom_qs] == ["1", "2"]
+        assert custom_qs.count() == 2
 
     def test_switch_db_instance(self):
         register_connection("testdb-1", "mongoenginetest2")

diff --git a/tests/queryset/test_transform.py b/tests/queryset/test_transform.py
@@ -104,18 +104,17 @@ class BlogPost(Document):
         post = BlogPost(**data)
         post.save()
 
-        assert "postTitle" in BlogPost.objects(title=data["title"])._query
-        assert not ("title" in BlogPost.objects(title=data["title"])._query)
-        assert BlogPost.objects(title=data["title"]).count() == 1
+        qs = BlogPost.objects(title=data["title"])
+        assert qs._query == {"postTitle": data["title"]}
+        assert qs.count() == 1
 
-        assert "_id" in BlogPost.objects(pk=post.id)._query
-        assert BlogPost.objects(pk=post.id).count() == 1
+        qs = BlogPost.objects(pk=post.id)
+        assert qs._query == {"_id": post.id}
+        assert qs.count() == 1
 
-        assert (
-            "postComments.commentContent"
-            in BlogPost.objects(comments__content="test")._query
-        )
-        assert BlogPost.objects(comments__content="test").count() == 1
+        qs = BlogPost.objects(comments__content="test")
+        assert qs._query == {"postComments.commentContent": "test"}
+        assert qs.count() == 1
 
         BlogPost.drop_collection()