@Test
  public void removeInstancesTest() throws Exception {
    // create a HelixClusterManager with DynamicClusterChangeHandler
    Properties properties = new Properties();
    properties.putAll(props);
    ClusterMapConfig clusterMapConfig = new ClusterMapConfig(new VerifiableProperties(properties));
    Set<DataNodeId> removedDataNodesFromListener = new HashSet<>();
    Set<ReplicaId> removedReplicasFromListener = new HashSet<>();
    ClusterMapChangeListener listener = new ClusterMapChangeListener() {
      @Override
      public void onReplicaAddedOrRemoved(List<ReplicaId> addedReplicas, List<ReplicaId> removedReplicas) {
        removedReplicasFromListener.addAll(removedReplicas);
      }

      @Override
      public void onDataNodeRemoved(DataNodeId removedDataNode) {
        removedDataNodesFromListener.add(removedDataNode);
      }
    };
    HelixClusterManager helixClusterManager =
        new HelixClusterManager(clusterMapConfig, selfInstanceName, helixManagerFactory, new MetricRegistry());
    helixClusterManager.registerClusterMapListener(listener);

    // Pick up a nodes from hardware layout to remove
    Set<DataNode> removedDataNodes = new HashSet<>();
    for (Datacenter dc : testHardwareLayout.getHardwareLayout().getDatacenters()) {
      removedDataNodes.add(dc.getDataNodes().iterator().next());
    }

    // Get the node information from HelixClusterManager before removal
    Set<AmbryDataNode> removedAmbryDataNodes = new HashSet<>();
    Set<AmbryReplica> removedReplicas = new HashSet<>();
    for (DataNode removedNode : removedDataNodes) {
      AmbryDataNode ambryDataNode = helixClusterManager.getDataNodeId(removedNode.getHostname(), removedNode.getPort());
      assertNotNull("Removed node should exist in HelixClusterManager", ambryDataNode);
      removedAmbryDataNodes.add(ambryDataNode);
      removedReplicas.addAll(helixClusterManager.getReplicaIds(ambryDataNode));
    }

    // Invoke the onDataNodeDelete() API to remove the node.
    for (DataNode removedNode : removedDataNodes) {
      HelixClusterChangeHandler clusterChangeHandler =
          helixClusterManager.getHelixClusterChangeHandler(removedNode.getDatacenterName());
      clusterChangeHandler.onDataNodeDelete(getInstanceName(removedNode.getHostname(), removedNode.getPort()));
    }

    // Verify HelixClusterManager contains the one less node per dc.
    assertEquals("Number of data nodes after instance addition is not correct",
        testHardwareLayout.getAllExistingDataNodes().size() - removedDataNodes.size(),
        helixClusterManager.getDataNodeIds().size());
    // Verify number of partitions in cluster manager is still same
    assertEquals("Number of partitions after partition addition is not correct",
        testPartitionLayout.getPartitionCount(), helixClusterManager.getAllPartitionIds(null).size());
    // Verify capacity stats are updated
    HelixClusterManagerQueryHelper clusterManagerCallback = helixClusterManager.getManagerQueryHelper();
    // Since that we removed one node from each dc, so the raw capacity is decreased by (# of nodes) * (# of disks) * (disk capacity)
    long rawCapacityInStaticLayout = (testHardwareLayout.getAllExistingDataNodes().size() - removedDataNodes.size())
        * testHardwareLayout.getDiskCount() * testHardwareLayout.getDiskCapacityInBytes();
    assertEquals("Raw capacity of entire cluster is not expected", rawCapacityInStaticLayout,
        clusterManagerCallback.getRawCapacity());

    // Verify nodes are not present in cluster map
    for (DataNode removedDataNode : removedDataNodes) {
      AmbryDataNode ambryNode =
          helixClusterManager.getDataNodeId(removedDataNode.getHostname(), removedDataNode.getPort());
      assertNull("Node should not be present", ambryNode);
    }
    // Verify replicas and disks of removed nodes are not present in cluster map
    for (AmbryDataNode ambryNode : removedAmbryDataNodes) {
      List<AmbryReplica> ambryReplicas = helixClusterManager.getReplicaIds(ambryNode);
      assertEquals("There should be no replicas for this node", 0, ambryReplicas.size());
      Set<AmbryDisk> ambryDisks = new HashSet<>(clusterManagerCallback.getDisks(ambryNode));
      assertEquals("There should be no disks for this node", 0, ambryDisks.size());
    }
    // Verify replicas of removed nodes are not present in any partition
    for (PartitionId updatedPartition : helixClusterManager.getAllPartitionIds(DEFAULT_PARTITION_CLASS)) {
      Set<ReplicaId> replicas = new HashSet<>(updatedPartition.getReplicaIds());
      for (AmbryReplica removedReplica : removedReplicas) {
        assertFalse("Removed replicas should be no longer present in any partition", replicas.contains(removedReplica));
      }
    }
    // Verify the clustermap change listener is invoked
    assertEquals(removedAmbryDataNodes, removedDataNodesFromListener);
    assertEquals(removedReplicas, removedReplicasFromListener);

    helixClusterManager.close();
  }
Test removing instances from cluster dynamically
ambry-clustermap/src/test/java/com/github/ambry/clustermap/ClusterChangeHandlerTest.java