Java dispatch support for multiple groups of single nodes

1 files changed, 54 insertions, 4 deletions

diff --git a/container-search/src/test/java/com/yahoo/search/dispatch/LoadBalancerTest.java b/container-search/src/test/java/com/yahoo/search/dispatch/LoadBalancerTest.java
index 448a8d0e894..443ab52d8cc 100644
--- a/container-search/src/test/java/com/yahoo/search/dispatch/LoadBalancerTest.java
+++ b/container-search/src/test/java/com/yahoo/search/dispatch/LoadBalancerTest.java
@@ -4,13 +4,14 @@ package com.yahoo.search.dispatch;
 import com.yahoo.search.dispatch.SearchCluster.Group;
 import com.yahoo.search.dispatch.SearchCluster.Node;
 import junit.framework.AssertionFailedError;
-import org.hamcrest.Matchers;
 import org.junit.Test;
 
 import java.util.Arrays;
 import java.util.Optional;
 
+import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.not;
 import static org.junit.Assert.assertThat;
 
 public class LoadBalancerTest {
@@ -24,18 +25,21 @@ public class LoadBalancerTest {
         Group group = grp.orElseGet(() -> {
             throw new AssertionFailedError("Expected a SearchCluster.Group");
         });
-        assertThat(group.nodes().size(), Matchers.equalTo(1));
+        assertThat(group.nodes().size(), equalTo(1));
     }
 
     @Test
-    public void requreThatLoadBalancerIgnoresMultiGroupSetups() {
+    public void requreThatLoadBalancerServesMultiGroupSetups() {
         Node n1 = new SearchCluster.Node(0, "test-node1", 0, 0);
         Node n2 = new SearchCluster.Node(1, "test-node2", 1, 1);
         SearchCluster cluster = new SearchCluster(88.0, Arrays.asList(n1, n2), null, 1, null);
         LoadBalancer lb = new LoadBalancer(cluster);
 
         Optional<Group> grp = lb.getGroupForQuery(null);
-        assertThat(grp.isPresent(), is(false));
+        Group group = grp.orElseGet(() -> {
+            throw new AssertionFailedError("Expected a SearchCluster.Group");
+        });
+        assertThat(group.nodes().size(), equalTo(1));
     }
 
     @Test
@@ -61,4 +65,50 @@ public class LoadBalancerTest {
         Optional<Group> grp = lb.getGroupForQuery(null);
         assertThat(grp.isPresent(), is(false));
     }
+
+    @Test
+    public void requreThatLoadBalancerReturnsDifferentGroups() {
+        Node n1 = new SearchCluster.Node(0, "test-node1", 0, 0);
+        Node n2 = new SearchCluster.Node(1, "test-node2", 1, 1);
+        SearchCluster cluster = new SearchCluster(88.0, Arrays.asList(n1, n2), null, 1, null);
+        LoadBalancer lb = new LoadBalancer(cluster);
+
+        // get first group
+        Optional<Group> grp = lb.getGroupForQuery(null);
+        Group group = grp.get();
+        int id1 = group.id();
+        // release allocation
+        lb.releaseGroup(group);
+
+        // get second group
+        grp = lb.getGroupForQuery(null);
+        group = grp.get();
+        assertThat(group.id(), not(equalTo(id1)));
+    }
+
+    @Test
+    public void requreThatLoadBalancerReturnsGroupWithShortestQueue() {
+        Node n1 = new SearchCluster.Node(0, "test-node1", 0, 0);
+        Node n2 = new SearchCluster.Node(1, "test-node2", 1, 1);
+        SearchCluster cluster = new SearchCluster(88.0, Arrays.asList(n1, n2), null, 1, null);
+        LoadBalancer lb = new LoadBalancer(cluster);
+
+        // get first group
+        Optional<Group> grp = lb.getGroupForQuery(null);
+        Group group = grp.get();
+        int id1 = group.id();
+
+        // get second group
+        grp = lb.getGroupForQuery(null);
+        group = grp.get();
+        int id2 = group.id();
+        assertThat(id2, not(equalTo(id1)));
+        // release second allocation
+        lb.releaseGroup(group);
+
+        // get third group
+        grp = lb.getGroupForQuery(null);
+        group = grp.get();
+        assertThat(group.id(), equalTo(id2));
+    }
 }