Add UTs

d2/src/test/java/com/linkedin/d2/balancer/strategies/relative/StateUpdaterTest.java

@@ -20,9 +20,11 @@
 import com.linkedin.d2.balancer.clients.TrackerClient;
 import com.linkedin.r2.util.NamedThreadFactory;
 import com.linkedin.test.util.retry.ThreeRetries;
+import com.linkedin.util.degrader.ErrorType;
 import java.net.URI;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -42,10 +44,7 @@
 
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyLong;
-import static org.testng.Assert.assertEquals;
-import static org.testng.Assert.assertFalse;
-import static org.testng.Assert.assertTrue;
-import static org.testng.Assert.fail;
+import static org.testng.Assert.*;
 
 
 /**
@@ -524,6 +523,34 @@ Object[][] getTrackerClients()
         };
   }
 
+  @Test
+  public void testStreamError()
+  {
+    Map<ErrorType, Integer> errorTypeCounts = new HashMap<>();
+    errorTypeCounts.put(ErrorType.STREAM_ERROR, 11);
+    List<TrackerClient> trackerClients = TrackerClientMockHelper.mockTrackerClients(2,
+        Arrays.asList(20, 20), Arrays.asList(10, 10), Arrays.asList(200L, 500L), Arrays.asList(100L, 200L), Arrays.asList(0, 11),
+        false, Arrays.asList(false, false), Arrays.asList(Collections.emptyMap(), errorTypeCounts));
+
+    PartitionState state = new PartitionStateTestDataBuilder()
+        .setClusterGenerationId(DEFAULT_CLUSTER_GENERATION_ID)
+        .setTrackerClientStateMap(trackerClients,
+            Arrays.asList(StateUpdater.MAX_HEALTH_SCORE, StateUpdater.MAX_HEALTH_SCORE),
+            Arrays.asList(TrackerClientState.HealthState.HEALTHY, TrackerClientState.HealthState.HEALTHY),
+            Arrays.asList(30, 30))
+        .build();
+
+    ConcurrentMap<Integer, PartitionState> partitionLoadBalancerStateMap = new ConcurrentHashMap<>();
+    partitionLoadBalancerStateMap.put(DEFAULT_PARTITION_ID, state);
+    setup(new D2RelativeStrategyProperties().setHighErrorRate(0.5), partitionLoadBalancerStateMap);
+
+    _stateUpdater.updateState();
+
+    assertEquals(_stateUpdater.getPointsMap(DEFAULT_PARTITION_ID).get(trackerClients.get(0).getUri()).intValue(), HEALTHY_POINTS);
+    assertEquals(_stateUpdater.getPointsMap(DEFAULT_PARTITION_ID).get(trackerClients.get(1).getUri()).intValue(),
+        (int) (HEALTHY_POINTS - RelativeLoadBalancerStrategyFactory.DEFAULT_DOWN_STEP * RelativeLoadBalancerStrategyFactory.DEFAULT_POINTS_PER_WEIGHT));
+  }
+
   @Test
   public void testCallCountBelowMinCallCount()
   {

d2/src/test/java/com/linkedin/d2/balancer/strategies/relative/TrackerClientMockHelper.java

@@ -27,6 +27,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
 import org.mockito.Mockito;
 
 import static org.mockito.Matchers.anyInt;
@@ -77,6 +78,18 @@ public static List<TrackerClient> mockTrackerClients(int numTrackerClients,
     return mockTrackerClients(numTrackerClients, callCountList, outstandingCallCountList, latencyList, outstandingLatencyList, errorCountList, false, doNotLoadBalance);
   }
 
+  public static List<TrackerClient> mockTrackerClients(int numTrackerClients, List<Integer> callCountList,
+      List<Integer> outstandingCallCountList, List<Long> latencyList, List<Long> outstandingLatencyList,
+      List<Integer> errorCountList, boolean doNotSlowStart, List<Boolean> doNotLoadBalance) {
+    List<Map<ErrorType, Integer>> errorTypeCountsList = errorCountList.stream().map(count -> {
+      Map<ErrorType, Integer> errorTypeCounts = new HashMap<>();
+      errorTypeCounts.put(ErrorType.SERVER_ERROR, count);
+      return errorTypeCounts;
+    }).collect(Collectors.toList());
+    return mockTrackerClients(numTrackerClients, callCountList, outstandingCallCountList, latencyList,
+        outstandingLatencyList, errorCountList, doNotSlowStart, doNotLoadBalance, errorTypeCountsList);
+  }
+
   /**
    * Mock a list of {@link TrackerClient} for testing
    *
@@ -86,11 +99,13 @@ public static List<TrackerClient> mockTrackerClients(int numTrackerClients,
    * @param latencyList The latency of each host
    * @param outstandingLatencyList The outstanding latency of each host
    * @param errorCountList The error count of each host
+   * @param errorTypeCounts The error count by type of each host
    * @return A list of mocked {@link TrackerClient}
    */
   public static List<TrackerClient> mockTrackerClients(int numTrackerClients, List<Integer> callCountList,
       List<Integer> outstandingCallCountList, List<Long> latencyList, List<Long> outstandingLatencyList,
-      List<Integer> errorCountList, boolean doNotSlowStart, List<Boolean> doNotLoadBalance)
+      List<Integer> errorCountList, boolean doNotSlowStart, List<Boolean> doNotLoadBalance,
+      List<Map<ErrorType, Integer>> errorTypeCounts)
   {
     List<TrackerClient> trackerClients = new ArrayList<>();
     for (int index = 0; index < numTrackerClients; index ++)
@@ -99,8 +114,6 @@ public static List<TrackerClient> mockTrackerClients(int numTrackerClients, List
       TrackerClient trackerClient = Mockito.mock(TrackerClient.class);
       CallTracker callTracker = Mockito.mock(CallTracker.class);
       LongStats longStats = new LongStats(callCountList.get(index), latencyList.get(index), 0, 0, 0, 0, 0, 0, 0);
-      Map<ErrorType, Integer> errorTypeCounts = new HashMap<>();
-      errorTypeCounts.put(ErrorType.SERVER_ERROR, errorCountList.get(index));
 
       CallTrackerImpl.CallTrackerStats callStats = new CallTrackerImpl.CallTrackerStats(
           RelativeLoadBalancerStrategyFactory.DEFAULT_UPDATE_INTERVAL_MS,
@@ -115,8 +128,8 @@ public static List<TrackerClient> mockTrackerClients(int numTrackerClients, List
           RelativeLoadBalancerStrategyFactory.DEFAULT_UPDATE_INTERVAL_MS - outstandingLatencyList.get(index),
       outstandingCallCountList.get(index),
       longStats,
-      errorTypeCounts,
-      errorTypeCounts);
+      errorTypeCounts.get(index),
+      errorTypeCounts.get(index));
 
       Mockito.when(trackerClient.getCallTracker()).thenReturn(callTracker);
       Mockito.when(callTracker.getCallStats()).thenReturn(callStats);

gradle.properties

@@ -1,4 +1,4 @@
-version=29.54.0
+version=29.55.0
 group=com.linkedin.pegasus
 org.gradle.configureondemand=true
 org.gradle.parallel=true