@Test
public void nameResolverReturnsEmptySubLists_becomeErrorByDefault() throws Exception {
  String errorDescription = "NameResolver returned no usable address";

  // Pass a FakeNameResolverFactory with an empty list and LB config
  FakeNameResolverFactory nameResolverFactory =
      new FakeNameResolverFactory.Builder(expectedUri).build();
  Map<String, Object> rawServiceConfig =
      parseConfig("{\"loadBalancingConfig\": [ {\"mock_lb\": { \"setting1\": \"high\" } } ] }");
  ManagedChannelServiceConfig parsedServiceConfig =
      createManagedChannelServiceConfig(rawServiceConfig, null);
  nameResolverFactory.nextConfigOrError.set(ConfigOrError.fromConfig(parsedServiceConfig));
  channelBuilder.nameResolverFactory(nameResolverFactory);
  createChannel();

  // LoadBalancer received the error
  verify(mockLoadBalancerProvider).newLoadBalancer(any(Helper.class));
  verify(mockLoadBalancer).handleNameResolutionError(statusCaptor.capture());
  Status status = statusCaptor.getValue();
  assertSame(Status.Code.UNAVAILABLE, status.getCode());
  assertThat(status.getDescription()).startsWith(errorDescription);

  // A resolution retry has been scheduled
  assertEquals(1, timer.numPendingTasks(NAME_RESOLVER_REFRESH_TASK_FILTER));
}
 

@SuppressWarnings("deprecation")
@Test
public void getInstance() {
  Helper helper = mock(Helper.class);
  when(helper.getSynchronizationContext()).thenReturn(syncContext);
  when(helper.getScheduledExecutorService()).thenReturn(clock.getScheduledExecutorService());
  when(helper.getAuthority()).thenReturn("fakeauthority");
  when(helper.getChannelLogger()).thenReturn(mock(ChannelLogger.class));

  assertThat(GrpclbLoadBalancerFactory.getInstance().newLoadBalancer(helper))
      .isInstanceOf(io.grpc.grpclb.GrpclbLoadBalancer.class);

  verify(helper).getSynchronizationContext();
  verify(helper).getScheduledExecutorService();
  verify(helper).getAuthority();
  verify(helper).getChannelLogger();
  verifyNoMoreInteractions(helper);
}
 

@Test
public void getState_withRequestConnect() {
  channelBuilder.nameResolverFactory(
      new FakeNameResolverFactory.Builder(expectedUri).setResolvedAtStart(false).build());
  requestConnection = false;
  createChannel();

  assertEquals(IDLE, channel.getState(false));
  verify(mockLoadBalancerFactory, never()).newLoadBalancer(any(Helper.class));

  // call getState() with requestConnection = true
  assertEquals(IDLE, channel.getState(true));
  ArgumentCaptor<Helper> helperCaptor = ArgumentCaptor.forClass(null);
  verify(mockLoadBalancerFactory).newLoadBalancer(helperCaptor.capture());
  helper = helperCaptor.getValue();

  helper.updateBalancingState(CONNECTING, mockPicker);
  assertEquals(CONNECTING, channel.getState(false));
  assertEquals(CONNECTING, channel.getState(true));
  verifyNoMoreInteractions(mockLoadBalancerFactory);
}
 

@Test
public void idleTimeoutAndReconnect() {
  long idleTimeoutMillis = 2000L;
  channelBuilder.idleTimeout(idleTimeoutMillis, TimeUnit.MILLISECONDS);
  createChannel();

  timer.forwardNanos(TimeUnit.MILLISECONDS.toNanos(idleTimeoutMillis));
  assertEquals(IDLE, channel.getState(true /* request connection */));

  ArgumentCaptor<Helper> helperCaptor = ArgumentCaptor.forClass(Helper.class);
  // Two times of requesting connection will create loadBalancer twice.
  verify(mockLoadBalancerFactory, times(2)).newLoadBalancer(helperCaptor.capture());
  Helper helper2 = helperCaptor.getValue();

  // Updating on the old helper (whose balancer has been shutdown) does not change the channel
  // state.
  helper.updateBalancingState(CONNECTING, mockPicker);
  assertEquals(IDLE, channel.getState(false));

  helper2.updateBalancingState(CONNECTING, mockPicker);
  assertEquals(CONNECTING, channel.getState(false));
}
 

private static Subchannel createSubchannelSafely(
    final Helper helper, final EquivalentAddressGroup addressGroup, final Attributes attrs,
    final SubchannelStateListener stateListener) {
  final AtomicReference<Subchannel> resultCapture = new AtomicReference<>();
  helper.getSynchronizationContext().execute(
      new Runnable() {
        @Override
        public void run() {
          Subchannel s = helper.createSubchannel(CreateSubchannelArgs.newBuilder()
              .setAddresses(addressGroup)
              .setAttributes(attrs)
              .build());
          s.start(stateListener);
          resultCapture.set(s);
        }
      });
  return resultCapture.get();
}
 

@Test
public void updateSubchannelAddresses_newAddressConnects() {
  ClientCall<String, Integer> call = channel.newCall(method, CallOptions.DEFAULT);
  call.start(mockCallListener, new Metadata()); // Create LB
  ArgumentCaptor<Helper> helperCaptor = ArgumentCaptor.forClass(null);
  verify(mockLoadBalancerFactory).newLoadBalancer(helperCaptor.capture());
  Helper helper = helperCaptor.getValue();
  Subchannel subchannel = createSubchannelSafely(helper, servers.get(0), Attributes.EMPTY);

  subchannel.requestConnection();
  MockClientTransportInfo t0 = newTransports.poll();
  t0.listener.transportReady();

  helper.updateSubchannelAddresses(subchannel, servers.get(1));

  subchannel.requestConnection();
  MockClientTransportInfo t1 = newTransports.poll();
  t1.listener.transportReady();
}
 

@Test
public void updateSubchannelAddresses_existingAddressDoesNotConnect() {
  ClientCall<String, Integer> call = channel.newCall(method, CallOptions.DEFAULT);
  call.start(mockCallListener, new Metadata()); // Create LB
  ArgumentCaptor<Helper> helperCaptor = ArgumentCaptor.forClass(null);
  verify(mockLoadBalancerFactory).newLoadBalancer(helperCaptor.capture());
  Helper helper = helperCaptor.getValue();
  Subchannel subchannel = createSubchannelSafely(helper, servers.get(0), Attributes.EMPTY);

  subchannel.requestConnection();
  MockClientTransportInfo t0 = newTransports.poll();
  t0.listener.transportReady();

  List<SocketAddress> changedList = new ArrayList<>(servers.get(0).getAddresses());
  changedList.add(new FakeSocketAddress("aDifferentServer"));
  helper.updateSubchannelAddresses(subchannel, new EquivalentAddressGroup(changedList));

  subchannel.requestConnection();
  assertNull(newTransports.poll());
}
 

@Test
public void updateOobChannelAddresses_existingAddressDoesNotConnect() {
  ClientCall<String, Integer> call = channel.newCall(method, CallOptions.DEFAULT);
  call.start(mockCallListener, new Metadata()); // Create LB
  ArgumentCaptor<Helper> helperCaptor = ArgumentCaptor.forClass(null);
  verify(mockLoadBalancerFactory).newLoadBalancer(helperCaptor.capture());
  Helper helper = helperCaptor.getValue();
  ManagedChannel oobChannel = helper.createOobChannel(servers.get(0), "localhost");

  oobChannel.newCall(method, CallOptions.DEFAULT).start(mockCallListener, new Metadata());
  MockClientTransportInfo t0 = newTransports.poll();
  t0.listener.transportReady();

  List<SocketAddress> changedList = new ArrayList<>(servers.get(0).getAddresses());
  changedList.add(new FakeSocketAddress("aDifferentServer"));
  helper.updateOobChannelAddresses(oobChannel, new EquivalentAddressGroup(changedList));

  oobChannel.newCall(method, CallOptions.DEFAULT).start(mockCallListener, new Metadata());
  assertNull(newTransports.poll());
}
 

@Test
public void newPolicyNameTheSameAsCurrentPolicy_shouldShutdownPendingLb() {
  gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[0]));
  LoadBalancer lb0 = balancers.get(lbPolicies[0]);

  gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[0]));
  assertThat(balancers.get(lbPolicies[0])).isSameInstanceAs(lb0);

  Helper helper0 = helpers.get(lb0);
  SubchannelPicker picker = mock(SubchannelPicker.class);
  helper0.updateBalancingState(READY, picker);

  gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[1]));
  LoadBalancer lb1 = balancers.get(lbPolicies[1]);

  gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[0]));
  verify(lb1).shutdown();
  assertThat(balancers.get(lbPolicies[0])).isSameInstanceAs(lb0);

  verifyNoMoreInteractions(lb0, lb1);
}
 

@Test
public void getState_withRequestConnect() {
  channelBuilder.nameResolverFactory(
      new FakeNameResolverFactory.Builder(expectedUri).setResolvedAtStart(false).build());
  requestConnection = false;
  createChannel();

  assertEquals(IDLE, channel.getState(false));
  verify(mockLoadBalancerProvider, never()).newLoadBalancer(any(Helper.class));

  // call getState() with requestConnection = true
  assertEquals(IDLE, channel.getState(true));
  ArgumentCaptor<Helper> helperCaptor = ArgumentCaptor.forClass(null);
  verify(mockLoadBalancerProvider).newLoadBalancer(helperCaptor.capture());
  helper = helperCaptor.getValue();

  updateBalancingStateSafely(helper, CONNECTING, mockPicker);
  assertEquals(CONNECTING, channel.getState(false));
  assertEquals(CONNECTING, channel.getState(true));
  verify(mockLoadBalancerProvider).newLoadBalancer(any(Helper.class));
}
 

private Subchannel createSubchannel(final Helper helper, final int index,
    final Attributes attrs) {
  final AtomicReference<Subchannel> newSubchannel = new AtomicReference<>();
  syncContext.execute(
      new Runnable() {
        @Override
        public void run() {
          Subchannel s =
              helper.createSubchannel(
                  CreateSubchannelArgs.newBuilder()
                      .setAddresses(eagLists[index])
                      .setAttributes(attrs)
                      .build());
          s.start(mockStateListeners[index]);
          newSubchannel.set(s);
        }
      });
  return newSubchannel.get();
}
 

@Before
public void setUp() {
  doReturn(mock(ChannelLogger.class)).when(helper).getChannelLogger();
  doReturn(syncContext).when(helper).getSynchronizationContext();
  doReturn(fakeClock.getScheduledExecutorService()).when(helper).getScheduledExecutorService();
  when(orcaOobUtil.newOrcaReportingHelperWrapper(any(Helper.class),
      any(OrcaOobReportListener.class)))
      .thenAnswer(new Answer<OrcaReportingHelperWrapper>() {
        @Override
        public OrcaReportingHelperWrapper answer(InvocationOnMock invocation) {
          Helper h = invocation.getArgument(0);
          FakeOrcaReportingHelperWrapper res =
              new FakeOrcaReportingHelperWrapper(h);
          childHelperWrappers.put(h.getAuthority(), res);
          return res;
        }
      });
  lbRegistry.register(lbProvider);
  localityStore =
      new LocalityStoreImpl(logId, helper, lbRegistry, random, loadStatsStore,
          orcaPerRequestUtil, orcaOobUtil);
}
 

GrpclbState(
    Helper helper,
    SubchannelPool subchannelPool,
    TimeProvider time,
    BackoffPolicy.Provider backoffPolicyProvider) {
  this.helper = checkNotNull(helper, "helper");
  this.syncContext = checkNotNull(helper.getSynchronizationContext(), "syncContext");
  this.subchannelPool = checkNotNull(subchannelPool, "subchannelPool");
  this.time = checkNotNull(time, "time provider");
  this.timerService = checkNotNull(helper.getScheduledExecutorService(), "timerService");
  this.backoffPolicyProvider = checkNotNull(backoffPolicyProvider, "backoffPolicyProvider");
  this.serviceName = checkNotNull(helper.getAuthority(), "helper returns null authority");
  this.logger = checkNotNull(helper.getChannelLogger(), "logger");
}
 

@Override
public LoadBalancer newLoadBalancer(Helper helper) {
  HelperImpl wrappedHelper = new HelperImpl(helper);
  LoadBalancer delegateBalancer = delegateFactory.newLoadBalancer(wrappedHelper);
  wrappedHelper.init(delegateBalancer);
  return new LoadBalancerImpl(wrappedHelper, delegateBalancer);
}
 

private ChildLoadBalancerHelper(
    String target,
    Helper rlsHelper,
    SubchannelStateManager subchannelStateManager,
    SubchannelPicker picker) {
  this.target = checkNotNull(target, "target");
  this.rlsHelper = checkNotNull(rlsHelper, "rlsHelper");
  this.subchannelStateManager = checkNotNull(subchannelStateManager, "subchannelStateManager");
  this.picker = checkNotNull(picker, "picker");
}
 

AutoConfiguredLoadBalancer(Helper helper) {
  this.helper = helper;
  delegateProvider = registry.getProvider(DEFAULT_POLICY);
  if (delegateProvider == null) {
    throw new IllegalStateException("Could not find LoadBalancer " + DEFAULT_POLICY
        + ". The build probably threw away META-INF/services/io.grpc.LoadBalancerProvider");
  }
  delegate = delegateProvider.newLoadBalancer(helper);
}
 

@Override
public LoadBalancer newLoadBalancer(Helper helper) {
  edsLbHelpers.add(helper);
  LoadBalancer edsLoadBalancer = mock(LoadBalancer.class);
  edsLoadBalancers.add(edsLoadBalancer);
  return edsLoadBalancer;
}
 

@Test
public void nameResolverReturnsEmptySubLists() {
  String errorDescription = "returned an empty list";

  // Pass a FakeNameResolverFactory with an empty list
  createChannel();

  // LoadBalancer received the error
  verify(mockLoadBalancerFactory).newLoadBalancer(any(Helper.class));
  verify(mockLoadBalancer).handleNameResolutionError(statusCaptor.capture());
  Status status = statusCaptor.getValue();
  assertSame(Status.Code.UNAVAILABLE, status.getCode());
  Truth.assertThat(status.getDescription()).contains(errorDescription);
}
 

@Test
public void getState_withRequestConnect_IdleWithLbRunning() {
  channelBuilder.nameResolverFactory(
      new FakeNameResolverFactory.Builder(expectedUri).setResolvedAtStart(false).build());
  createChannel();
  verify(mockLoadBalancerFactory).newLoadBalancer(any(Helper.class));

  helper.updateBalancingState(IDLE, mockPicker);

  assertEquals(IDLE, channel.getState(true));
  verifyNoMoreInteractions(mockLoadBalancerFactory);
  verify(mockPicker).requestConnection();
}
 

private static void updateSubchannelAddressesSafely(
    final Helper helper, final Subchannel subchannel, final EquivalentAddressGroup addrs) {
  helper.getSynchronizationContext().execute(
      new Runnable() {
        @Override
        public void run() {
          subchannel.updateAddresses(Collections.singletonList(addrs));
        }
      });
}
 

private static Subchannel createSubchannelSafely(
    final Helper helper, final EquivalentAddressGroup addressGroup, final Attributes attrs) {
  final AtomicReference<Subchannel> resultCapture = new AtomicReference<Subchannel>();
  helper.getSynchronizationContext().execute(
      new Runnable() {
        @Override
        public void run() {
          resultCapture.set(helper.createSubchannel(addressGroup, attrs));
        }
      });
  return resultCapture.get();
}
 

private void createChannel(ClientInterceptor... interceptors) {
  checkState(channel == null);

  channel =
      new ManagedChannelImpl(
          channelBuilder,
          mockTransportFactory,
          new FakeBackoffPolicyProvider(),
          balancerRpcExecutorPool,
          timer.getStopwatchSupplier(),
          Arrays.asList(interceptors),
          timer.getTimeProvider());

  int numExpectedTasks = 0;

  // Force-exit the initial idle-mode
  channel.syncContext.execute(new Runnable() {
    @Override
    public void run() {
      channel.exitIdleMode();
    }
  });
  if (channelBuilder.idleTimeoutMillis != ManagedChannelImpl.IDLE_TIMEOUT_MILLIS_DISABLE) {
    numExpectedTasks += 1;
  }

  if (getNameResolverRefresh() != null) {
    numExpectedTasks += 1;
  }

  assertEquals(numExpectedTasks, timer.numPendingTasks());
  ArgumentCaptor<Helper> helperCaptor = ArgumentCaptor.forClass(null);
  verify(mockLoadBalancerProvider).newLoadBalancer(helperCaptor.capture());
}
 

private static void requestConnectionSafely(Helper helper, final Subchannel subchannel) {
  helper.getSynchronizationContext().execute(
      new Runnable() {
        @Override
        public void run() {
          subchannel.requestConnection();
        }
      });
}
 

@Test
public void channelsAndSubchannels_instrumented_state() throws Exception {
  createChannel();

  ArgumentCaptor<Helper> helperCaptor = ArgumentCaptor.forClass(null);
  verify(mockLoadBalancerProvider).newLoadBalancer(helperCaptor.capture());
  helper = helperCaptor.getValue();

  assertEquals(IDLE, getStats(channel).state);
  updateBalancingStateSafely(helper, CONNECTING, mockPicker);
  assertEquals(CONNECTING, getStats(channel).state);

  AbstractSubchannel subchannel =
      (AbstractSubchannel) createSubchannelSafely(
          helper, addressGroup, Attributes.EMPTY, subchannelStateListener);

  assertEquals(IDLE, getStats(subchannel).state);
  requestConnectionSafely(helper, subchannel);
  assertEquals(CONNECTING, getStats(subchannel).state);

  MockClientTransportInfo transportInfo = transports.poll();

  assertEquals(CONNECTING, getStats(subchannel).state);
  transportInfo.listener.transportReady();
  assertEquals(READY, getStats(subchannel).state);

  assertEquals(CONNECTING, getStats(channel).state);
  updateBalancingStateSafely(helper, READY, mockPicker);
  assertEquals(READY, getStats(channel).state);

  channel.shutdownNow();
  assertEquals(SHUTDOWN, getStats(channel).state);
  assertEquals(SHUTDOWN, getStats(subchannel).state);
}
 

@Test
public void requestConnectionForwardedToLatestPolicies() {
  gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[0]));
  LoadBalancer lb0 = balancers.get(lbPolicies[0]);
  Helper helper0 = helpers.get(lb0);
  SubchannelPicker picker = mock(SubchannelPicker.class);
  helper0.updateBalancingState(READY, picker);

  gracefulSwitchLb.requestConnection();
  verify(lb0).requestConnection();

  gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[1]));
  LoadBalancer lb1 = balancers.get(lbPolicies[1]);
  gracefulSwitchLb.requestConnection();
  verify(lb1).requestConnection();

  gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[2]));
  verify(lb1).shutdown();
  LoadBalancer lb2 = balancers.get(lbPolicies[2]);
  gracefulSwitchLb.requestConnection();
  verify(lb2).requestConnection();

  // lb2 reports READY
  helpers.get(lb2).updateBalancingState(READY, mock(SubchannelPicker.class));
  verify(lb0).shutdown();

  gracefulSwitchLb.requestConnection();
  verify(lb2, times(2)).requestConnection();

  gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[3]));
  LoadBalancer lb3 = balancers.get(lbPolicies[3]);
  gracefulSwitchLb.requestConnection();
  verify(lb3).requestConnection();

  verifyNoMoreInteractions(lb0, lb1, lb2, lb3);
}
 

@SuppressWarnings("deprecation")
@Test
public void getInstance() {
  Helper helper = mock(Helper.class);
  assertThat(
      RoundRobinLoadBalancerFactory.getInstance().newLoadBalancer(helper).getClass().getName())
      .isEqualTo("io.grpc.util.RoundRobinLoadBalancer");
  verifyZeroInteractions(helper);
}
 

private static void updateAddressesSafely(
    Helper helper, final Subchannel subchannel, final List<EquivalentAddressGroup> addrs) {
  helper.getSynchronizationContext().execute(
      new Runnable() {
        @Override
        public void run() {
          subchannel.updateAddresses(addrs);
        }
      });
}
 

GrpclbState(
    GrpclbConfig config,
    Helper helper,
    SubchannelPool subchannelPool,
    TimeProvider time,
    Stopwatch stopwatch,
    BackoffPolicy.Provider backoffPolicyProvider) {
  this.config = checkNotNull(config, "config");
  this.helper = checkNotNull(helper, "helper");
  this.syncContext = checkNotNull(helper.getSynchronizationContext(), "syncContext");
  if (config.getMode() == Mode.ROUND_ROBIN) {
    this.subchannelPool = checkNotNull(subchannelPool, "subchannelPool");
    subchannelPool.registerListener(
        new PooledSubchannelStateListener() {
          @Override
          public void onSubchannelState(
              Subchannel subchannel, ConnectivityStateInfo newState) {
            handleSubchannelState(subchannel, newState);
          }
        });
  } else {
    this.subchannelPool = null;
  }
  this.time = checkNotNull(time, "time provider");
  this.stopwatch = checkNotNull(stopwatch, "stopwatch");
  this.timerService = checkNotNull(helper.getScheduledExecutorService(), "timerService");
  this.backoffPolicyProvider = checkNotNull(backoffPolicyProvider, "backoffPolicyProvider");
  if (config.getServiceName() != null) {
    this.serviceName = config.getServiceName();
  } else {
    this.serviceName = checkNotNull(helper.getAuthority(), "helper returns null authority");
  }
  this.logger = checkNotNull(helper.getChannelLogger(), "logger");
}
 

private static void updateBalancingStateSafely(
    final Helper helper, final ConnectivityState state, final SubchannelPicker picker) {
  helper.getSynchronizationContext().execute(
      new Runnable() {
        @Override
        public void run() {
          helper.updateBalancingState(state, picker);
        }
      });
}
 

@Test
public void switchWhileOldPolicyGoesFromReadyToNotReadyWhileNewPolicyStillIdle() {
  gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[0]));
  LoadBalancer lb0 = balancers.get(lbPolicies[0]);
  InOrder inOrder = inOrder(lb0, mockHelper);
  Helper helper0 = helpers.get(lb0);
  SubchannelPicker picker = mock(SubchannelPicker.class);
  helper0.updateBalancingState(READY, picker);

  gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[1]));
  verify(lb0, never()).shutdown();

  LoadBalancer lb1 = balancers.get(lbPolicies[1]);
  Helper helper1 = helpers.get(lb1);

  picker = mock(SubchannelPicker.class);
  helper0.updateBalancingState(CONNECTING, picker);

  verify(mockHelper, never()).updateBalancingState(CONNECTING, picker);
  inOrder.verify(mockHelper).updateBalancingState(CONNECTING, BUFFER_PICKER);
  inOrder.verify(lb0).shutdown(); // shutdown after update

  picker = mock(SubchannelPicker.class);
  helper1.updateBalancingState(CONNECTING, picker);
  inOrder.verify(mockHelper).updateBalancingState(CONNECTING, picker);

  inOrder.verifyNoMoreInteractions();
  verifyNoMoreInteractions(lb1);
}