Merge pull request #2477 from lukasloetkolben/feature_allow_disk_use

Feature: "allow_disk_use"

tests/queryset/test_queryset.py

@@ -21,6 +21,7 @@ from mongoengine.queryset import (
     QuerySetManager,
     queryset_manager,
 )
+from tests.utils import requires_mongodb_gte_44
 
 
 class db_ops_tracker(query_counter):
@@ -5656,6 +5657,31 @@ class TestQueryset(unittest.TestCase):
         qs = self.Person.objects().timeout(False)
         assert qs._cursor_args == {"no_cursor_timeout": True}
 
+    @requires_mongodb_gte_44
+    def test_allow_disk_use(self):
+        qs = self.Person.objects()
+        assert qs._cursor_args == {}
+
+        qs = self.Person.objects().allow_disk_use(False)
+        assert qs._cursor_args == {}
+
+        qs = self.Person.objects().allow_disk_use(True)
+        assert qs._cursor_args == {"allow_disk_use": True}
+
+        # Test if allow_disk_use changes the results
+        self.Person.drop_collection()
+        self.Person.objects.create(name="Foo", age=12)
+        self.Person.objects.create(name="Baz", age=17)
+        self.Person.objects.create(name="Bar", age=13)
+
+        qs_disk = self.Person.objects().order_by("age").allow_disk_use(True)
+        qs = self.Person.objects().order_by("age")
+
+        assert qs_disk.count() == qs.count()
+
+        for index in range(qs_disk.count()):
+            assert qs_disk[index] == qs[index]
+
 
 if __name__ == "__main__":
     unittest.main()

tests/utils.py

@@ -1,3 +1,4 @@
+import operator
 import unittest
 
 import pytest
@@ -33,6 +34,10 @@ def get_as_pymongo(doc):
     return doc.__class__.objects.as_pymongo().get(id=doc.id)
 
 
+def requires_mongodb_gte_44(func):
+    return _decorated_with_ver_requirement(func, (4, 4), oper=operator.ge)
+
+
 def _decorated_with_ver_requirement(func, mongo_version_req, oper):
     """Return a MongoDB version requirement decorator.