Fix WhiteListPolicy

Sylvain Lebresne · Sylvain Lebresne · commit 375fbbdede7a · 2014-03-25T12:49:14.000+01:00
diff --git a/driver-core/src/main/java/com/datastax/driver/core/policies/WhiteListPolicy.java b/driver-core/src/main/java/com/datastax/driver/core/policies/WhiteListPolicy.java
@@ -15,7 +15,7 @@
  */
 package com.datastax.driver.core.policies;
 
-import java.net.InetAddress;
+import java.net.InetSocketAddress;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Iterator;
@@ -49,7 +49,7 @@
  */
 public class WhiteListPolicy implements LoadBalancingPolicy {
     private final LoadBalancingPolicy childPolicy;
-    private final Set<InetAddress> whiteList;
+    private final Set<InetSocketAddress> whiteList;
 
     /**
      * Create a new policy that wraps the provided child policy but only "allow" hosts
@@ -59,7 +59,7 @@ public class WhiteListPolicy implements LoadBalancingPolicy {
      * @param whiteList the white listed hosts. Only hosts from this list may get connected
      * to (whether they will get connected to or not depends on the child policy).
      */
-    public WhiteListPolicy(LoadBalancingPolicy childPolicy, Collection<InetAddress> whiteList) {
+    public WhiteListPolicy(LoadBalancingPolicy childPolicy, Collection<InetSocketAddress> whiteList) {
         this.childPolicy = childPolicy;
         this.whiteList = ImmutableSet.copyOf(whiteList);
     }
@@ -77,7 +77,7 @@ public WhiteListPolicy(LoadBalancingPolicy childPolicy, Collection<InetAddress>
     public void init(Cluster cluster, Collection<Host> hosts) {
         List<Host> whiteHosts = new ArrayList<Host>(hosts.size());
         for (Host host : hosts)
-            if (whiteList.contains(host.getAddress()))
+            if (whiteList.contains(host.getSocketAddress()))
                 whiteHosts.add(host);
 
         if (whiteHosts.isEmpty())
@@ -95,7 +95,7 @@ public void init(Cluster cluster, Collection<Host> hosts) {
      */
     @Override
     public HostDistance distance(Host host) {
-        return whiteList.contains(host.getAddress())
+        return whiteList.contains(host.getSocketAddress())
              ? childPolicy.distance(host)
              : HostDistance.IGNORED;
     }
@@ -121,25 +121,25 @@ public Iterator<Host> newQueryPlan(String loggedKeyspace, Statement statement) {
 
     @Override
     public void onUp(Host host) {
-        if (whiteList.contains(host.getAddress()))
+        if (whiteList.contains(host.getSocketAddress()))
             childPolicy.onUp(host);
     }
 
     @Override
     public void onDown(Host host) {
-        if (whiteList.contains(host.getAddress()))
+        if (whiteList.contains(host.getSocketAddress()))
             childPolicy.onDown(host);
     }
 
     @Override
     public void onAdd(Host host) {
-        if (whiteList.contains(host.getAddress()))
+        if (whiteList.contains(host.getSocketAddress()))
             childPolicy.onAdd(host);
     }
 
     @Override
     public void onRemove(Host host) {
-        if (whiteList.contains(host.getAddress()))
+        if (whiteList.contains(host.getSocketAddress()))
             childPolicy.onRemove(host);
     }
 }
diff --git a/driver-core/src/test/java/com/datastax/driver/core/LoadBalancingPolicyTest.java b/driver-core/src/test/java/com/datastax/driver/core/LoadBalancingPolicyTest.java
@@ -16,6 +16,7 @@
 package com.datastax.driver.core;
 
 import java.net.InetAddress;
+import java.net.InetSocketAddress;
 import java.util.*;
 
 import org.testng.annotations.Test;
@@ -74,7 +75,7 @@ public void roundRobinTest() throws Throwable {
     @Test(groups = "long")
     public void whiteListPolicyTest() throws Throwable {
 
-        List<InetAddress> whiteList = Arrays.asList(InetAddress.getByName(CCMBridge.IP_PREFIX + '2'));
+        List<InetSocketAddress> whiteList = Arrays.asList(new InetSocketAddress(InetAddress.getByName(CCMBridge.IP_PREFIX + '2'), 9042));
 
         Cluster.Builder builder = Cluster.builder().withLoadBalancingPolicy(new WhiteListPolicy(new RoundRobinPolicy(), whiteList));
         CCMBridge.CCMCluster c = CCMBridge.buildCluster(3, builder);
