public synchronized Map<String, TaskCount> getTableTaskCount(String taskName) {
    String helixJobName = getHelixJobName(taskName);
    JobConfig jobConfig = _taskDriver.getJobConfig(helixJobName);
    Preconditions.checkArgument(jobConfig != null, "Task: %s does not exist", taskName);
    Map<String, TaskConfig> taskConfigMap = jobConfig.getTaskConfigMap();
    Map<String, TaskCount> taskCountMap = new HashMap<>();
    JobContext jobContext = _taskDriver.getJobContext(helixJobName);
    if (jobContext == null) {
      for (TaskConfig taskConfig : taskConfigMap.values()) {
        String tableName = taskConfig.getConfigMap().getOrDefault(MinionConstants.TABLE_NAME_KEY, UNKNOWN_TABLE_NAME);
        taskCountMap.computeIfAbsent(tableName, k -> new TaskCount()).addTaskState(null);
      }
      return taskCountMap;
    }
    Map<String, Integer> taskIdPartitionMap = jobContext.getTaskIdPartitionMap();
    for (Map.Entry<String, TaskConfig> entry : taskConfigMap.entrySet()) {
      String taskId = entry.getKey();
      TaskPartitionState state = null;
      Integer partition = taskIdPartitionMap.get(taskId);
      if (partition != null) {
        state = jobContext.getPartitionState(partition);
      }
      TaskConfig taskConfig = entry.getValue();
      String tableName = taskConfig.getConfigMap().getOrDefault(MinionConstants.TABLE_NAME_KEY, UNKNOWN_TABLE_NAME);
      taskCountMap.computeIfAbsent(tableName, k -> new TaskCount()).addTaskState(state);
    }
    return taskCountMap;
  }
@Test
  public void testGetTableTaskCount() {
    String taskName = "Task_TestTask_12345";
    String helixJobName = PinotHelixTaskResourceManager.getHelixJobName(taskName);
    TaskDriver taskDriver = mock(TaskDriver.class);
    JobConfig jobConfig = mock(JobConfig.class);
    when(taskDriver.getJobConfig(anyString())).thenReturn(jobConfig);
    Map<String, TaskConfig> taskConfigMap = new HashMap<>();
    taskConfigMap.put("taskId0", new TaskConfig("", new HashMap<>()));
    taskConfigMap.put("taskId1",
        new TaskConfig("", new HashMap<>(Collections.singletonMap("tableName", "table1_OFFLINE"))));
    when(jobConfig.getTaskConfigMap()).thenReturn(taskConfigMap);
    JobContext jobContext = mock(JobContext.class);
    when(taskDriver.getJobContext(helixJobName)).thenReturn(jobContext);
    Map<String, Integer> taskIdPartitionMap = new HashMap<>();
    taskIdPartitionMap.put("taskId0", 0);
    taskIdPartitionMap.put("taskId1", 1);
    when(jobContext.getTaskIdPartitionMap()).thenReturn(taskIdPartitionMap);
    when(jobContext.getTaskIdForPartition(0)).thenReturn("taskId0");
    when(jobContext.getTaskIdForPartition(1)).thenReturn("taskId1");
    when(jobContext.getPartitionState(0)).thenReturn(TaskPartitionState.RUNNING);
    when(jobContext.getPartitionState(1)).thenReturn(TaskPartitionState.COMPLETED);

    PinotHelixTaskResourceManager mgr =
        new PinotHelixTaskResourceManager(mock(PinotHelixResourceManager.class), taskDriver);
    Map<String, PinotHelixTaskResourceManager.TaskCount> tableTaskCount = mgr.getTableTaskCount(taskName);
    assertEquals(tableTaskCount.size(), 2);
    PinotHelixTaskResourceManager.TaskCount taskCount = tableTaskCount.get("table1_OFFLINE");
    assertEquals(taskCount.getTotal(), 1);
    assertEquals(taskCount.getCompleted(), 1);
    assertEquals(taskCount.getRunning(), 0);
    assertEquals(taskCount.getWaiting(), 0);
    assertEquals(taskCount.getError(), 0);
    assertEquals(taskCount.getUnknown(), 0);
    taskCount = tableTaskCount.get("unknown");
    assertEquals(taskCount.getTotal(), 1);
    assertEquals(taskCount.getCompleted(), 0);
    assertEquals(taskCount.getRunning(), 1);
    assertEquals(taskCount.getWaiting(), 0);
    assertEquals(taskCount.getError(), 0);
    assertEquals(taskCount.getUnknown(), 0);
  }