@Override
  @VisibleForTesting
  public void getLoadBalancedClusterAndUriProperties(String clusterName,
      Callback<Pair<ClusterProperties, UriProperties>> callback)
  {
    boolean waitForUpdatedValue = _timeout > 0;
    // if timeout is 0, we must not add the timeout callback, otherwise it would trigger immediately
    if (waitForUpdatedValue)
    {
      Callback<Pair<ClusterProperties, UriProperties>> finalCallback = callback;
      try
      {
        callback =
            new TimeoutCallback<>(_executor, _timeout, _unit, new Callback<Pair<ClusterProperties, UriProperties>>()
            {
              @Override
              public void onError(Throwable e)
              {
                if (e instanceof TimeoutException)
                {
                  handleTimeoutFromGetClusterAndUriProperties(clusterName, finalCallback);
                }
                else
                {
                  finalCallback.onError(
                      new ServiceUnavailableException(clusterName, "PEGA_1011. " + e.getMessage(), e));
                }
              }

              @Override
              public void onSuccess(Pair<ClusterProperties, UriProperties> result)
              {
                finalCallback.onSuccess(result);
              }
            }, "Timeout while fetching cluster");
      }
      catch (RejectedExecutionException e)
      {
        _log.debug("Executor rejected new tasks. It has shut down or its queue size has reached max limit");
      }
    }
    getLoadBalancedClusterAndUriProperties(clusterName, waitForUpdatedValue, callback);
  }
@Test
  @SuppressWarnings("unchecked")
  public void testGetLoadBalancedClusterAndUriProperties() throws InterruptedException, ExecutionException
  {
    MockStore<ServiceProperties> serviceRegistry = new MockStore<>();
    MockStore<ClusterProperties> clusterRegistry = new MockStore<>();
    MockStore<UriProperties> uriRegistry = new MockStore<>();
    SimpleLoadBalancerState state =
        spy(new SimpleLoadBalancerState(new SynchronousExecutorService(), uriRegistry, clusterRegistry, serviceRegistry,
                                        new HashMap<>(), new HashMap<>()));

    doAnswer(invocation ->
             {
               Thread.sleep(10);
               return null;
             }).when(state).listenToCluster(any(), any());

    SimpleLoadBalancer loadBalancer = spy(new SimpleLoadBalancer(state, 1, TimeUnit.MILLISECONDS, _d2Executor));
    FutureCallback<Pair<ClusterProperties, UriProperties>> callback = spy(new FutureCallback<>());
    // case1: listenToCluster timeout, and simpleLoadBalancer not hit the cache value
    loadBalancer.getLoadBalancedClusterAndUriProperties(CLUSTER1_NAME, callback);
    try
    {
      callback.get();
    }
    catch (ExecutionException e)
    {
      Assert.assertTrue(e.getCause() instanceof ServiceUnavailableException);
    }
    verify(loadBalancer).handleTimeoutFromGetClusterAndUriProperties(eq(CLUSTER1_NAME), eq(callback));
    verify(callback).onError(any(ServiceUnavailableException.class));

    // case2: listenToCluster timeout, and simpleLoadBalancer hit the cache value from state
    LoadBalancerStateItem<ClusterProperties> clusterItem = new LoadBalancerStateItem<>(CLUSTER_PROPERTIES, 1, 1);
    LoadBalancerStateItem<UriProperties> uriItem = new LoadBalancerStateItem<>(URI_PROPERTIES, 1, 1);
    when(state.getClusterProperties(CLUSTER1_NAME)).thenReturn(clusterItem);
    when(state.getUriProperties(CLUSTER1_NAME)).thenReturn(uriItem);
    callback = spy(new FutureCallback<>());
    loadBalancer.getLoadBalancedClusterAndUriProperties(CLUSTER1_NAME, callback);
    callback.get();
    verify(callback).onSuccess(eq(Pair.of(CLUSTER_PROPERTIES, URI_PROPERTIES)));


    // case3: getLoadBalancedClusterAndUriProperties without timeout
    state =
        spy(new SimpleLoadBalancerState(new SynchronousExecutorService(), uriRegistry, clusterRegistry, serviceRegistry,
                                        new HashMap<>(), new HashMap<>()));
    loadBalancer = spy(new SimpleLoadBalancer(state, 5, TimeUnit.SECONDS, _d2Executor));
    clusterRegistry.put(CLUSTER1_NAME, CLUSTER_PROPERTIES);
    uriRegistry.put(CLUSTER1_NAME, URI_PROPERTIES);
    callback = spy(new FutureCallback<>());
    loadBalancer.getLoadBalancedClusterAndUriProperties(CLUSTER1_NAME, callback);
    callback.get();
    verify(loadBalancer, never()).handleTimeoutFromGetClusterAndUriProperties(any(), any());
    verify(callback).onSuccess(eq(Pair.of(CLUSTER_PROPERTIES, URI_PROPERTIES)));
  }