Skip to content

Commit 50ffa80

Browse files
author
erdenezul
authored
Merge pull request #2008 from bagerard/refactor_deprecated_pymongo_test
refactored deprecated pymongo methods in tests
2 parents 7ef688b + b4fe0b3 commit 50ffa80

17 files changed

+106
-82
lines changed

‎mongoengine/connection.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
from pymongo import MongoClient, ReadPreference, uri_parser
22
import six
33

4-
from mongoengine.python_support import IS_PYMONGO_3
4+
from mongoengine.pymongo_support import IS_PYMONGO_3
55

66
__all__ = ['MongoEngineConnectionError', 'connect', 'register_connection',
7-
'DEFAULT_CONNECTION_NAME']
7+
'DEFAULT_CONNECTION_NAME', 'get_db']
88

99

1010
DEFAULT_CONNECTION_NAME = 'default'

‎mongoengine/context_managers.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55

66
from mongoengine.common import _import_class
77
from mongoengine.connection import DEFAULT_CONNECTION_NAME, get_db
8+
from mongoengine.pymongo_support import count_documents
89

910
__all__ = ('switch_db', 'switch_collection', 'no_dereference',
1011
'no_sub_classes', 'query_counter', 'set_write_concern')
@@ -237,7 +238,7 @@ def _get_count(self):
237238
and substracting the queries issued by this context. In fact everytime this is called, 1 query is
238239
issued so we need to balance that
239240
"""
240-
count = self.db.system.profile.find(self._ignored_query).count() - self._ctx_query_counter
241+
count = count_documents(self.db.system.profile, self._ignored_query) - self._ctx_query_counter
241242
self._ctx_query_counter += 1 # Account for the query we just issued to gather the information
242243
return count
243244

‎mongoengine/document.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
switch_db)
1919
from mongoengine.errors import (InvalidDocumentError, InvalidQueryError,
2020
SaveConditionError)
21-
from mongoengine.python_support import IS_PYMONGO_3
21+
from mongoengine.pymongo_support import IS_PYMONGO_3, list_collection_names
2222
from mongoengine.queryset import (NotUniqueError, OperationError,
2323
QuerySet, transform)
2424

@@ -228,7 +228,7 @@ def _get_capped_collection(cls):
228228

229229
# If the collection already exists and has different options
230230
# (i.e. isn't capped or has different max/size), raise an error.
231-
if collection_name in db.collection_names():
231+
if collection_name in list_collection_names(db, include_system_collections=True):
232232
collection = db[collection_name]
233233
options = collection.options()
234234
if (

‎mongoengine/pymongo_support.py‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
"""
2+
Helper functions, constants, and types to aid with PyMongo v2.7 - v3.x support.
3+
"""
4+
import pymongo
5+
6+
_PYMONGO_37 = (3, 7)
7+
8+
PYMONGO_VERSION = tuple(pymongo.version_tuple[:2])
9+
10+
IS_PYMONGO_3 = PYMONGO_VERSION[0] >= 3
11+
IS_PYMONGO_GTE_37 = PYMONGO_VERSION >= _PYMONGO_37
12+
13+
14+
def count_documents(collection, filter):
15+
"""Pymongo>3.7 deprecates count in favour of count_documents"""
16+
if IS_PYMONGO_GTE_37:
17+
return collection.count_documents(filter)
18+
else:
19+
count = collection.find(filter).count()
20+
return count
21+
22+
23+
def list_collection_names(db, include_system_collections=False):
24+
"""Pymongo>3.7 deprecates collection_names in favour of list_collection_names"""
25+
if IS_PYMONGO_GTE_37:
26+
collections = db.list_collection_names()
27+
else:
28+
collections = db.collection_names()
29+
30+
if not include_system_collections:
31+
collections = [c for c in collections if not c.startswith('system.')]
32+
33+
return collections

‎mongoengine/python_support.py‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,8 @@
11
"""
2-
Helper functions, constants, and types to aid with Python v2.7 - v3.x and
3-
PyMongo v2.7 - v3.x support.
2+
Helper functions, constants, and types to aid with Python v2.7 - v3.x support
43
"""
5-
import pymongo
64
import six
75

8-
9-
IS_PYMONGO_3 = pymongo.version_tuple[0] >= 3
10-
116
# six.BytesIO resolves to StringIO.StringIO in Py2 and io.BytesIO in Py3.
127
StringIO = six.BytesIO
138

‎mongoengine/queryset/base.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121
from mongoengine.context_managers import set_write_concern, switch_db
2222
from mongoengine.errors import (InvalidQueryError, LookUpError,
2323
NotUniqueError, OperationError)
24-
from mongoengine.python_support import IS_PYMONGO_3
24+
from mongoengine.pymongo_support import IS_PYMONGO_3
2525
from mongoengine.queryset import transform
2626
from mongoengine.queryset.field_list import QueryFieldList
2727
from mongoengine.queryset.visitor import Q, QNode

‎mongoengine/queryset/transform.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
from mongoengine.common import _import_class
1111
from mongoengine.connection import get_connection
1212
from mongoengine.errors import InvalidQueryError
13-
from mongoengine.python_support import IS_PYMONGO_3
13+
from mongoengine.pymongo_support import IS_PYMONGO_3
1414

1515
__all__ = ('query', 'update')
1616

‎tests/document/class_methods.py‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
import unittest
33

44
from mongoengine import *
5+
from mongoengine.pymongo_support import list_collection_names
56

67
from mongoengine.queryset import NULLIFY, PULL
78
from mongoengine.connection import get_db
@@ -27,9 +28,7 @@ class Person(Document):
2728
self.Person = Person
2829

2930
def tearDown(self):
30-
for collection in self.db.collection_names():
31-
if 'system.' in collection:
32-
continue
31+
for collection in list_collection_names(self.db):
3332
self.db.drop_collection(collection)
3433

3534
def test_definition(self):
@@ -66,10 +65,10 @@ def test_drop_collection(self):
6665
"""
6766
collection_name = 'person'
6867
self.Person(name='Test').save()
69-
self.assertIn(collection_name, self.db.collection_names())
68+
self.assertIn(collection_name, list_collection_names(self.db))
7069

7170
self.Person.drop_collection()
72-
self.assertNotIn(collection_name, self.db.collection_names())
71+
self.assertNotIn(collection_name, list_collection_names(self.db))
7372

7473
def test_register_delete_rule(self):
7574
"""Ensure that register delete rule adds a delete rule to the document
@@ -340,7 +339,7 @@ class Person(Document):
340339
meta = {'collection': collection_name}
341340

342341
Person(name="Test User").save()
343-
self.assertIn(collection_name, self.db.collection_names())
342+
self.assertIn(collection_name, list_collection_names(self.db))
344343

345344
user_obj = self.db[collection_name].find_one()
346345
self.assertEqual(user_obj['name'], "Test User")
@@ -349,7 +348,7 @@ class Person(Document):
349348
self.assertEqual(user_obj.name, "Test User")
350349

351350
Person.drop_collection()
352-
self.assertNotIn(collection_name, self.db.collection_names())
351+
self.assertNotIn(collection_name, list_collection_names(self.db))
353352

354353
def test_collection_name_and_primary(self):
355354
"""Ensure that a collection with a specified name may be used.

‎tests/document/delta.py‎

Lines changed: 5 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,16 +3,14 @@
33

44
from bson import SON
55
from mongoengine import *
6-
from mongoengine.connection import get_db
6+
from mongoengine.pymongo_support import list_collection_names
7+
from tests.utils import MongoDBTestCase
78

8-
__all__ = ("DeltaTest",)
99

10-
11-
class DeltaTest(unittest.TestCase):
10+
class DeltaTest(MongoDBTestCase):
1211

1312
def setUp(self):
14-
connect(db='mongoenginetest')
15-
self.db = get_db()
13+
super(DeltaTest, self).setUp()
1614

1715
class Person(Document):
1816
name = StringField()
@@ -25,9 +23,7 @@ class Person(Document):
2523
self.Person = Person
2624

2725
def tearDown(self):
28-
for collection in self.db.collection_names():
29-
if 'system.' in collection:
30-
continue
26+
for collection in list_collection_names(self.db):
3127
self.db.drop_collection(collection)
3228

3329
def test_delta(self):

‎tests/document/inheritance.py‎

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -6,23 +6,18 @@
66

77
from mongoengine import (BooleanField, Document, EmbeddedDocument,
88
EmbeddedDocumentField, GenericReferenceField,
9-
IntField, ReferenceField, StringField, connect)
10-
from mongoengine.connection import get_db
9+
IntField, ReferenceField, StringField)
10+
from mongoengine.pymongo_support import list_collection_names
11+
from tests.utils import MongoDBTestCase
1112
from tests.fixtures import Base
1213

1314
__all__ = ('InheritanceTest', )
1415

1516

16-
class InheritanceTest(unittest.TestCase):
17-
18-
def setUp(self):
19-
connect(db='mongoenginetest')
20-
self.db = get_db()
17+
class InheritanceTest(MongoDBTestCase):
2118

2219
def tearDown(self):
23-
for collection in self.db.collection_names():
24-
if 'system.' in collection:
25-
continue
20+
for collection in list_collection_names(self.db):
2621
self.db.drop_collection(collection)
2722

2823
def test_constructor_cls(self):

0 commit comments

Comments
 (0)