Fix assertion to resolve flakiness in upstreamLocalityStatsList order (…

…grpc#11514)
ronsigal · Sep 6, 2024 · f6d2f20 · f6d2f20
1 parent 721d063
commit f6d2f20
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/xds/src/test/java/io/grpc/xds/ClusterImplLoadBalancerTest.java b/xds/src/test/java/io/grpc/xds/ClusterImplLoadBalancerTest.java
@@ -440,15 +440,15 @@ public void pickFirstLoadReport_onUpdateAddress() {
     streamTracer2.streamClosed(Status.UNAVAILABLE);
 
     clusterStats = Iterables.getOnlyElement(loadStatsManager.getClusterStatsReports(CLUSTER));
-    UpstreamLocalityStats localityStats1 = Iterables.get(clusterStats.upstreamLocalityStatsList(),
-        0);
-    assertThat(localityStats1.locality()).isEqualTo(locality1);
+    List<UpstreamLocalityStats> upstreamLocalityStatsList =
+        clusterStats.upstreamLocalityStatsList();
+    UpstreamLocalityStats localityStats1 = Iterables.find(upstreamLocalityStatsList,
+        upstreamLocalityStats -> upstreamLocalityStats.locality().equals(locality1));
     assertThat(localityStats1.totalIssuedRequests()).isEqualTo(0L);
     assertThat(localityStats1.totalSuccessfulRequests()).isEqualTo(0L);
     assertThat(localityStats1.totalErrorRequests()).isEqualTo(0L);
-    UpstreamLocalityStats localityStats2 = Iterables.get(clusterStats.upstreamLocalityStatsList(),
-        1);
-    assertThat(localityStats2.locality()).isEqualTo(locality2);
+    UpstreamLocalityStats localityStats2 = Iterables.find(upstreamLocalityStatsList,
+        upstreamLocalityStats -> upstreamLocalityStats.locality().equals(locality2));
     assertThat(localityStats2.totalIssuedRequests()).isEqualTo(1L);
     assertThat(localityStats2.totalSuccessfulRequests()).isEqualTo(0L);
     assertThat(localityStats2.totalErrorRequests()).isEqualTo(1L);