@Test
  public void testStatsManagerWithProblematicStores() throws Exception {
    DataNodeId dataNodeId = new MockDataNodeId(Collections.singletonList(new Port(6667, PortType.PLAINTEXT)),
        Collections.singletonList("/tmp"), "DC1");
    Map<PartitionId, Store> problematicStoreMap = new HashMap<>();
    PartitionId partitionId1 = new MockPartitionId(1, MockClusterMap.DEFAULT_PARTITION_CLASS,
        Collections.singletonList((MockDataNodeId) dataNodeId), 0);
    PartitionId partitionId2 = new MockPartitionId(2, MockClusterMap.DEFAULT_PARTITION_CLASS,
        Collections.singletonList((MockDataNodeId) dataNodeId), 0);
    problematicStoreMap.put(partitionId1, null);
    Store exceptionStore = new MockStore(new MockStoreStats(new HashMap<>(), true));
    problematicStoreMap.put(partitionId2, exceptionStore);
    StatsManager testStatsManager =
        new StatsManager(new MockStorageManager(problematicStoreMap, dataNodeId), mockClusterMap,
            Arrays.asList(partitionId1.getReplicaIds().get(0), partitionId2.getReplicaIds().get(0)),
            new MetricRegistry(), statsManagerConfig, new MockTime(), null, null, inMemoryAccountService, dataNodeId);
    List<PartitionId> unreachablePartitions = new ArrayList<>();
    Map<Long, Map<Short, Map<Short, ContainerStorageStats>>> hostAccountStorageStatsMap = new HashMap<>();
    for (PartitionId partitionId : problematicStoreMap.keySet()) {
      testStatsManager.collectAndAggregateAccountStorageStats(hostAccountStorageStatsMap, partitionId,
          unreachablePartitions);
    }
    assertEquals("Aggregated map should not contain any value", 0L, hostAccountStorageStatsMap.size());
    assertEquals("Unreachable store count mismatch with expected value", 2, unreachablePartitions.size());

    StatsManager.AccountStatsPublisher publisher = testStatsManager.new AccountStatsPublisher(accountStatsStore);
    publisher.run();
    HostAccountStorageStatsWrapper statsWrapper = accountStatsStore.queryHostAccountStorageStatsByHost("localhost", 0);

    List<String> unreachableStores = statsWrapper.getHeader().getUnreachableStores();
    assertTrue("The unreachable store list should contain Partition1 and Partition2",
        unreachableStores.containsAll(Arrays.asList(partitionId1.toPathString(), partitionId2.toPathString())));

    // test for the scenario where some stores are healthy and some are bad
    Map<PartitionId, Store> mixedStoreMap = new HashMap<>(storeMap);
    unreachablePartitions.clear();
    PartitionId partitionId3 = new MockPartitionId(3, MockClusterMap.DEFAULT_PARTITION_CLASS,
        Collections.singletonList((MockDataNodeId) dataNodeId), 0);
    PartitionId partitionId4 = new MockPartitionId(4, MockClusterMap.DEFAULT_PARTITION_CLASS,
        Collections.singletonList((MockDataNodeId) dataNodeId), 0);
    mixedStoreMap.put(partitionId3, null);
    mixedStoreMap.put(partitionId4, exceptionStore);
    testStatsManager = new StatsManager(new MockStorageManager(mixedStoreMap, dataNodeId), mockClusterMap,
        Arrays.asList(partitionId3.getReplicaIds().get(0), partitionId4.getReplicaIds().get(0)), new MetricRegistry(),
        statsManagerConfig, new MockTime(), null, null, inMemoryAccountService, dataNodeId);
    hostAccountStorageStatsMap.clear();
    for (PartitionId partitionId : mixedStoreMap.keySet()) {
      testStatsManager.collectAndAggregateAccountStorageStats(hostAccountStorageStatsMap, partitionId,
          unreachablePartitions);
    }
    assertEquals("Unreachable store count mismatch with expected value", 2, unreachablePartitions.size());
    // test fetchSnapshot method in StatsManager
    unreachablePartitions.clear();
    // partition 0, 1, 2 are healthy stores, partition 3, 4 are bad ones.
    for (PartitionId partitionId : mixedStoreMap.keySet()) {
      Map<Short, Map<Short, ContainerStorageStats>> containerStatsMapForPartition =
          hostAccountStorageStatsMap.get(partitionId.getId());
      if (partitionId.getId() < 3) {
        assertEquals("Actual map does not match with expected snapshot with partition id " + partitionId.toPathString(),
            hostAccountStorageStats.getStorageStats().get(partitionId.getId()), containerStatsMapForPartition);
      }
    }
  }
Test to verify the behavior when dealing with {@link Store} that is null and when {@link StoreException} is thrown.
@throws Exception
ambry-server/src/test/java/com/github/ambry/server/StatsManagerTest.java