Merge pull request bitly#32 from ajdavis/close-test-sockets

Harden the unittests a little: close connections between tests, and stop the IOLoop before asserting.

asyncmongo/pool.py

@@ -39,14 +39,17 @@ def get_connection_pool(self, pool_id, *args, **kwargs):
     @classmethod
     def close_idle_connections(self, pool_id=None):
         """close idle connections to mongo"""
+        if not hasattr(self, '_pools'):
+            return
+
         if pool_id:
             if pool_id not in self._pools:
                 raise ProgrammingError("pool %r does not exist" % pool_id)
             else:
                 pool = self._pools[pool_id]
                 pool.close()
         else:
-            for pool in self._pools.items():
+            for pool_id, pool in self._pools.items():
                 pool.close()
 
 class ConnectionPool(object):

test/test_authentication.py

@@ -21,26 +21,26 @@ def test_authentication(self):
             db = asyncmongo.Client(pool_id='testauth', host='127.0.0.1', port=27017, dbname='test', dbuser='testuser', dbpass='testpass', maxconnections=2)
 
             def update_callback(response, error):
+                tornado.ioloop.IOLoop.instance().stop()
                 logging.info(response)
                 assert len(response) == 1
                 test_shunt.register_called('update')
-                tornado.ioloop.IOLoop.instance().stop()
-
+
             db.test_stats.update({"_id" : TEST_TIMESTAMP}, {'$inc' : {'test_count' : 1}}, upsert=True, callback=update_callback)
 
             tornado.ioloop.IOLoop.instance().start()
             test_shunt.assert_called('update')
 
             def query_callback(response, error):
+                tornado.ioloop.IOLoop.instance().stop()
                 logging.info(response)
                 logging.info(error)
                 assert error is None
                 assert isinstance(response, dict)
                 assert response['_id'] == TEST_TIMESTAMP
                 assert response['test_count'] == 1
                 test_shunt.register_called('retrieved')
-                tornado.ioloop.IOLoop.instance().stop()
-
+
             db.test_stats.find_one({"_id" : TEST_TIMESTAMP}, callback=query_callback)
             tornado.ioloop.IOLoop.instance().start()
             test_shunt.assert_called('retrieved')

test/test_connection.py

@@ -14,21 +14,21 @@ def test_query(self):
         db = asyncmongo.Client(pool_id='test_query', host='127.0.0.1', port=27017, dbname='test', mincached=3)
 
         def insert_callback(response, error):
+            tornado.ioloop.IOLoop.instance().stop()
             logging.info(response)
             assert len(response) == 1
             test_shunt.register_called('inserted')
-            tornado.ioloop.IOLoop.instance().stop()
-
+
         db.test_users.insert({"_id" : "test_connection.%d" % TEST_TIMESTAMP}, safe=True, callback=insert_callback)
 
         tornado.ioloop.IOLoop.instance().start()
         test_shunt.assert_called('inserted')
 
         def callback(response, error):
+            tornado.ioloop.IOLoop.instance().stop()
             assert len(response) == 1
             test_shunt.register_called('got_record')
-            tornado.ioloop.IOLoop.instance().stop()
-
+
         db.test_users.find({}, limit=1, callback=callback)
 
         tornado.ioloop.IOLoop.instance().start()

test/test_duplicate_insert.py

@@ -13,22 +13,22 @@ def test_duplicate_insert(self):
         db = asyncmongo.Client(pool_id='dup_insert', host='127.0.0.1', port=27017, dbname='test')
 
         def insert_callback(response, error):
+            tornado.ioloop.IOLoop.instance().stop()
             logging.info(response)
             assert len(response) == 1
             test_shunt.register_called('inserted')
-            tornado.ioloop.IOLoop.instance().stop()
-
+
         db.test_users.insert({"_id" : "duplicate_insert.%d" % TEST_TIMESTAMP}, callback=insert_callback)
 
         tornado.ioloop.IOLoop.instance().start()
         test_shunt.assert_called('inserted')
 
         def duplicate_callback(response, error):
+            tornado.ioloop.IOLoop.instance().stop()
             logging.info(response)
             if error:
                 test_shunt.register_called('dupe')
-            tornado.ioloop.IOLoop.instance().stop()
-
+
         db.test_users.insert({"_id" : "duplicate_insert.%d" % TEST_TIMESTAMP}, callback=duplicate_callback)
 
         tornado.ioloop.IOLoop.instance().start()

test/test_insert_delete.py

@@ -13,33 +13,33 @@ def test_insert(self):
         db = asyncmongo.Client(pool_id='testinsert', host='127.0.0.1', port=27017, dbname='test')
 
         def insert_callback(response, error):
+            tornado.ioloop.IOLoop.instance().stop()
             logging.info(response)
             assert len(response) == 1
             test_shunt.register_called('inserted')
-            tornado.ioloop.IOLoop.instance().stop()
-
+
         db.test_users.insert({"_id" : "insert.%d" % TEST_TIMESTAMP}, callback=insert_callback)
 
         tornado.ioloop.IOLoop.instance().start()
         test_shunt.assert_called('inserted')
 
         def query_callback(response, error):
+            tornado.ioloop.IOLoop.instance().stop()
             logging.info(response)
             assert len(response) == 1
             test_shunt.register_called('retrieved')
-            tornado.ioloop.IOLoop.instance().stop()
-
+
         db.test_users.find_one({"_id" : "insert.%d" % TEST_TIMESTAMP}, callback=query_callback)
         tornado.ioloop.IOLoop.instance().start()
         test_shunt.assert_called('retrieved')
 
 
         def delete_callback(response, error):
+            tornado.ioloop.IOLoop.instance().stop()
             logging.info(response)
             assert len(response) == 1
             test_shunt.register_called('deleted')
-            tornado.ioloop.IOLoop.instance().stop()
-
+
         db.test_users.remove({"_id" : "insert.%d" % TEST_TIMESTAMP}, callback=delete_callback)
         tornado.ioloop.IOLoop.instance().start()
         test_shunt.assert_called('deleted')

test/test_pooled_db.py

@@ -18,30 +18,29 @@ def test_pooled_db(self):
         test_users_collection = client.connection('test_users')
 
         def insert_callback(response, error):
+            tornado.ioloop.IOLoop.instance().stop()
             logging.info(response)
             assert len(response) == 1
             test_shunt.register_called('inserted')
-            tornado.ioloop.IOLoop.instance().stop()
-
+
         test_users_collection.insert({"_id" : "record_test.%d" % TEST_TIMESTAMP}, safe=True, callback=insert_callback)
 
         tornado.ioloop.IOLoop.instance().start()
         test_shunt.assert_called('inserted')
 
         def pool_callback(response, error):
-            assert len(response) == 1
-            test_shunt.register_called('pool1')
             if test_shunt.is_called('pool2'):
                 tornado.ioloop.IOLoop.instance().stop()
-
-        def pool_callback2(response, error):
             assert len(response) == 1
-            test_shunt.register_called('pool2')
-
+            test_shunt.register_called('pool1')
+
+        def pool_callback2(response, error):
             if test_shunt.is_called('pool1'):
                 # don't expect 2 finishes second
                 tornado.ioloop.IOLoop.instance().stop()
-
+            assert len(response) == 1
+            test_shunt.register_called('pool2')
+
         test_users_collection.find({}, limit=1, callback=pool_callback)
         test_users_collection.find({}, limit=1, callback=pool_callback2)
 

test/test_safe_updates.py

@@ -13,11 +13,11 @@ def test_update_safe(self):
         db = asyncmongo.Client(pool_id='testinsert', host='127.0.0.1', port=27017, dbname='test', maxconnections=2)
 
         def update_callback(response, error):
+            tornado.ioloop.IOLoop.instance().stop()
             logging.info(response)
             assert len(response) == 1
             test_shunt.register_called('update')
-            tornado.ioloop.IOLoop.instance().stop()
-
+
         # all of these should be called, but only one should have a callback
         # we also are checking that connections in the pool never increases >1 with max_connections=2
         # this is because connections for safe=False calls get put back in the pool immediated
@@ -31,13 +31,13 @@ def update_callback(response, error):
         test_shunt.assert_called('update')
 
         def query_callback(response, error):
+            tornado.ioloop.IOLoop.instance().stop()
             logging.info(response)
             assert isinstance(response, dict)
             assert response['_id'] == TEST_TIMESTAMP
             assert response['test_count'] == 5
             test_shunt.register_called('retrieved')
-            tornado.ioloop.IOLoop.instance().stop()
-
+
         db.test_stats.find_one({"_id" : TEST_TIMESTAMP}, callback=query_callback)
         tornado.ioloop.IOLoop.instance().start()
         test_shunt.assert_called('retrieved')

test/test_shunt.py

@@ -18,6 +18,7 @@
     sys.path.insert(0, app_dir)
 
 import asyncmongo
+import asyncmongo.pool
 # make sure we get the local asyncmongo
 assert asyncmongo.__file__.startswith(app_dir)
 
@@ -41,6 +42,7 @@ def setUp(self):
         sleep_time = 1 + (len(self.mongods) * 2)
         logging.info('waiting for mongod to start (sleeping %d seconds)' % sleep_time)
         time.sleep(sleep_time)
+        asyncmongo.pool.ConnectionPools.close_idle_connections()
 
     def tearDown(self):
         """teardown method that cleans up child mongod instances, and removes their temporary data files"""

test/test_slave_only.py

@@ -21,11 +21,11 @@ def test_query_slave(self):
             time.sleep(4)
 
             def update_callback(response, error):
+                tornado.ioloop.IOLoop.instance().stop()
                 logging.info(response)
                 assert len(response) == 1
                 test_shunt.register_called('update')
-                tornado.ioloop.IOLoop.instance().stop()
-
+
             masterdb.test_stats.update({"_id" : TEST_TIMESTAMP}, {'$inc' : {'test_count' : 1}}, upsert=True, callback=update_callback)
 
             tornado.ioloop.IOLoop.instance().start()
@@ -35,15 +35,15 @@ def update_callback(response, error):
             time.sleep(2.5)
 
             def query_callback(response, error):
+                tornado.ioloop.IOLoop.instance().stop()
                 logging.info(response)
                 logging.info(error)
                 assert error is None
                 assert isinstance(response, dict)
                 assert response['_id'] == TEST_TIMESTAMP
                 assert response['test_count'] == 1
                 test_shunt.register_called('retrieved')
-                tornado.ioloop.IOLoop.instance().stop()
-
+
             slavedb.test_stats.find_one({"_id" : TEST_TIMESTAMP}, callback=query_callback)
             tornado.ioloop.IOLoop.instance().start()
             test_shunt.assert_called('retrieved')