BasicLoadBalancingPolicyPreferredRemoteDcsTest adjustments

Author: Bouncheck

core/src/test/java/com/datastax/oss/driver/internal/core/loadbalancing/BasicLoadBalancingPolicyPreferredRemoteDcsTest.java

@@ -53,6 +53,8 @@ public void should_prioritize_single_replica() {
     when(request.getRoutingKeyspace()).thenReturn(KEYSPACE);
     when(request.getRoutingKey()).thenReturn(ROUTING_KEY);
     when(tokenMap.getReplicas(KEYSPACE, ROUTING_KEY)).thenReturn(ImmutableSet.of(node3));
+    when(tokenMap.getReplicas(eq(KEYSPACE), any(), eq(ROUTING_KEY)))
+        .thenReturn(ImmutableSet.of(node3));
 
     // node3 always first, round-robin on the rest
     assertThat(policy.newQueryPlan(request, session))
@@ -81,6 +83,8 @@ public void should_prioritize_and_shuffle_replicas() {
     when(request.getRoutingKey()).thenReturn(ROUTING_KEY);
     when(tokenMap.getReplicas(KEYSPACE, ROUTING_KEY))
         .thenReturn(ImmutableSet.of(node1, node2, node3, node6, node9));
+    when(tokenMap.getReplicas(eq(KEYSPACE), any(), eq(ROUTING_KEY)))
+        .thenReturn(ImmutableSet.of(node1, node2, node3, node6, node9));
 
     // node 6 and 9 being in a remote DC, they don't get a boost for being a replica
     assertThat(policy.newQueryPlan(request, session))

core/src/test/java/com/datastax/oss/driver/internal/core/metadata/DefaultTopologyMonitorTest.java

@@ -502,7 +502,7 @@ public void should_get_peer_address_info_peers_v2() {
     AdminRow peer1 = mockPeersV2Row(2, node1.getHostId());
     topologyMonitor.isSchemaV2 = true;
     topologyMonitor.stubQueries(
-        new StubbedQuery("SELECT * FROM system.local", mockResult(local)),
+        new StubbedQuery("SELECT * FROM system.local WHERE key='local'", mockResult(local)),
         new StubbedQuery("SELECT * FROM system.peers_v2", mockResult(peer2, peer1)));
     when(context.getSslEngineFactory()).thenReturn(Optional.empty());
 
@@ -534,7 +534,7 @@ public void should_get_peer_address_info_peers_dse() {
     AdminRow peer1 = mockPeersRowDse(2, node1.getHostId());
     topologyMonitor.isSchemaV2 = true;
     topologyMonitor.stubQueries(
-        new StubbedQuery("SELECT * FROM system.local", mockResult(local)),
+        new StubbedQuery("SELECT * FROM system.local WHERE key='local'", mockResult(local)),
         new StubbedQuery("SELECT * FROM system.peers_v2", Maps.newHashMap(), null, true),
         new StubbedQuery("SELECT * FROM system.peers", mockResult(peer2, peer1)));
     when(context.getSslEngineFactory()).thenReturn(Optional.empty());
@@ -567,7 +567,7 @@ public void should_get_peer_address_info_peers_dse_with_ssl() {
     AdminRow peer1 = mockPeersRowDseWithSsl(2, node1.getHostId());
     topologyMonitor.isSchemaV2 = true;
     topologyMonitor.stubQueries(
-        new StubbedQuery("SELECT * FROM system.local", mockResult(local)),
+        new StubbedQuery("SELECT * FROM system.local WHERE key='local'", mockResult(local)),
         new StubbedQuery("SELECT * FROM system.peers_v2", Maps.newHashMap(), null, true),
         new StubbedQuery("SELECT * FROM system.peers", mockResult(peer2, peer1)));
     when(context.getSslEngineFactory()).thenReturn(Optional.of(sslEngineFactory));