From 1e028c404d369b7e2a1d350e5a1cd589fd61966a Mon Sep 17 00:00:00 2001
From: Terry Wilson <tmwilson@google.com>
Date: Mon, 10 Apr 2023 14:05:01 -0700
Subject: [PATCH] xds: Wait for sync context before assertions in federation
 test (#10021)

---
 .../io/grpc/xds/XdsClientFederationTest.java     | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/xds/src/test/java/io/grpc/xds/XdsClientFederationTest.java b/xds/src/test/java/io/grpc/xds/XdsClientFederationTest.java
index 225c498187..aff98a043e 100644
--- a/xds/src/test/java/io/grpc/xds/XdsClientFederationTest.java
+++ b/xds/src/test/java/io/grpc/xds/XdsClientFederationTest.java
@@ -33,6 +33,7 @@ import java.util.HashSet;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.UUID;
+import java.util.concurrent.ExecutionException;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
@@ -125,7 +126,7 @@ public class XdsClientFederationTest {
    * starts {@link LoadReportClient}s to do that.
    */
   @Test
-  public void lrsClientsStartedForLocalityStats() throws InterruptedException {
+  public void lrsClientsStartedForLocalityStats() throws InterruptedException, ExecutionException {
     trafficdirector.setLdsConfig(ControlPlaneRule.buildServerListener(),
         ControlPlaneRule.buildClientListener("test-server"));
     directpathPa.setLdsConfig(ControlPlaneRule.buildServerListener(),
@@ -144,16 +145,18 @@ public class XdsClientFederationTest {
     for (Entry<ServerInfo, LoadReportClient> entry : xdsClient.getServerLrsClientMap().entrySet()) {
       xdsClient.addClusterLocalityStats(entry.getKey(), "clusterName", "edsServiceName",
           Locality.create("", "", ""));
+      waitForSyncContext(xdsClient);
       assertThat(entry.getValue().lrsStream).isNotNull();
     }
   }
 
+
   /**
    * Assures that when an {@link XdsClient} is asked to add cluster locality stats it appropriately
    * starts {@link LoadReportClient}s to do that.
    */
   @Test
-  public void lrsClientsStartedForDropStats() throws InterruptedException {
+  public void lrsClientsStartedForDropStats() throws InterruptedException, ExecutionException {
     trafficdirector.setLdsConfig(ControlPlaneRule.buildServerListener(),
         ControlPlaneRule.buildClientListener("test-server"));
     directpathPa.setLdsConfig(ControlPlaneRule.buildServerListener(),
@@ -171,6 +174,7 @@ public class XdsClientFederationTest {
     // corresponding LRS client should be started
     for (Entry<ServerInfo, LoadReportClient> entry : xdsClient.getServerLrsClientMap().entrySet()) {
       xdsClient.addClusterDropStats(entry.getKey(), "clusterName", "edsServiceName");
+      waitForSyncContext(xdsClient);
       assertThat(entry.getValue().lrsStream).isNotNull();
     }
   }
@@ -203,6 +207,14 @@ public class XdsClientFederationTest {
     assertThat(loadStatManagers).containsNoDuplicates();
   }
 
+  // Waits for all SynchronizationContext tasks to finish, assuming no new ones get added.
+  private static void waitForSyncContext(XdsClient xdsClient)
+      throws InterruptedException, ExecutionException {
+    // We use this method just to have a task added to the sync context queue and wait for it to
+    // finish.
+    xdsClient.getSubscribedResourcesMetadataSnapshot().get();
+  }
+
   private Map<String, ?> defaultBootstrapOverride() {
     return ImmutableMap.of(
         "node", ImmutableMap.of(