@Test
public void testGetCandidateHdfsDirMap() throws Exception {

  FileSetInfo fileSetInfo1 =
      new FileSetInfo("abc", "123", 0,
          ImmutableList.copyOf(new FileSetInfo.ServingInfo[]{}), new Options());
  FileSetInfo fileSetInfo2 =
      new FileSetInfo("def", "456", 0,
          ImmutableList.copyOf(new FileSetInfo.ServingInfo[]{}), new Options());
  FileSetInfo fileSetInfo3 =
      new FileSetInfo("ghi", "789", 0,
          ImmutableList.copyOf(new FileSetInfo.ServingInfo[]{}), new Options());
  MemberModifier.field(ZooKeeperManager.class, "fileSetInfoMap")
      .set(zkManager, ImmutableMap.builder()
          .put(fileSetInfo1.fileSetName, fileSetInfo1)
          .put(fileSetInfo2.fileSetName, fileSetInfo2)
          .put(fileSetInfo3.fileSetName, fileSetInfo3)
          .build());
  when(zk.getData(eq(FILE_SET_DIR + "/" + fileSetInfo1.fileSetName),
      anyBoolean(), any(Stat.class))).thenReturn(fileSetInfo1.toJson());
  when(zk.getData(eq(LOCKS_DIR + "/" + fileSetInfo1.fileSetName),
      anyBoolean(), any(Stat.class))).thenReturn(fileSetInfo1.toJson());
  when(zk.getData(eq(FILE_SET_DIR + "/" + fileSetInfo2.fileSetName),
      anyBoolean(), any(Stat.class))).thenReturn(fileSetInfo2.toJson());
  when(zk.getData(eq(LOCKS_DIR + "/" + fileSetInfo2.fileSetName),
      anyBoolean(), any(Stat.class))).thenThrow(new KeeperException.NoNodeException());
  when(zk.getData(eq(FILE_SET_DIR + "/" + fileSetInfo3.fileSetName),
      anyBoolean(), any(Stat.class))).thenThrow(new KeeperException.NoNodeException());
  when(zk.getData(eq(LOCKS_DIR + "/" + fileSetInfo3.fileSetName),
      anyBoolean(), any(Stat.class))).thenReturn(fileSetInfo3.toJson());
  Map<String, Pair<FileSetInfo, FileSetInfo>> results = zkManager.getCandidateHdfsDirMap();
  assertEquals(3, results.size());
  assertEquals(fileSetInfo1, results.get(fileSetInfo1.fileSetName).getLeft());
  assertEquals(fileSetInfo1, results.get(fileSetInfo1.fileSetName).getRight());
  assertEquals(fileSetInfo2, results.get(fileSetInfo2.fileSetName).getLeft());
  assertNull(results.get(fileSetInfo2.fileSetName).getRight());
  assertNull(results.get(fileSetInfo3.fileSetName).getLeft());
  assertEquals(fileSetInfo3, results.get(fileSetInfo3.fileSetName).getRight());
}