Add tests/integration/standard/test_row_factories.py

Author: aboudreault

tests/integration/standard/test_row_factories.py

@@ -19,7 +19,7 @@
 except ImportError:
     import unittest # noqa
 
-from cassandra.cluster import Cluster, ResultSet
+from cassandra.cluster import Cluster, ResultSet, ExecutionProfile, EXEC_PROFILE_DEFAULT
 from cassandra.query import tuple_factory, named_tuple_factory, dict_factory, ordered_dict_factory
 from cassandra.util import OrderedDict
 
@@ -31,8 +31,7 @@ def setup_module():
 class NameTupleFactory(BasicSharedKeyspaceUnitTestCase):
 
     def setUp(self):
-        super(NameTupleFactory, self).setUp()
-        self.session.row_factory = named_tuple_factory
+        self.common_setup(1, execution_profiles={EXEC_PROFILE_DEFAULT: ExecutionProfile(row_factory=named_tuple_factory)})
         ddl = '''
                 CREATE TABLE {0}.{1} (
                     k int PRIMARY KEY,
@@ -78,38 +77,22 @@ class RowFactoryTests(BasicSharedKeyspaceUnitTestCaseWFunctionTable):
     """
     Test different row_factories and access code
     """
-    def setUp(self):
-        super(RowFactoryTests, self).setUp()
-        self.insert1 = '''
-            INSERT INTO {0}.{1}
-                ( k , v )
-            VALUES
-                ( 1 , 1 )
-        '''.format(self.keyspace_name, self.function_table_name)
-
-        self.insert2 = '''
-            INSERT INTO {0}.{1}
-                ( k , v )
-            VALUES
-                ( 2 , 2 )
-        '''.format(self.keyspace_name, self.function_table_name)
-
-        self.select = '''
-            SELECT * FROM {0}.{1}
-        '''.format(self.keyspace_name, self.function_table_name)
-
-    def tearDown(self):
-        self.drop_function_table()
+    @classmethod
+    def setUpClass(cls):
+        cls.common_setup(rf=1, create_class_table=True)
+        q = "INSERT INTO {0}.{1} (k, v) VALUES (%s, %s)".format(cls.ks_name, cls.ks_name)
+        cls.session.execute(q, (1, 1))
+        cls.session.execute(q, (2, 2))
+        cls.select = "SELECT * FROM {0}.{1}".format(cls.ks_name, cls.ks_name)
+
+    def _results_from_row_factory(self, row_factory):
+        cluster = Cluster(protocol_version=PROTOCOL_VERSION,
+                          execution_profiles={EXEC_PROFILE_DEFAULT: ExecutionProfile(row_factory=row_factory)})
+        with cluster:
+            return cluster.connect().execute(self.select)
 
     def test_tuple_factory(self):
-        session = self.session
-        session.row_factory = tuple_factory
-
-        session.execute(self.insert1)
-        session.execute(self.insert2)
-
-        result = session.execute(self.select)
-
+        result = self._results_from_row_factory(tuple_factory)
         self.assertIsInstance(result, ResultSet)
         self.assertIsInstance(result[0], tuple)
 
@@ -122,14 +105,7 @@ def test_tuple_factory(self):
         self.assertEqual(result[1][0], 2)
 
     def test_named_tuple_factory(self):
-        session = self.session
-        session.row_factory = named_tuple_factory
-
-        session.execute(self.insert1)
-        session.execute(self.insert2)
-
-        result = session.execute(self.select)
-
+        result = self._results_from_row_factory(named_tuple_factory)
         self.assertIsInstance(result, ResultSet)
         result = list(result)
 
@@ -141,17 +117,10 @@ def test_named_tuple_factory(self):
         self.assertEqual(result[1].k, result[1].v)
         self.assertEqual(result[1].k, 2)
 
-    def test_dict_factory(self):
-        session = self.session
-        session.row_factory = dict_factory
-
-        session.execute(self.insert1)
-        session.execute(self.insert2)
-
-        result = session.execute(self.select)
-
+    def _test_dict_factory(self, row_factory, row_type):
+        result = self._results_from_row_factory(row_factory)
         self.assertIsInstance(result, ResultSet)
-        self.assertIsInstance(result[0], dict)
+        self.assertIsInstance(result[0], row_type)
 
         for row in result:
             self.assertEqual(row['k'], row['v'])
@@ -161,25 +130,11 @@ def test_dict_factory(self):
         self.assertEqual(result[1]['k'], result[1]['v'])
         self.assertEqual(result[1]['k'], 2)
 
-    def test_ordered_dict_factory(self):
-        session = self.session
-        session.row_factory = ordered_dict_factory
-
-        session.execute(self.insert1)
-        session.execute(self.insert2)
-
-        result = session.execute(self.select)
-
-        self.assertIsInstance(result, ResultSet)
-        self.assertIsInstance(result[0], OrderedDict)
-
-        for row in result:
-            self.assertEqual(row['k'], row['v'])
+    def test_dict_factory(self):
+        self._test_dict_factory(dict_factory, dict)
 
-        self.assertEqual(result[0]['k'], result[0]['v'])
-        self.assertEqual(result[0]['k'], 1)
-        self.assertEqual(result[1]['k'], result[1]['v'])
-        self.assertEqual(result[1]['k'], 2)
+    def test_ordered_dict_factory(self):
+        self._test_dict_factory(ordered_dict_factory, OrderedDict)
 
     def test_generator_row_factory(self):
         """
@@ -201,7 +156,12 @@ def _gen_row_factory(rows):
         session = self.session
         session.row_factory = generator_row_factory
 
-        session.execute(self.insert1)
+        session.execute('''
+             INSERT INTO {0}.{1}
+                 ( k , v )
+             VALUES
+                 ( 1 , 1 )
+        '''.format(self.keyspace_name, self.function_table_name))
         result = session.execute(self.select)
         self.assertIsInstance(result, ResultSet)
         first_row = result.one()
@@ -218,8 +178,7 @@ def setup_class(cls):
         cls.session = cls.cluster.connect()
         cls._cass_version, cls._cql_version = get_server_versions()
         ddl = '''
-            CREATE TABLE test1rf.table_num_col ( key blob PRIMARY KEY, "626972746864617465" blob )
-            WITH COMPACT STORAGE'''
+            CREATE TABLE test1rf.table_num_col ( key blob PRIMARY KEY, "626972746864617465" blob )'''
         cls.session.execute(ddl)
 
     @classmethod
@@ -252,8 +211,9 @@ def test_can_select_with_dict_factory(self):
         """
         can SELECT numeric column  using  dict_factory
         """
-        self.session.row_factory = dict_factory
-        try:
-            self.session.execute('SELECT * FROM test1rf.table_num_col')
-        except ValueError as e:
-            self.fail("Unexpected ValueError exception: %s" % e.message)
+        with Cluster(protocol_version=PROTOCOL_VERSION,
+                     execution_profiles={EXEC_PROFILE_DEFAULT: ExecutionProfile(row_factory=dict_factory)}) as cluster:
+            try:
+                cluster.connect().execute('SELECT * FROM test1rf.table_num_col')
+            except ValueError as e:
+                self.fail("Unexpected ValueError exception: %s" % e.message)