@Test
  public void testRestoreDiskAvailability() throws Exception {
    ClusterSetupHelper clusterHelper = new ClusterSetupHelper();
    clusterHelper.setupCluster();
    clusterHelper.createAllMountPaths();
    clusterHelper.turnFullAutoOnForLocalhost();

    HelixClusterManager clusterMap = clusterHelper.clusterMap;
    HelixParticipant helixParticipant = clusterHelper.helixParticipant;
    HelixAdmin helixAdmin = clusterHelper.helixAdmin;
    DataNodeId localNode = clusterHelper.localNode;
    String clusterName = clusterHelper.clusterName;
    String instanceName = clusterHelper.instanceName;

    try {
      List<DiskId> allDisks = new ArrayList<>(clusterHelper.getDiskToReplicaMap().keySet());
      // 1. Start storage manager without restoring disk availability, it will not update the disk capacity
      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));
      // Start the storage manager, it will create all the replicas' directories
      storageManager.start();
      storageManager.shutdown();

      int instanceCapacity = storageManager.getCapacityOfHealthyDisks(allDisks, Collections.emptyList());
      InstanceConfig instanceConfig = helixAdmin.getInstanceConfig(clusterName, instanceName);
      assertFalse(instanceConfig.getInstanceCapacityMap().containsKey("DISK"));

      // 2. Start storage manager with restoring disk availability, it will update disk capacity
      generateConfigs(true, true, false, 2, true, true);
      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();
      storageManager.shutdown();
      instanceConfig = helixAdmin.getInstanceConfig(clusterName, instanceName);
      assertTrue(instanceConfig.getInstanceCapacityMap().containsKey("DISK"));
      assertEquals(instanceCapacity, instanceConfig.getInstanceCapacityMap().get("DISK").intValue());

      // 3. Mark a disk as unavailable and then restore the disk's availability
      DiskId badDiskId = allDisks.get(0);
      helixParticipant.setDisksState(Collections.singletonList(badDiskId), HardwareState.UNAVAILABLE);

      int instanceCapacityWithBadDisk =
          storageManager.getCapacityOfHealthyDisks(allDisks, Collections.singletonList(badDiskId));
      helixParticipant.updateDiskCapacity(instanceCapacityWithBadDisk);
      instanceConfig = helixAdmin.getInstanceConfig(clusterName, instanceName);
      assertEquals(instanceCapacityWithBadDisk, instanceConfig.getInstanceCapacityMap().get("DISK").intValue());
      // Delete mount path for bad disk so it can't be recovered
      Utils.deleteFileOrDirectory(new File(badDiskId.getMountPath()));
      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();
      storageManager.shutdown();
      instanceConfig = helixAdmin.getInstanceConfig(clusterName, instanceName);
      assertEquals(instanceCapacityWithBadDisk, instanceConfig.getInstanceCapacityMap().get("DISK").intValue());
      for (DiskId diskId : clusterHelper.getDiskToReplicaMap().keySet()) {
        if (diskId.getMountPath().equals(badDiskId.getMountPath())) {
          // this is the bad disk, it should still be available
          assertEquals(HardwareState.UNAVAILABLE, diskId.getState());
          break;
        }
      }

      // 4. Recreate the mount path, this would bring back the bad disk
      File file = new File(badDiskId.getMountPath());
      assertTrue(file.mkdirs());
      file.deleteOnExit();
      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();
      storageManager.shutdown();
      instanceConfig = helixAdmin.getInstanceConfig(clusterName, instanceName);
      assertEquals(instanceCapacity, instanceConfig.getInstanceCapacityMap().get("DISK").intValue());
      for (DiskId diskId : clusterHelper.getDiskToReplicaMap().keySet()) {
        if (diskId.getMountPath().equals(badDiskId.getMountPath())) {
          // this is the bad disk, it should still be available
          assertEquals(HardwareState.AVAILABLE, diskId.getState());
          break;
        }
      }
    } finally {
      clusterHelper.shutdownCluster();
    }
  }
Test the case when the disk is fixed and we restore the disk's availability
@throws Exception
ambry-store/src/test/java/com/github/ambry/store/StorageManagerTest.java