@Test
  public void testDiskFailureHandler() throws Exception {
    ClusterSetupHelper clusterHelper = new ClusterSetupHelper();
    clusterHelper.setupCluster();
    HelixParticipant helixParticipant = clusterHelper.helixParticipant;
    HelixAdmin helixAdmin = clusterHelper.helixAdmin;
    HelixClusterManager clusterMap = clusterHelper.clusterMap;
    DataNodeId localNode = clusterHelper.localNode;
    String clusterName = clusterHelper.clusterName;

    try {
      List<? extends ReplicaId> replicas = clusterMap.getReplicaIds(localNode);
      Map<DiskId, List<ReplicaId>> diskToReplicas = clusterHelper.getDiskToReplicaMap();
      clusterHelper.createAllMountPaths();
      StorageManager storageManager =
          new StorageManager(storeConfig, diskManagerConfig, Utils.newScheduler(1, false), metricRegistry,
              new MockIdFactory(), clusterMap, localNode, new DummyMessageStoreHardDelete(),
              Collections.singletonList(helixParticipant), SystemTime.getInstance(), new DummyMessageStoreRecovery(),
              new InMemAccountService(false, false));
      storageManager.start();
      // starting the storage manager won't start Disk failure handler right away, since there is a 10 minutes
      // delay to run the handler in a scheduler
      StorageManager.DiskFailureHandler handler = storageManager.new DiskFailureHandler();

      assertEquals(new ArrayList<>(storageManager.getDiskToDiskManager().keySet()), handler.getAllDisks());
      assertEquals("There shouldn't be any failed disk", 0, handler.getFailedDisks().size());
      clusterHelper.turnFullAutoOnForLocalhost();

      // Set every replicas to error state
      sendStateTransitionMessages(helixParticipant.getHelixManager(), "10000", replicas, "OFFLINE", "BOOTSTRAP");
      Thread.sleep(1000);

      long failureCountBefore = storageManager.getStoreMainMetrics().handleDiskFailureCount.getCount();
      // The case where there is no failed disk, running handler doesn't change anything.
      handler.run();
      assertEquals(failureCountBefore, storageManager.getStoreMainMetrics().handleDiskFailureCount.getCount());

      DiskId diskToFail = diskToReplicas.entrySet().iterator().next().getKey();
      List<ReplicaId> replicasOnFailedDisk = diskToReplicas.get(diskToFail);

      // The case where all but one replicas on the given disk are shutdown. Since the disk is not considered as failed,
      // running the handler doesn't change anything.
      for (int i = 0; i < replicasOnFailedDisk.size() - 1; i++) {
        storageManager.getStore(replicasOnFailedDisk.get(i).getPartitionId(), false).shutdown();
      }
      handler.run();
      assertEquals(failureCountBefore, storageManager.getStoreMainMetrics().handleDiskFailureCount.getCount());

      // The case where all replicas on this disk are down
      storageManager.getStore(replicasOnFailedDisk.get(replicasOnFailedDisk.size() - 1).getPartitionId(), false)
          .shutdown();
      verifyDiskFailureSuccess(storageManager, handler, helixAdmin, clusterMap, clusterName, localNode, diskToReplicas,
          diskToFail);

      // The case to run this again, since there is no new failed disk, running the handler again won't change anything.
      failureCountBefore = storageManager.getStoreMainMetrics().handleDiskFailureCount.getCount();
      long successCount = storageManager.getStoreMainMetrics().handleDiskFailureSuccessCount.getCount();
      handler.run();
      assertEquals(failureCountBefore, storageManager.getStoreMainMetrics().handleDiskFailureCount.getCount());
      assertEquals(successCount, storageManager.getStoreMainMetrics().handleDiskFailureSuccessCount.getCount());

      // The case to fail another disk
      List<DiskId> healthyDisks = new ArrayList<>(diskToReplicas.keySet());
      healthyDisks.remove(diskToFail);
      diskToFail = healthyDisks.iterator().next();
      replicasOnFailedDisk = diskToReplicas.get(diskToFail);
      for (ReplicaId replica : replicasOnFailedDisk) {
        storageManager.getStore(replica.getPartitionId(), false).shutdown();
      }
      verifyDiskFailureSuccess(storageManager, handler, helixAdmin, clusterMap, clusterName, localNode, diskToReplicas,
          diskToFail);

      // The case where a disk is empty
      // first set the disk back to available
      helixParticipant.setDisksState(Collections.singletonList(diskToFail), HardwareState.AVAILABLE);
      Thread.sleep(500);// clustermap will be updated
      // Add this disk to the disk manager in storage manager, this disk should have no replicas anymore
      storageManager.addDisk(diskToFail);
      // Create a new handler so the failed disk in this new handler would be empty
      handler = storageManager.new DiskFailureHandler();
      failureCountBefore = storageManager.getStoreMainMetrics().handleDiskFailureCount.getCount();
      handler.run();
      assertEquals(failureCountBefore, storageManager.getStoreMainMetrics().handleDiskFailureCount.getCount());

      storageManager.shutdown();
    } finally {
      clusterHelper.shutdownCluster();
    }
  }
Test disk failure handler with real helix clustermap and helix participant.
ambry-store/src/test/java/com/github/ambry/store/StorageManagerTest.java