From 5c09616aae5ff1620ed6bb5867717e0451c2d00f Mon Sep 17 00:00:00 2001
From: yifeizhuang <yifeizhuang@gmail.com>
Date: Mon, 3 Apr 2023 11:27:32 -0700
Subject: [PATCH] xds: fix flaky wrr test (#10005)

---
 .../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 eedcea935c8..8ab45ef8500 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 void edfScheduler() {
       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);
     }
   }