Fix failing test after JAVA-1035.

Alexandre Dutra · Alexandre Dutra · commit d18546a4e50e · 2016-01-16T12:19:03.000+01:00
diff --git a/driver-core/src/test/java/com/datastax/driver/core/TokenIntegrationTest.java b/driver-core/src/test/java/com/datastax/driver/core/TokenIntegrationTest.java
@@ -284,13 +284,18 @@ private void checkRangesPerHost(String keyspace, int replicationFactor) {
     public void should_expose_tokens_per_host() {
         for (Host host : cluster.getMetadata().allHosts()) {
             assertThat(host.getTokens()).hasSize(numTokens);
-
             // Check against the info in the system tables, which is a bit weak since it's exactly how the metadata is
             // constructed in the first place, but there's not much else we can do.
             // Note that this relies on all queries going to node 1, which is why we use a WhiteList LBP in setup().
-            Row row = session.execute("select tokens from system.peers where peer = '" + host.getBroadcastAddress().getHostAddress() + "'").one();
-            if (row == null)
+            boolean isControlHost = host.getSocketAddress().equals(cluster.manager.controlConnection.connectionRef.get().address);
+            Row row;
+            if (isControlHost) {
                 row = session.execute("select tokens from system.local").one();
+            } else {
+                // non-control hosts are populated from system.peers and their broadcast address should be known
+                assertThat(host.getBroadcastAddress()).isNotNull();
+                row = session.execute("select tokens from system.peers where peer = '" + host.getBroadcastAddress().getHostAddress() + "'").one();
+            }
             Set<String> tokenStrings = row.getSet("tokens", String.class);
             assertThat(tokenStrings).hasSize(numTokens);
             Iterable<Token> tokensFromSystemTable = Iterables.transform(tokenStrings, new Function<String, Token>() {
