Set execution profile for tests, remove batch counter cqlengine test

beltran · beltran · commit daffcd3bfeb7 · 2018-04-23T10:34:20.000-04:00
diff --git a/tests/integration/cqlengine/query/test_batch_query.py b/tests/integration/cqlengine/query/test_batch_query.py
@@ -220,7 +220,7 @@ def tearDown(self):
         drop_table(TestMultiKeyModel)
         drop_table(CounterBatchQueryModel)
 
-    @execute_count(6)
+    @execute_count(4)
     def test_cassandra_batch_type(self):
         """
         Tests the different types of `class: cassandra.query.BatchType`
@@ -239,15 +239,6 @@ def test_cassandra_batch_type(self):
         obj = TestMultiKeyModel.objects(partition=1)
         self.assertEqual(2, len(obj))
 
-        with BatchQuery(batch_type=cassandra_BatchType.COUNTER) as b:
-            CounterBatchQueryModel.batch(b).create(k=1, v=1)
-            CounterBatchQueryModel.batch(b).create(k=1, v=2)
-            CounterBatchQueryModel.batch(b).create(k=1, v=10)
-
-        obj = CounterBatchQueryModel.objects(k=1)
-        self.assertEqual(1, len(obj))
-        self.assertEqual(obj[0].v, 13)
-
         with BatchQuery(batch_type=cassandra_BatchType.LOGGED) as b:
             TestMultiKeyModel.batch(b).create(partition=1, cluster=1)
             TestMultiKeyModel.batch(b).create(partition=1, cluster=2)
diff --git a/tests/integration/cqlengine/test_ttl.py b/tests/integration/cqlengine/test_ttl.py
@@ -28,7 +28,7 @@
 from uuid import uuid4
 from cassandra.cqlengine import columns
 from cassandra.cqlengine.connection import get_session
-from tests.integration import CASSANDRA_VERSION, greaterthancass20
+from tests.integration import CASSANDRA_VERSION
 
 
 class BaseTTLTest(BaseCassEngTestCase):
diff --git a/tests/integration/simulacron/test_connection.py b/tests/integration/simulacron/test_connection.py
@@ -358,10 +358,12 @@ def test_driver_recovers_nework_isolation(self):
         listener = TrackDownListener()
 
         cluster = Cluster(['127.0.0.1'],
-                          load_balancing_policy=RoundRobinPolicy(),
                           idle_heartbeat_timeout=idle_heartbeat_timeout,
                           idle_heartbeat_interval=idle_heartbeat_interval,
-                          executor_threads=16)
+                          executor_threads=16,
+                          execution_profiles={
+                              EXEC_PROFILE_DEFAULT: ExecutionProfile(load_balancing_policy=RoundRobinPolicy())}
+                          )
         session = cluster.connect(wait_for_all_pools=True)
 
         cluster.register_listener(listener)
diff --git a/tests/integration/simulacron/test_policies.py b/tests/integration/simulacron/test_policies.py
@@ -384,6 +384,8 @@ def test_setting_retry_policy_to_statement(self):
 
     def set_cluster(self, retry_policy):
         self.cluster = Cluster(protocol_version=PROTOCOL_VERSION, compression=False,
-                          default_retry_policy=retry_policy)
+                               execution_profiles={
+                                   EXEC_PROFILE_DEFAULT: ExecutionProfile(retry_policy=retry_policy)}
+                               )
         self.session = self.cluster.connect(wait_for_all_pools=True)
         self.addCleanup(self.cluster.shutdown)
