More test fixes

TheRealFalcon · TheRealFalcon · commit 8473c044c917 · 2020-03-27T11:01:36.000-05:00
diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py
@@ -507,6 +507,9 @@ def use_cluster(cluster_name, nodes, ipformat=None, start=True, workloads=None,
                 CCM_CLUSTER = DseCluster(path, cluster_name, **ccm_options)
                 CCM_CLUSTER.set_configuration_options({'start_native_transport': True})
                 CCM_CLUSTER.set_configuration_options({'batch_size_warn_threshold_in_kb': 5})
+                if Version(dse_version) >= Version('5.0'):
+                    CCM_CLUSTER.set_configuration_options({'enable_user_defined_functions': True})
+                    CCM_CLUSTER.set_configuration_options({'enable_scripted_user_defined_functions': True})
                 if Version(dse_version) >= Version('5.1'):
                     # For Inet4Address
                     CCM_CLUSTER.set_dse_configuration_options({
diff --git a/tests/integration/advanced/graph/test_graph.py b/tests/integration/advanced/graph/test_graph.py
@@ -19,7 +19,7 @@
 from cassandra.protocol import SyntaxException
 from cassandra.policies import WhiteListRoundRobinPolicy
 from cassandra.cluster import NoHostAvailable
-from cassandra.cluster import EXEC_PROFILE_GRAPH_DEFAULT, GraphExecutionProfile,
+from cassandra.cluster import EXEC_PROFILE_GRAPH_DEFAULT, GraphExecutionProfile
 from cassandra.graph import single_object_row_factory, Vertex, graph_object_row_factory, \
     graph_graphson2_row_factory, graph_graphson3_row_factory
 from cassandra.util import SortedSet
diff --git a/tests/integration/advanced/test_unixsocketendpoint.py b/tests/integration/advanced/test_unixsocketendpoint.py
@@ -65,7 +65,7 @@ def setUpClass(cls):
         lbp = UnixSocketWhiteListRoundRobinPolicy([UNIX_SOCKET_PATH])
         ep = ExecutionProfile(load_balancing_policy=lbp)
         endpoint = UnixSocketEndPoint(UNIX_SOCKET_PATH)
-        cls.cluster = TestCluster([endpoint], execution_profiles={EXEC_PROFILE_DEFAULT: ep})
+        cls.cluster = TestCluster(contact_points=[endpoint], execution_profiles={EXEC_PROFILE_DEFAULT: ep})
 
     @classmethod
     def tearDownClass(cls):
diff --git a/tests/integration/standard/test_query.py b/tests/integration/standard/test_query.py
@@ -24,7 +24,7 @@
 from cassandra import ConsistencyLevel, Unavailable, InvalidRequest, cluster
 from cassandra.query import (PreparedStatement, BoundStatement, SimpleStatement,
                              BatchStatement, BatchType, dict_factory, TraceUnavailable)
-from cassandra.cluster import NoHostAvailable, ExecutionProfile, EXEC_PROFILE_DEFAULT
+from cassandra.cluster import NoHostAvailable, ExecutionProfile, EXEC_PROFILE_DEFAULT, Cluster
 from cassandra.policies import HostDistance, RoundRobinPolicy, WhiteListRoundRobinPolicy
 from tests.integration import use_singledc, PROTOCOL_VERSION, BasicSharedKeyspaceUnitTestCase, \
     greaterthanprotocolv3, MockLoggingHandler, get_supported_protocol_versions, local, get_cluster, setup_keyspace, \
@@ -488,7 +488,7 @@ def test_prepared_metadata_generation(self):
         base_line = None
         for proto_version in get_supported_protocol_versions():
             beta_flag = True if proto_version in ProtocolVersion.BETA_VERSIONS else False
-            cluster = TestCluster()
+            cluster = Cluster(protocol_version=proto_version, allow_beta_protocol_version=beta_flag)
 
             session = cluster.connect()
             select_statement = session.prepare("SELECT * FROM system.local")
