@Test
  public void testSetDiskState() throws Exception {
    assumeTrue(dataNodeConfigSourceType == DataNodeConfigSourceType.PROPERTY_STORE);
    ClusterMapConfig clusterMapConfig = new ClusterMapConfig(new VerifiableProperties(props));
    String instanceName = ClusterMapUtils.getInstanceName("localhost", clusterMapConfig.clusterMapPort);
    HelixParticipant helixParticipant =
        new HelixParticipant(mock(HelixClusterManager.class), clusterMapConfig, new HelixFactory(),
            new MetricRegistry(), getDefaultZkConnectStr(clusterMapConfig), true);
    HelixAdmin helixAdmin = helixParticipant.getHelixAdmin();
    DataNodeConfig dataNodeConfig = getDataNodeConfigInHelix(helixAdmin, instanceName);
    // by default, all disks are available
    List<String> mountPaths = new ArrayList<>();
    for (Map.Entry<String, DataNodeConfig.DiskConfig> entry : dataNodeConfig.getDiskConfigs().entrySet()) {
      mountPaths.add(entry.getKey());
      assertTrue(entry.getValue().getState() == HardwareState.AVAILABLE);
    }
    AmbryDataNode dataNode = mock(AmbryDataNode.class);
    List<DiskId> disks = new ArrayList<>();
    for (String mountPath : mountPaths) {
      disks.add(new AmbryDisk(clusterMapConfig, dataNode, mountPath, HardwareState.UNAVAILABLE,
          MIN_DISK_CAPACITY_IN_BYTES * 10));
    }
    DiskId failDisk = mock(MockDiskId.class);
    DiskId missingMountDisk = new AmbryDisk(clusterMapConfig, dataNode, "/missingpath", HardwareState.AVAILABLE,
        MIN_DISK_CAPACITY_IN_BYTES * 10);

    // Test some failure cases
    try {
      helixParticipant.setDisksState(null, HardwareState.UNAVAILABLE);
      fail("Empty list should fail");
    } catch (Exception e) {
    }
    try {
      helixParticipant.setDisksState(Collections.emptyList(), HardwareState.UNAVAILABLE);
      fail("Empty list should fail");
    } catch (Exception e) {
    }
    try {
      helixParticipant.setDisksState(Collections.singletonList(failDisk), HardwareState.UNAVAILABLE);
      fail("Disk is not right");
    } catch (Exception e) {
    }
    try {
      helixParticipant.setDisksState(Collections.singletonList(missingMountDisk), HardwareState.UNAVAILABLE);
      fail("Mount path is missing");
    } catch (Exception e) {
    }

    // Nothing is changed
    assertTrue(helixParticipant.setDisksState(disks, HardwareState.AVAILABLE));
    Thread.sleep(50);
    // All the disks are changed
    assertTrue(helixParticipant.setDisksState(disks, HardwareState.UNAVAILABLE));
    Thread.sleep(50);
    dataNodeConfig = getDataNodeConfigInHelix(helixAdmin, instanceName);
    for (Map.Entry<String, DataNodeConfig.DiskConfig> entry : dataNodeConfig.getDiskConfigs().entrySet()) {
      assertTrue(entry.getValue().getState() == HardwareState.UNAVAILABLE);
    }
    // change it back to available
    assertTrue(helixParticipant.setDisksState(disks, HardwareState.AVAILABLE));
    Thread.sleep(50);
    dataNodeConfig = getDataNodeConfigInHelix(helixAdmin, instanceName);
    for (Map.Entry<String, DataNodeConfig.DiskConfig> entry : dataNodeConfig.getDiskConfigs().entrySet()) {
      assertTrue(entry.getValue().getState() == HardwareState.AVAILABLE);
    }
    helixParticipant.close();
  }
Test setting disk state in property store.
@throws Exception
ambry-clustermap/src/test/java/com/github/ambry/clustermap/HelixParticipantTest.java