Bläddra i källkod

MongoDB backend tests now passing with pymongo 3.x. Closes #2589

Ask Solem 10 år sedan
förälder
incheckning
9bccd0291e
2 ändrade filer med 15 tillägg och 10 borttagningar
  1. 10 6
      celery/backends/mongodb.py
  2. 5 4
      celery/tests/backends/test_mongodb.py

+ 10 - 6
celery/backends/mongodb.py

@@ -28,7 +28,7 @@ from kombu.utils import cached_property
 
 from celery import states
 from celery.exceptions import ImproperlyConfigured
-from celery.five import string_t
+from celery.five import items, string_t
 from celery.utils.timeutils import maybe_timedelta
 
 from .base import BaseBackend
@@ -92,17 +92,21 @@ class MongoBackend(BaseBackend):
             self.options = dict(config, **config.pop('options', None) or {})
 
             # Set option defaults
-            if pymongo.version_tuple >= (3, ):
-                self.options.setdefault('maxPoolSize', self.max_pool_size)
-            else:
-                self.options.setdefault('max_pool_size', self.max_pool_size)
-                self.options.setdefault('auto_start_request', False)
+            for key, value in items(self._prepare_client_options()):
+                self.options.setdefault(key, value)
 
         self.url = url
         if self.url:
             # Specifying backend as an URL
             self.host = self.url
 
+    def _prepare_client_options(self):
+            if pymongo.version_tuple >= (3, ):
+                return {'maxPoolSize': self.max_pool_size}
+            else:  # pragma: no cover
+                return {'max_pool_size': max_pool_size,
+                        'auto_start_request': False}
+
     def _get_connection(self):
         """Connect to the MongoDB server."""
         if self._connection is None:

+ 5 - 4
celery/tests/backends/test_mongodb.py

@@ -98,8 +98,9 @@ class test_MongoBackend(AppCase):
 
             connection = self.backend._get_connection()
             mock_Connection.assert_called_once_with(
-                host='mongodb://localhost:27017', max_pool_size=10,
-                auto_start_request=False)
+                host='mongodb://localhost:27017',
+                **self.backend._prepare_client_options()
+            )
             self.assertEqual(sentinel.connection, connection)
 
     def test_get_connection_no_connection_mongodb_uri(self):
@@ -113,8 +114,8 @@ class test_MongoBackend(AppCase):
 
             connection = self.backend._get_connection()
             mock_Connection.assert_called_once_with(
-                host=mongodb_uri, max_pool_size=10,
-                auto_start_request=False)
+                host=mongodb_uri, **self.backend._prepare_client_options()
+            )
             self.assertEqual(sentinel.connection, connection)
 
     @patch('celery.backends.mongodb.MongoBackend._get_connection')