Integration test for PYTHON-323 and other general test fixes

tests/integration/cqlengine/connections/test_connection.py

@@ -26,6 +26,8 @@ from cassandra.query import dict_factory
 
 from tests.integration import PROTOCOL_VERSION, execute_with_long_wait_retry
 from tests.integration.cqlengine.base import BaseCassEngTestCase
+from tests.integration.cqlengine import DEFAULT_KEYSPACE, setup_connection
+from cassandra.cqlengine import models
 
 
 class TestConnectModel(Model):
@@ -38,6 +40,7 @@ class ConnectionTest(BaseCassEngTestCase):
 
     @classmethod
     def setUpClass(cls):
+        cls.original_cluster = connection.cluster
         cls.keyspace1 = 'ctest1'
         cls.keyspace2 = 'ctest2'
         super(ConnectionTest, cls).setUpClass()
@@ -52,7 +55,11 @@ class ConnectionTest(BaseCassEngTestCase):
     def tearDownClass(cls):
         execute_with_long_wait_retry(cls.setup_session, "DROP KEYSPACE {0}".format(cls.keyspace1))
         execute_with_long_wait_retry(cls.setup_session, "DROP KEYSPACE {0}".format(cls.keyspace2))
+        models.DEFAULT_KEYSPACE = DEFAULT_KEYSPACE
+        cls.original_cluster.shutdown()
         cls.setup_cluster.shutdown()
+        setup_connection(DEFAULT_KEYSPACE)
+        models.DEFAULT_KEYSPACE
 
     def setUp(self):
         self.c = Cluster(protocol_version=PROTOCOL_VERSION)

tests/integration/cqlengine/management/test_compaction_settings.py

@@ -83,7 +83,7 @@ class AlterTableTest(BaseCassEngTestCase):
 
         table_meta = _get_table_metadata(tmp)
 
-        self.assertRegex(table_meta.export_as_string(), '.*SizeTieredCompactionStrategy.*')
+        self.assertRegexpMatches(table_meta.export_as_string(), '.*SizeTieredCompactionStrategy.*')
 
     def test_alter_options(self):
 
@@ -97,11 +97,11 @@ class AlterTableTest(BaseCassEngTestCase):
         drop_table(AlterTable)
         sync_table(AlterTable)
         table_meta = _get_table_metadata(AlterTable)
-        self.assertRegex(table_meta.export_as_string(), ".*'sstable_size_in_mb': '64'.*")
+        self.assertRegexpMatches(table_meta.export_as_string(), ".*'sstable_size_in_mb': '64'.*")
         AlterTable.__options__['compaction']['sstable_size_in_mb'] = '128'
         sync_table(AlterTable)
         table_meta = _get_table_metadata(AlterTable)
-        self.assertRegex(table_meta.export_as_string(), ".*'sstable_size_in_mb': '128'.*")
+        self.assertRegexpMatches(table_meta.export_as_string(), ".*'sstable_size_in_mb': '128'.*")
 
 
 class OptionsTest(BaseCassEngTestCase):

tests/integration/cqlengine/management/test_management.py

@@ -207,7 +207,7 @@ class TablePropertiesTests(BaseCassEngTestCase):
         option = 'no way will this ever be an option'
         try:
             ModelWithTableProperties.__options__[option] = 'what was I thinking?'
-            self.assertRaisesRegex(KeyError, "Invalid table option.*%s.*" % option, sync_table, ModelWithTableProperties)
+            self.assertRaisesRegexp(KeyError, "Invalid table option.*%s.*" % option, sync_table, ModelWithTableProperties)
         finally:
             ModelWithTableProperties.__options__.pop(option, None)
 

tests/integration/cqlengine/query/test_queryset.py

@@ -34,7 +34,7 @@ from cassandra.cqlengine.management import sync_table, drop_table
 from cassandra.cqlengine.models import Model
 from cassandra.cqlengine import columns
 from cassandra.cqlengine import query
-from cassandra.cqlengine.query import QueryException
+from cassandra.cqlengine.query import QueryException, BatchQuery
 from datetime import timedelta
 from datetime import tzinfo
 
@@ -1087,3 +1087,49 @@ class TestModelQueryWithDBField(BaseCassEngTestCase):
                 i = model.create(**values)
             self.assertEqual(model.objects(k0=i.k0, k1=i.k1).count(), len(clustering_values))
             self.assertEqual(model.objects(k0=i.k0, k1=i.k1, v1=0).count(), 1)
+
+
+class TestModelSmall(Model):
+
+    test_id = columns.Integer(primary_key=True)
+
+
+class TestModelQueryWithFetchSize(BaseCassEngTestCase):
+    """
+    Test FetchSize, and ensure that results are returned correctly
+    regardless of the paging size
+
+    @since 3.1
+    @jira_ticket PYTHON-324
+    @expected_result results are properly retrieved and the correct size
+
+    @test_category object_mapper
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        super(TestModelQueryWithFetchSize, cls).setUpClass()
+        sync_table(TestModelSmall)
+
+    @classmethod
+    def tearDownClass(cls):
+        super(TestModelQueryWithFetchSize, cls).tearDownClass()
+        drop_table(TestModelSmall)
+
+    def test_defaultFetchSize(self):
+        with BatchQuery() as b:
+            for i in range(5100):
+                TestModelSmall.batch(b).create(test_id=i)
+        self.assertEqual(len(TestModelSmall.objects.fetch_size(1)), 5100)
+        self.assertEqual(len(TestModelSmall.objects.fetch_size(500)), 5100)
+        self.assertEqual(len(TestModelSmall.objects.fetch_size(4999)), 5100)
+        self.assertEqual(len(TestModelSmall.objects.fetch_size(5000)), 5100)
+        self.assertEqual(len(TestModelSmall.objects.fetch_size(5001)), 5100)
+        self.assertEqual(len(TestModelSmall.objects.fetch_size(5100)), 5100)
+        self.assertEqual(len(TestModelSmall.objects.fetch_size(5101)), 5100)
+        self.assertEqual(len(TestModelSmall.objects.fetch_size(1)), 5100)
+
+        with self.assertRaises(QueryException):
+            TestModelSmall.objects.fetch_size(0)
+        with self.assertRaises(QueryException):
+            TestModelSmall.objects.fetch_size(-1)

tests/integration/cqlengine/test_batch_query.py

@@ -223,7 +223,7 @@ class BatchQueryCallbacksTests(BaseCassEngTestCase):
                 batch.execute()
             batch.execute()
         self.assertEqual(len(w), 2)  # package filter setup to warn always
-        self.assertRegex(str(w[0].message), r"^Batch.*multiple.*")
+        self.assertRegexpMatches(str(w[0].message), r"^Batch.*multiple.*")
 
     def test_disable_multiple_callback_warning(self):
         """