Apply review comments

Author: lukasz-antoniak

integration-tests/src/test/java/com/datastax/dse/driver/api/core/cql/continuous/ContinuousPagingIT.java

@@ -281,7 +281,7 @@ public void prepared_statement_paging_should_be_resilient_to_schema_change() {
       // dropped.
       Row row = it.next();
       assertThat(row.getString("k")).isNotNull();
-      if (ccmRule.isDistributionAtMinimalVersion(BackendType.DSE, Version.parse("6.0.0"))) {
+      if (ccmRule.isDistributionOf(BackendType.DSE, (dist, cass) -> dist.compareTo(Version.parse("6.0.0")) >= 0)) {
         // DSE 6 only, v should be null here since dropped.
         // Not reliable for 5.1 since we may have gotten page queued before schema changed.
         assertThat(row.isNull("v")).isTrue();

integration-tests/src/test/java/com/datastax/dse/driver/api/core/graph/remote/GraphTraversalRemoteITBase.java

@@ -644,9 +644,8 @@ public void should_allow_use_of_dsl_graph_binary() {
    */
   @Test
   public void should_return_correct_results_when_bulked() {
-    Assumptions.assumeThat(
-            CcmBridge.isDistributionOf(BackendType.DSE)
-                && CcmBridge.getDistributionVersion().compareTo(Version.parse("5.1.2")) > 0)
+    Assumptions.assumeThat(CcmBridge.isDistributionOf(BackendType.DSE,
+                    (dist, cass) -> dist.compareTo(Version.parse("5.1.2")) > 0))
         .isTrue();
 
     List<String> results = graphTraversalSource().E().label().barrier().toList();

integration-tests/src/test/java/com/datastax/dse/driver/api/core/graph/statement/GraphTraversalITBase.java

@@ -599,9 +599,8 @@ public void should_allow_use_of_dsl_graph_binary() throws Exception {
    */
   @Test
   public void should_return_correct_results_when_bulked() {
-    Assumptions.assumeThat(
-            CcmBridge.isDistributionOf(BackendType.DSE)
-                && CcmBridge.getDistributionVersion().compareTo(Version.parse("5.1.2")) > 0)
+    Assumptions.assumeThat(CcmBridge.isDistributionOf(BackendType.DSE,
+            (dist, cass) -> dist.compareTo(Version.parse("5.1.2")) > 0))
         .isTrue();
 
     GraphResultSet rs =

integration-tests/src/test/java/com/datastax/oss/driver/core/metadata/SchemaChangesIT.java

@@ -55,8 +55,8 @@ public class SchemaChangesIT {
 
   static {
     CustomCcmRule.Builder builder = CustomCcmRule.builder();
-    if (!CcmBridge.isDistributionOf(BackendType.DSE)
-        && CcmBridge.getCassandraVersion().nextStable().compareTo(Version.V4_0_0) >= 0) {
+    if (!CcmBridge.isDistributionOf(BackendType.DSE,
+            (dist, cass) -> cass.nextStable().compareTo(Version.V4_0_0) >= 0)) {
       builder.withCassandraConfiguration("enable_materialized_views", true);
     }
     CCM_RULE = builder.build();

integration-tests/src/test/java/com/datastax/oss/driver/core/metadata/SchemaIT.java

@@ -335,7 +335,7 @@ public void should_exclude_virtual_keyspaces_from_token_map() {
 
   private void skipIfDse60() {
     // Special case: DSE 6.0 reports C* 4.0 but does not support virtual tables
-    if (!ccmRule.isDistributionAtMinimalVersion(BackendType.DSE, DSE_MIN_VIRTUAL_TABLES)) {
+    if (!ccmRule.isDistributionOf(BackendType.DSE, (dist, cass) -> dist.compareTo(DSE_MIN_VIRTUAL_TABLES) >= 0)) {
       throw new AssumptionViolatedException("DSE 6.0 does not support virtual tables");
     }
   }

integration-tests/src/test/java/com/datastax/oss/driver/mapper/DeleteReactiveIT.java

@@ -58,8 +58,8 @@ public class DeleteReactiveIT extends InventoryITBase {
   @ClassRule public static TestRule chain = RuleChain.outerRule(ccmRule).around(sessionRule);
 
   private static CustomCcmRule.Builder configureCcm(CustomCcmRule.Builder builder) {
-    if (!CcmBridge.isDistributionOf(BackendType.DSE)
-        && CcmBridge.getCassandraVersion().nextStable().compareTo(Version.V4_0_0) >= 0) {
+    if (!CcmBridge.isDistributionOf(BackendType.DSE,
+            (dist, cass) -> cass.nextStable().compareTo(Version.V4_0_0) >= 0)) {
       builder.withCassandraConfiguration("enable_sasi_indexes", true);
     }
     return builder;

integration-tests/src/test/java/com/datastax/oss/driver/mapper/InventoryITBase.java

@@ -89,12 +89,12 @@ protected static List<String> createStatements(BaseCcmRule ccmRule) {
     return builder.build();
   }
 
-  private static final Version MINIMUM_SASI_VERSION = Version.parse("3.4.0");
-  private static final Version BROKEN_SASI_VERSION = Version.parse("6.8.0");
+  private static final Version MINIMUM_SASI_VERSION = Objects.requireNonNull(Version.parse("3.4.0"));
+  private static final Version BROKEN_SASI_VERSION = Objects.requireNonNull(Version.parse("6.8.0"));
 
   protected static boolean isSasiBroken(BaseCcmRule ccmRule) {
     // creating SASI indexes is broken in DSE 6.8.0
-    return ccmRule.isDistributionAtVersion(BackendType.DSE, BROKEN_SASI_VERSION);
+    return ccmRule.isDistributionOf(BackendType.DSE, (dist, cass) -> dist.compareTo(BROKEN_SASI_VERSION) == 0);
   }
 
   protected static boolean supportsSASI(BaseCcmRule ccmRule) {

osgi-tests/src/test/java/com/datastax/oss/driver/internal/osgi/support/CcmStagedReactor.java

@@ -39,7 +39,7 @@ public class CcmStagedReactor extends AllConfinedStagedReactor {
 
   static {
     CcmBridge.Builder builder = CcmBridge.builder().withNodes(1);
-    if (CcmBridge.isDistributionAtMinimalVersion(BackendType.DSE, DSE_5_0)) {
+    if (CcmBridge.isDistributionOf(BackendType.DSE, (dist, cass) -> dist.compareTo(DSE_5_0) >= 0)) {
       builder.withDseWorkloads("graph");
     }
     CCM_BRIDGE = builder.build();

test-infra/src/main/java/com/datastax/oss/driver/api/testinfra/ccm/BaseCcmRule.java

@@ -80,12 +80,8 @@ public boolean isDistributionOf(BackendType type) {
     return CcmBridge.isDistributionOf(type);
   }
 
-  public boolean isDistributionAtVersion(BackendType type, Version version) {
-    return CcmBridge.isDistributionAtVersion(type, version);
-  }
-
-  public boolean isDistributionAtMinimalVersion(BackendType type, Version version) {
-    return CcmBridge.isDistributionAtMinimalVersion(type, version);
+  public boolean isDistributionOf(BackendType type, CcmBridge.VersionComparator comparator) {
+    return CcmBridge.isDistributionOf(type, comparator);
   }
 
   public Version getDistributionVersion() {

test-infra/src/main/java/com/datastax/oss/driver/api/testinfra/ccm/CcmBridge.java

@@ -180,12 +180,8 @@ public static boolean isDistributionOf(BackendType type) {
     return DISTRIBUTION == type;
   }
 
-  public static boolean isDistributionAtVersion(BackendType type, Version version) {
-    return isDistributionOf(type) && getDistributionVersion().compareTo(version) == 0;
-  }
-
-  public static boolean isDistributionAtMinimalVersion(BackendType type, Version version) {
-    return isDistributionOf(type) && getDistributionVersion().compareTo(version) >= 0;
+  public static boolean isDistributionOf(BackendType type, VersionComparator comparator) {
+    return isDistributionOf(type) && comparator.accept(getDistributionVersion(), getCassandraVersion());
   }
 
   public static Version getDistributionVersion() {
@@ -641,4 +637,8 @@ public CcmBridge build() {
           dseWorkloads);
     }
   }
+
+  public interface VersionComparator {
+    boolean accept(Version distribution, Version cassandra);
+  }
 }