From bbe5a0227dd54fe97c59f8674a6e486e513fc38e Mon Sep 17 00:00:00 2001
From: yifeizhuang <yifeizhuang@gmail.com>
Date: Mon, 3 Apr 2023 09:14:44 -0700
Subject: [PATCH] xds: fix flaky wrr test (#10004)

---
 .../java/io/grpc/xds/WeightedRoundRobinLoadBalancerTest.java   | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/xds/src/test/java/io/grpc/xds/WeightedRoundRobinLoadBalancerTest.java b/xds/src/test/java/io/grpc/xds/WeightedRoundRobinLoadBalancerTest.java
index eedcea935c..8ab45ef850 100644
--- a/xds/src/test/java/io/grpc/xds/WeightedRoundRobinLoadBalancerTest.java
+++ b/xds/src/test/java/io/grpc/xds/WeightedRoundRobinLoadBalancerTest.java
@@ -638,7 +638,8 @@ public class WeightedRoundRobinLoadBalancerTest {
       pickCount.put(result, pickCount.getOrDefault(result, 0) + 1);
     }
     for (int i = 0; i < capacity; i++) {
-      assertThat(Math.abs(pickCount.get(i) / 1000.0 - weights[i] / totalWeight) ).isAtMost(0.01);
+      assertThat(Math.abs(pickCount.getOrDefault(i, 0) / 1000.0 - weights[i] / totalWeight) )
+          .isAtMost(0.01);
     }
   }