@Override
  public int getClusterCount(String clusterName, String scheme, int partitionId) throws ServiceUnavailableException
  {
    FutureCallback<Integer> clusterCountFutureCallback = new FutureCallback<>();

    _state.listenToCluster(clusterName, (type, name) ->
    {
      if (_state.getUriProperties(clusterName).getProperty() != null)
      {
        Set<URI> uris =
            _state.getUriProperties(clusterName).getProperty().getUriBySchemeAndPartition(scheme, partitionId);

        clusterCountFutureCallback.onSuccess((uris != null) ? uris.size() : 0);
      }
      else
      {
        // there won't be a UriProperties if there are no Uris announced for this scheme and/or partition. Return zero in this case.
        clusterCountFutureCallback.onSuccess(0);
      }
    });

    try
    {
      return clusterCountFutureCallback.get(_timeout, _unit);
    }
    catch (ExecutionException | TimeoutException | IllegalStateException | InterruptedException e)
    {
      if (e instanceof TimeoutException || e.getCause() instanceof TimeoutException)
      {
        int clusterCount = getClusterCountFromCache(clusterName, scheme, partitionId);
        if (clusterCount >= 0)
        {
          return clusterCount;
        }
      }
      die("ClusterInfo",
          "PEGA_1017, unable to retrieve cluster count for cluster: " + clusterName + ", scheme: " + scheme
              + ", partition: " + partitionId + ", exception: " + e);
      return -1;
    }
  }
@Test(dataProvider = "provideKeys")
  public void testClusterInfoProvider(int numHttp, int numHttps, int expectedNumHttp, int expectedNumHttps,
      int partitionIdForAdd, int partitionIdForCheck)
      throws InterruptedException, ExecutionException, ServiceUnavailableException
  {
    MockStore<ServiceProperties> serviceRegistry = new MockStore<>();
    MockStore<ClusterProperties> clusterRegistry = new MockStore<>();
    MockStore<UriProperties> uriRegistry = new MockStore<>();
    SimpleLoadBalancer loadBalancer = setupLoadBalancer(serviceRegistry, clusterRegistry, uriRegistry);

    populateUriRegistry(numHttp, numHttps, partitionIdForAdd, uriRegistry);
    clusterRegistry.put(CLUSTER1_NAME, new ClusterProperties(CLUSTER1_NAME));

    Assert.assertEquals(loadBalancer.getClusterCount(CLUSTER1_NAME, PropertyKeys.HTTP_SCHEME, partitionIdForCheck), expectedNumHttp,
        "Http cluster count for partitionId: " + partitionIdForCheck + " should be: " + expectedNumHttp);
    Assert.assertEquals(loadBalancer.getClusterCount(CLUSTER1_NAME, PropertyKeys.HTTPS_SCHEME, partitionIdForCheck), expectedNumHttps,
        "Https cluster count for partitionId: " + partitionIdForCheck + " should be: " + expectedNumHttps);
  }