test: Try to open/close database before each transaction

tools/integration_tests/tokenserver/test_node_assignment.py

@@ -43,9 +43,11 @@ def test_user_creation(self):
         self.assertEqual(self._count_users(), 1)
 
     def test_new_user_allocation(self):
+        self._db_connect()
         # Start with a clean database
         cursor = self._execute_sql('DELETE FROM nodes', ())
         cursor.close()
+        self.database.close()
 
         self._add_node(available=100, current_load=0, capacity=100, backoff=1,
                        node='https://node1')
@@ -73,9 +75,11 @@ def test_new_user_allocation(self):
         self.assertEqual(node['available'], 98)
 
     def test_successfully_releasing_node_capacity(self):
+        self._db_connect()
         # Start with a clean database
         cursor = self._execute_sql('DELETE FROM nodes', ())
         cursor.close()
+        self.database.close()
 
         node_id1 = self._add_node(available=0, current_load=99, capacity=100,
                                   node='https://node1')
@@ -116,9 +120,12 @@ def test_successfully_releasing_node_capacity(self):
         self.assertEqual(node5['available'], 0)
 
     def test_unsuccessfully_releasing_node_capacity(self):
+        self._db_connect()
+
         # Start with a clean database
         cursor = self._execute_sql('DELETE FROM nodes', ())
         cursor.close()
+        self.database.close()
 
         self._add_node(available=0, current_load=100, capacity=100,
                        node='https://node1')

tools/integration_tests/tokenserver/test_support.py

@@ -29,10 +29,7 @@ def setUpClass(cls):
         cls._build_auth_headers = cls._build_oauth_headers
 
     def setUp(self):
-        engine = create_engine(os.environ['SYNC_TOKENSERVER__DATABASE_URL'])
-        self.database = engine. \
-            execution_options(isolation_level='AUTOCOMMIT'). \
-            connect()
+        self._db_connect()
 
         host_url = urlparse.urlparse(self.TOKENSERVER_HOST)
         self.app = TestApp(self.TOKENSERVER_HOST, extra_environ={
@@ -54,8 +51,11 @@ def setUp(self):
 
         # Ensure we have a node with enough capacity to run the tests.
         self._add_node(capacity=100, node=self.NODE_URL, id=self.NODE_ID)
+        self.database.close()
 
     def tearDown(self):
+        self._db_connect()
+
         # And clean up at the end, for good measure.
         cursor = self._execute_sql(('DELETE FROM users'), ())
         cursor.close()
@@ -97,6 +97,7 @@ def _build_oauth_headers(self, generation=None, user='test',
 
     def _add_node(self, capacity=100, available=100, node=NODE_URL, id=None,
                   current_load=0, backoff=0, downed=0):
+        self._db_connect()
         query = 'INSERT INTO nodes (service, node, available, capacity, \
             current_load, backoff, downed'
         data = (self.service_id, node, available, capacity, current_load,
@@ -110,15 +111,18 @@ def _add_node(self, capacity=100, available=100, node=NODE_URL, id=None,
 
         cursor = self._execute_sql(query, data)
         cursor.close()
+        self.database.close()
 
         return self._last_insert_id()
 
     def _get_node(self, id):
+        self._db_connect()
         query = 'SELECT * FROM nodes WHERE id=%s'
         cursor = self._execute_sql(query, (id,))
         (id, service, node, available, current_load, capacity, downed,
          backoff) = cursor.fetchone()
         cursor.close()
+        self.database.close()
 
         return {
             'id': id,
@@ -132,23 +136,28 @@ def _get_node(self, id):
         }
 
     def _last_insert_id(self):
+        self._db_connect()
         cursor = self._execute_sql('SELECT LAST_INSERT_ID()', ())
         (id,) = cursor.fetchone()
         cursor.close()
+        self.database.close()
 
         return id
 
     def _add_service(self, service_name, pattern):
+        self._db_connect()
         query = 'INSERT INTO services (service, pattern) \
             VALUES(%s, %s)'
         cursor = self._execute_sql(query, (service_name, pattern))
         cursor.close()
+        self.database.close()
 
         return self._last_insert_id()
 
     def _add_user(self, email=None, generation=1234, client_state='aaaa',
                   created_at=None, nodeid=NODE_ID, keys_changed_at=1234,
                   replaced_at=None):
+        self._db_connect()
         query = '''
             INSERT INTO users (service, email, generation, client_state, \
                 created_at, nodeid, keys_changed_at, replaced_at)
@@ -162,16 +171,19 @@ def _add_user(self, email=None, generation=1234, client_state='aaaa',
                                     created_at, nodeid, keys_changed_at,
                                     replaced_at))
         cursor.close()
+        self.database.close()
 
         return self._last_insert_id()
 
     def _get_user(self, uid):
+        self._db_connect()
         query = 'SELECT * FROM users WHERE uid = %s'
         cursor = self._execute_sql(query, (uid,))
 
         (uid, service, email, generation, client_state, created_at,
          replaced_at, nodeid, keys_changed_at) = cursor.fetchone()
         cursor.close()
+        self.database.close()
 
         return {
             'uid': uid,
@@ -186,6 +198,7 @@ def _get_user(self, uid):
         }
 
     def _get_replaced_users(self, service_id, email):
+        self._db_connect()
         query = 'SELECT * FROM users WHERE service = %s AND email = %s AND \
             replaced_at IS NOT NULL'
         cursor = self._execute_sql(query, (service_id, email))
@@ -209,21 +222,26 @@ def _get_replaced_users(self, service_id, email):
             users.append(user_dict)
 
         cursor.close()
+        self.database.close()
         return users
 
     def _get_service_id(self, service):
+        self._db_connect()
         query = 'SELECT id FROM services WHERE service = %s'
         cursor = self._execute_sql(query, (service,))
         (service_id,) = cursor.fetchone()
         cursor.close()
+        self.database.close()
 
         return service_id
 
     def _count_users(self):
+        self._db_connect()
         query = 'SELECT COUNT(DISTINCT(uid)) FROM users'
         cursor = self._execute_sql(query, ())
         (count,) = cursor.fetchone()
         cursor.close()
+        self.database.close()
 
         return count
 
@@ -232,6 +250,12 @@ def _execute_sql(self, query, args):
 
         return cursor
 
+    def _db_connect(self):
+        engine = create_engine(os.environ['SYNC_TOKENSERVER__DATABASE_URL'])
+        self.database = engine. \
+            execution_options(isolation_level='AUTOCOMMIT'). \
+            connect()
+
     def unsafelyParseToken(self, token):
         # For testing purposes, don't check HMAC or anything...
         return json.loads(decode_token_bytes(token)[:-32].decode('utf8'))