Fix load balancing flaky test (#2223)

Author: JamesNK

test/FunctionalTests/Balancer/PickFirstBalancerTests.cs

@@ -138,6 +138,9 @@ Task<HelloReply> UnaryMethod(HelloRequest request, ServerCallContext context)
         Logger.LogInformation("Ending " + endpoint.Address);
         endpoint.Dispose();
 
+        // Wait for client to change to idle state in reaction to server stopping.
+        await BalancerWaitHelpers.WaitForChannelStateAsync(Logger, channel, ConnectivityState.Idle).DefaultTimeout();
+
         Logger.LogInformation("Restarting");
         using var endpointNew = BalancerHelpers.CreateGrpcEndpoint<HelloRequest, HelloReply>(50051, UnaryMethod, nameof(UnaryMethod));
 
@@ -180,6 +183,9 @@ Task<HelloReply> UnaryMethod(HelloRequest request, ServerCallContext context)
         Logger.LogInformation("Ending " + endpoint.Address);
         endpoint.Dispose();
 
+        // Wait for client to change to idle state in reaction to server stopping.
+        await BalancerWaitHelpers.WaitForChannelStateAsync(Logger, channel, ConnectivityState.Idle).DefaultTimeout();
+
         var ex = await ExceptionAssert.ThrowsAsync<RpcException>(
             () => client.UnaryCall(new HelloRequest { Name = "Balancer" }).ResponseAsync).DefaultTimeout();