1 files changed, 18 insertions, 3 deletions
diff --git a/tests/test_mongoengine.py b/tests/test_mongoengine.py
index 803607a..148d274 100644
--- a/tests/test_mongoengine.py
+++ b/tests/test_mongoengine.py
@@ -31,6 +31,9 @@ try:
 except ImportError:
     mongoengine = None
 
+if os.environ.get('SKIP_MONGOENGINE') == '1':
+    mongoengine = None
+
 if mongoengine:
     from factory.mongoengine import MongoEngineFactory
 
@@ -42,12 +45,14 @@ if mongoengine:
         address = mongoengine.EmbeddedDocumentField(Address)
 
     class AddressFactory(MongoEngineFactory):
-        FACTORY_FOR = Address
+        class Meta:
+            model = Address
 
         street = factory.Sequence(lambda n: 'street%d' % n)
 
     class PersonFactory(MongoEngineFactory):
-        FACTORY_FOR = Person
+        class Meta:
+            model = Person
 
         name = factory.Sequence(lambda n: 'name%d' % n)
         address = factory.SubFactory(AddressFactory)
@@ -59,10 +64,20 @@ class MongoEngineTestCase(unittest.TestCase):
     db_name = os.environ.get('MONGO_DATABASE', 'factory_boy_test')
     db_host = os.environ.get('MONGO_HOST', 'localhost')
     db_port = int(os.environ.get('MONGO_PORT', '27017'))
+    server_timeout_ms = int(os.environ.get('MONGO_TIMEOUT', '300'))
 
     @classmethod
     def setUpClass(cls):
-        cls.db = mongoengine.connect(cls.db_name, host=cls.db_host, port=cls.db_port)
+        from pymongo import read_preferences as mongo_rp
+        cls.db = mongoengine.connect(
+            db=cls.db_name,
+            host=cls.db_host,
+            port=cls.db_port,
+            # PyMongo>=2.1 requires an explicit read_preference.
+            read_preference=mongo_rp.ReadPreference.PRIMARY,
+            # PyMongo>=2.1 has a 20s timeout, use 100ms instead
+            serverselectiontimeoutms=cls.server_timeout_ms,
+        )
 
     @classmethod
     def tearDownClass(cls):