fix 3.11 failures after CASSANDRA-10786

mambocab · mambocab · commit c4efbe0769f8 · 2017-10-12T14:26:33.000-04:00
diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py
@@ -194,8 +194,10 @@ def get_supported_protocol_versions():
     3.10 -> 5(beta),4,3
 `   """
     version = Version(CASSANDRA_VERSION)
-    if version >= Version('3.10'):
+    if version >= Version('4.0'):
         return (3, 4, 5)
+    elif version >= Version('3.10'):
+        return (3, 4)
     elif version >= Version('3.0'):
         return (3, 4)
     elif version >= Version('2.2'):
diff --git a/tests/integration/long/test_failure_types.py b/tests/integration/long/test_failure_types.py
@@ -65,20 +65,12 @@ def setUp(self):
         """
         Test is skipped if run with native protocol version <4
         """
-        self.support_v5 = True
         if PROTOCOL_VERSION < 4:
             raise unittest.SkipTest(
                 "Native protocol 4,0+ is required for custom payloads, currently using %r"
                 % (PROTOCOL_VERSION,))
-        try:
-            self.cluster = Cluster(protocol_version=ProtocolVersion.MAX_SUPPORTED, allow_beta_protocol_version=True)
-            self.session = self.cluster.connect()
-        except NoHostAvailable:
-            log.info("Protocol Version 5 not supported,")
-            self.cluster = Cluster(protocol_version=PROTOCOL_VERSION)
-            self.session = self.cluster.connect()
-            self.support_v5 = False
-
+        self.cluster = Cluster(protocol_version=PROTOCOL_VERSION)
+        self.session = self.cluster.connect()
         self.nodes_currently_failing = []
         self.node1, self.node2, self.node3 = get_cluster().nodes.values()
 
@@ -157,7 +149,7 @@ def _perform_cql_statement(self, text, consistency_level, expected_exception, se
         else:
             with self.assertRaises(expected_exception) as cm:
                 self.execute_helper(session, statement)
-            if self.support_v5 and (isinstance(cm.exception, WriteFailure) or isinstance(cm.exception, ReadFailure)):
+            if ProtocolVersion.uses_error_code_map(PROTOCOL_VERSION) and (isinstance(cm.exception, WriteFailure) or isinstance(cm.exception, ReadFailure)):
                 if isinstance(cm.exception, ReadFailure):
                     self.assertEqual(list(cm.exception.error_code_map.values())[0], 1)
                 else:
