From 72ebaa52e9330a063d7d45a72ce65710498ef99f Mon Sep 17 00:00:00 2001 From: Tom Floyer Date: Sat, 8 Dec 2018 22:10:10 +0300 Subject: [PATCH] Ensure that QuerySet.only() works correctly after QuerySet.count() This test checks if .only() method of QuerySet instance works after using .count() method. --- tests/queryset/queryset.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/queryset/queryset.py b/tests/queryset/queryset.py index 9e251af7..b44c36df 100644 --- a/tests/queryset/queryset.py +++ b/tests/queryset/queryset.py @@ -4720,6 +4720,28 @@ class QuerySetTest(unittest.TestCase): 'password_salt').only('email').to_json() self.assertEqual('[{"email": "ross@example.com"}]', serialized_user) + def test_only_after_count(self): + """Test that only() works after count()""" + + class User(Document): + name = StringField() + age = IntField() + address = StringField() + User.drop_collection() + User(name="User", age=50, + address="Moscow, Russia").save() + + user_queryset = User.objects(age=50) + + result = user_queryset.only("name", "age").as_pymongo().first() + self.assertEqual(result, {"name": "User", "age": 50}) + + result = user_queryset.count() + self.assertEqual(result, 1) + + result = user_queryset.only("name", "age").as_pymongo().first() + self.assertEqual(result, {"name": "User", "age": 50}) + def test_no_dereference(self): class Organization(Document):