Update test to reflect changes in cluster init (JAVA-522).

olim7t · olim7t · commit 029d932bcfdd · 2014-11-12T19:25:22.000+01:00
diff --git a/driver-core/src/test/java/com/datastax/driver/core/LoadBalancingPolicyBootstrapTest.java b/driver-core/src/test/java/com/datastax/driver/core/LoadBalancingPolicyBootstrapTest.java
@@ -17,6 +17,7 @@
 
 import java.io.PrintWriter;
 import java.io.StringWriter;
+import java.util.Collection;
 
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
@@ -41,7 +42,8 @@ private void setup() {
 
     @Test(groups = "short")
     public void notificationsTest() throws Exception {
-        assertEquals(policy.adds, 2, "adds\n" + policy.history);
+        assertEquals(policy.inits, 2, "inits\n" + policy.history);
+        assertEquals(policy.adds, 0, "adds\n" + policy.history);
         assertEquals(policy.suspecteds, 0, "suspecteds\n" + policy.history);
         assertEquals(policy.removes, 0, "removes\n" + policy.history);
         assertEquals(policy.ups, 0, "ups\n" + policy.history);
@@ -55,6 +57,7 @@ private void tearDown() {
 
     static class CountingPolicy extends DelegatingLoadBalancingPolicy {
 
+        int inits;
         int adds;
         int suspecteds;
         int removes;
@@ -68,6 +71,12 @@ public CountingPolicy(LoadBalancingPolicy delegate) {
             super(delegate);
         }
 
+        @Override
+        public void init(Cluster cluster, Collection<Host> hosts) {
+            super.init(cluster, hosts);
+            inits += hosts.size();
+        }
+
         public void onAdd(Host host) {
             out.printf("add %s%n", host);
             adds++;
