Set route on send session instead of individual messages

Session route is "default" unless otherwise specified, which overrides
per-message routes. This means that the route parameter has not worked
as expected.

2 files changed, 4 insertions, 17 deletions

diff --git a/vespaclient-java/src/main/java/com/yahoo/vespastat/BucketStatsRetriever.java b/vespaclient-java/src/main/java/com/yahoo/vespastat/BucketStatsRetriever.java
index 1c855455a37..d42faf418a1 100644
--- a/vespaclient-java/src/main/java/com/yahoo/vespastat/BucketStatsRetriever.java
+++ b/vespaclient-java/src/main/java/com/yahoo/vespastat/BucketStatsRetriever.java
@@ -33,7 +33,6 @@ public class BucketStatsRetriever {
 
     private final MessageBusSyncSession session;
     private final MessageBusDocumentAccess documentAccess;
-    private final String route;
 
     public BucketStatsRetriever(
             DocumentAccessFactory documentAccessFactory,
@@ -42,7 +41,7 @@ public class BucketStatsRetriever {
         registerShutdownHook(registrar);
         this.documentAccess = documentAccessFactory.createDocumentAccess();
         this.session = documentAccess.createSyncSession(new SyncParameters.Builder().build());
-        this.route = route;
+        this.session.setRoute(route);
     }
 
     private void registerShutdownHook(ShutdownHookRegistrar registrar) {
@@ -102,19 +101,10 @@ public class BucketStatsRetriever {
 
 
     private <T extends Reply> T sendMessage(DocumentMessage msg, Class<T> expectedReply) throws BucketStatsException {
-        setRoute(msg, route);
         Reply reply = session.syncSend(msg);
         return validateReply(reply, expectedReply);
     }
 
-    private static void setRoute(DocumentMessage msg, String route) throws BucketStatsException {
-        try {
-            msg.setRoute(Route.parse(route));
-        } catch (Exception e) {
-            throw new BucketStatsException(String.format("Invalid route: '%s'.", route));
-        }
-    }
-
     private static <T extends Reply> T validateReply(Reply reply, Class<T> type) throws BucketStatsException {
         if (reply.hasErrors()) {
             throw new BucketStatsException(makeErrorMessage(reply));
diff --git a/vespaclient-java/src/test/java/com/yahoo/vespastat/BucketStatsRetrieverTest.java b/vespaclient-java/src/test/java/com/yahoo/vespastat/BucketStatsRetrieverTest.java
index a34e3a73205..475547546d3 100644
--- a/vespaclient-java/src/test/java/com/yahoo/vespastat/BucketStatsRetrieverTest.java
+++ b/vespaclient-java/src/test/java/com/yahoo/vespastat/BucketStatsRetrieverTest.java
@@ -20,6 +20,7 @@ import java.util.List;
 import static org.junit.Assert.assertEquals;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.argThat;
+import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -127,12 +128,8 @@ public class BucketStatsRetrieverTest {
         BucketStatsRetriever retriever = new BucketStatsRetriever(mockedFactory, route, t -> {});
         retriever.retrieveBucketList(new BucketId(0), bucketSpace);
 
-        verify(mockedSession).syncSend(argThat(new ArgumentMatcher<Message>() {
-            @Override
-            public boolean matches(Object o) {
-                return ((Message) o).getRoute().equals(Route.parse(route));
-            }
-        }));
+        // Route is set at session-level, not per message sent.
+        verify(mockedSession).setRoute(eq(route));
     }
 
     private BucketStatsRetriever createRetriever() {