### Question:
WagedRebalancer implements StatefulRebalancer<ResourceControllerDataProvider> { @Override public Map<String, IdealState> computeNewIdealStates(ResourceControllerDataProvider clusterData, Map<String, Resource> resourceMap, final CurrentStateOutput currentStateOutput) throws HelixRebalanceException { LOG.info("Start computing new ideal states for resources: {}", resourceMap.keySet().toString()); validateInput(clusterData, resourceMap); Map<String, IdealState> newIdealStates; try { newIdealStates = computeBestPossibleStates(clusterData, resourceMap, currentStateOutput, _rebalanceAlgorithm); } catch (HelixRebalanceException ex) { LOG.error("Failed to calculate the new assignments.", ex); _rebalanceFailureCount.increment(1L); HelixRebalanceException.Type failureType = ex.getFailureType(); if (failureType.equals(HelixRebalanceException.Type.INVALID_REBALANCER_STATUS) || failureType .equals(HelixRebalanceException.Type.UNKNOWN_FAILURE)) { throw ex; } else { LOG.warn( "Returning the last known-good best possible assignment from metadata store due to " + "rebalance failure of type: {}", failureType); Map<String, ResourceAssignment> assignmentRecord = getBestPossibleAssignment(_assignmentMetadataStore, new CurrentStateOutput(), resourceMap.keySet()); newIdealStates = convertResourceAssignment(clusterData, assignmentRecord); } } newIdealStates.values().parallelStream().forEach(idealState -> { String resourceName = idealState.getResourceName(); ResourceAssignment finalAssignment = _mappingCalculator .computeBestPossiblePartitionState(clusterData, idealState, resourceMap.get(resourceName), currentStateOutput); idealState.getRecord().getMapFields().clear(); for (Partition partition : finalAssignment.getMappedPartitions()) { Map<String, String> newStateMap = finalAssignment.getReplicaMap(partition); idealState.setInstanceStateMap(partition.getPartitionName(), newStateMap == null ? Collections.emptyMap() : newStateMap); } }); LOG.info("Finish computing new ideal states for resources: {}", resourceMap.keySet().toString()); return newIdealStates; } WagedRebalancer(HelixManager helixManager); protected  WagedRebalancer(AssignmentMetadataStore assignmentMetadataStore,
      RebalanceAlgorithm algorithm, Optional<MetricCollector> metricCollectorOptional); private  WagedRebalancer(AssignmentMetadataStore assignmentMetadataStore,
      RebalanceAlgorithm algorithm, MappingCalculator mappingCalculator, HelixManager manager,
      MetricCollector metricCollector, boolean isAsyncGlobalRebalanceEnabled); void setGlobalRebalanceAsyncMode(boolean isAsyncGlobalRebalanceEnabled); synchronized void updateRebalancePreference(
      Map<ClusterConfig.GlobalRebalancePreferenceKey, Integer> newPreference); @Override void reset(); @Override void close(); @Override Map<String, IdealState> computeNewIdealStates(ResourceControllerDataProvider clusterData,
      Map<String, Resource> resourceMap, final CurrentStateOutput currentStateOutput);  }### Answer:
@Test(dependsOnMethods = "testRebalance") public void testInvalidRebalancerStatus() throws IOException { AssignmentMetadataStore metadataStore = Mockito.mock(AssignmentMetadataStore.class); when(metadataStore.getBaseline()) .thenThrow(new RuntimeException("Mock Error. Metadata store fails.")); WagedRebalancer rebalancer = new WagedRebalancer(metadataStore, _algorithm, Optional.empty()); ResourceControllerDataProvider clusterData = setupClusterDataCache(); Map<String, Resource> resourceMap = clusterData.getIdealStates().keySet().stream().collect( Collectors.toMap(resourceName -> resourceName, resourceName -> new Resource(resourceName))); try { rebalancer.computeNewIdealStates(clusterData, resourceMap, new CurrentStateOutput()); Assert.fail("Rebalance shall fail."); } catch (HelixRebalanceException ex) { Assert.assertEquals(ex.getFailureType(), HelixRebalanceException.Type.INVALID_REBALANCER_STATUS); Assert.assertEquals(ex.getMessage(), "Failed to get the current baseline assignment because of unexpected error. Failure Type: INVALID_REBALANCER_STATUS"); } }