@Test
  public void testHostPartitionClassStorageStats() throws Exception {
    // First write some stats to account reports
    testMultiStoreStats();
    HostAccountStorageStatsWrapper accountStats1 = mySqlStore.queryHostAccountStorageStatsByHost(hostname1, port1);
    HostAccountStorageStatsWrapper accountStats2 = mySqlStore.queryHostAccountStorageStatsByHost(hostname2, port2);
    AccountStatsMySqlStore mySqlStore3 = createAccountStatsMySqlStore(clusterName2, hostname3, port3);
    HostAccountStorageStatsWrapper accountStats3 = mySqlStore3.queryHostAccountStorageStatsByHost(hostname3, port3);
// From this account stats, create partition class stats;
    Set<Long> allPartitionKeys = new HashSet<Long>() {
      {
        addAll(accountStats1.getStats().getStorageStats().keySet());
        addAll(accountStats2.getStats().getStorageStats().keySet());
        addAll(accountStats3.getStats().getStorageStats().keySet());
      }
    };
    List<String> partitionClassNames = Arrays.asList("default", "new");
    Map<Long, String> partitionIdToClassName = new HashMap<>();
    int ind = 0;
    for (long partitionId : allPartitionKeys) {
      partitionIdToClassName.put(partitionId, partitionClassNames.get(ind % partitionClassNames.size()));
      ind++;
    }
    HostPartitionClassStorageStatsWrapper partitionClassStats1 =
        convertHostAccountStorageStatsToHostPartitionClassStorageStats(accountStats1, partitionIdToClassName);
    HostPartitionClassStorageStatsWrapper partitionClassStats2 =
        convertHostAccountStorageStatsToHostPartitionClassStorageStats(accountStats2, partitionIdToClassName);
    HostPartitionClassStorageStatsWrapper partitionClassStats3 =
        convertHostAccountStorageStatsToHostPartitionClassStorageStats(accountStats3, partitionIdToClassName);
    mySqlStore.storeHostPartitionClassStorageStats(partitionClassStats1);
    mySqlStore.storeHostPartitionClassStorageStats(partitionClassStats2);
    mySqlStore3.storeHostPartitionClassStorageStats(partitionClassStats3);

    Map<String, Set<Integer>> partitionNameAndIds = mySqlStore.queryPartitionNameAndIds();
    assertEquals(new HashSet<>(partitionClassNames), partitionNameAndIds.keySet());
    Map<Long, String> dbPartitionKeyToClassName = partitionNameAndIds.entrySet()
        .stream()
        .flatMap(ent -> ent.getValue().stream().map(pid -> new Pair<>(ent.getKey(), (long) pid)))
        .collect(Collectors.toMap(Pair::getSecond, Pair::getFirst));
    assertEquals(partitionIdToClassName, dbPartitionKeyToClassName);

    // Fetch HostPartitionClassStorageStats
    HostPartitionClassStorageStatsWrapper obtainedStats1 =
        mySqlStore.queryHostPartitionClassStorageStatsByHost(hostname1, port1, partitionNameAndIds);
    assertEquals(partitionClassStats1.getStats().getStorageStats(), obtainedStats1.getStats().getStorageStats());
    HostPartitionClassStorageStatsWrapper obtainedStats2 =
        mySqlStore.queryHostPartitionClassStorageStatsByHost(hostname2, port2, partitionNameAndIds);
    assertEquals(partitionClassStats2.getStats().getStorageStats(), obtainedStats2.getStats().getStorageStats());
    HostPartitionClassStorageStatsWrapper obtainedStats3 =
        mySqlStore3.queryHostPartitionClassStorageStatsByHost(hostname3, port3, partitionNameAndIds);
    assertEquals(partitionClassStats3.getStats().getStorageStats(), obtainedStats3.getStats().getStorageStats());

    mySqlStore3.shutdown();
  }
Test methods to store and fetch partition class, partition name partition id and partition class storage stats.
@throws Exception
ambry-mysql/src/integration-test/java/com/github/ambry/accountstats/AccountStatsMySqlStoreIntegrationTest.java