@Test
  public void testDiskFailureErrorCases() 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 {
      Map<DiskId, List<ReplicaId>> diskToReplicas = clusterHelper.getDiskToReplicaMap();

      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();

      // Above are the same initialization code from the previous test method
      long failureCountBefore = storageManager.getStoreMainMetrics().handleDiskFailureCount.getCount();
      Assert.assertEquals(0, failureCountBefore);

      // First shutdown all the replicas on the disk, so we know there will be a failure to be detected
      DiskId diskToFail = diskToReplicas.entrySet().iterator().next().getKey();
      List<ReplicaId> replicasOnFailedDisk = diskToReplicas.get(diskToFail);
      for (int i = 0; i < replicasOnFailedDisk.size(); i++) {
        storageManager.getStore(replicasOnFailedDisk.get(i).getPartitionId(), false).shutdown();
      }
      // Test case 1. Failed to enter maintenance mode
      long diskFailureErrorCountBefore = storageManager.getStoreMainMetrics().handleDiskFailureErrorCount.getCount();
      long diskFailureSuccessCountBefore =
          storageManager.getStoreMainMetrics().handleDiskFailureSuccessCount.getCount();
      helixAdmin.manuallyEnableMaintenanceMode(clusterName, true, "Block test", null);
      handler.run();
      // there should be one disk failure
      assertEquals(failureCountBefore + 1, storageManager.getStoreMainMetrics().handleDiskFailureCount.getCount());
      // but it should fail due to the failure to enter maintenance mode
      assertEquals(diskFailureErrorCountBefore + 1,
          storageManager.getStoreMainMetrics().handleDiskFailureErrorCount.getCount());
      assertEquals(diskFailureSuccessCountBefore,
          storageManager.getStoreMainMetrics().handleDiskFailureSuccessCount.getCount());
      // The diskToFail should be removed from the failedDisk list
      Assert.assertFalse(handler.getFailedDisks().contains(diskToFail));
      failureCountBefore++;
      diskFailureErrorCountBefore++;
      // Get out of maintenance mode
      helixAdmin.manuallyEnableMaintenanceMode(clusterName, false, "Block Test", null);

      // Test case 2. Not all replicas are in ERROR state
      List<ReplicaId> replicasToTransitionToError = replicasOnFailedDisk.subList(0, replicasOnFailedDisk.size() - 1);
      sendStateTransitionMessages(helixParticipant.getHelixManager(), "10000", replicasToTransitionToError, "OFFLINE",
          "BOOTSTRAP");
      Thread.sleep(1000);
      handler.run();
      assertEquals(failureCountBefore + 1, storageManager.getStoreMainMetrics().handleDiskFailureCount.getCount());
      assertEquals(diskFailureErrorCountBefore + 1,
          storageManager.getStoreMainMetrics().handleDiskFailureErrorCount.getCount());
      assertEquals(diskFailureSuccessCountBefore,
          storageManager.getStoreMainMetrics().handleDiskFailureSuccessCount.getCount());
      // The diskToFail should be removed from the failedDisk list
      Assert.assertFalse(handler.getFailedDisks().contains(diskToFail));
      failureCountBefore++;
      diskFailureErrorCountBefore++;

      // Test case 3. A success case
      // Set ever replicas to error statey
      sendStateTransitionMessages(helixParticipant.getHelixManager(), "10000",
          replicasOnFailedDisk.subList(replicasOnFailedDisk.size() - 1, replicasOnFailedDisk.size()), "OFFLINE",
          "BOOTSTRAP");
      Thread.sleep(1000);
      verifyDiskFailureSuccess(storageManager, handler, helixAdmin, clusterMap, clusterName, localNode, diskToReplicas,
          diskToFail);
      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