Update IoOveruseMonitor to use UidStatsCollector.

- IoOveruseMonitor process per-UID I/O stats and detects I/O overuse.
Update this class to use UidStatsCollector instead of directly using
process or I/O stats collectors.
- Remove the logic to resolve package/service name from UID because this
is already done by UidStatsCollector.
- Address some minor nits in IoOveruseMonitor.

Test: atest libwatchdog_test
Bug: 199782126

Change-Id: I74dfc03204ebd72b3404da8b9efe3210eb146338

4 files changed, 189 insertions, 188 deletions

diff --git a/cpp/watchdog/server/src/IoOveruseMonitor.cpp b/cpp/watchdog/server/src/IoOveruseMonitor.cpp
index b97dd98bd4..23e14ee63e 100644
--- a/cpp/watchdog/server/src/IoOveruseMonitor.cpp
+++ b/cpp/watchdog/server/src/IoOveruseMonitor.cpp
@@ -170,10 +170,11 @@ void IoOveruseMonitor::terminate() {
 }
 
 Result<void> IoOveruseMonitor::onPeriodicCollection(
-        time_t time, SystemState systemState, const android::wp<UidIoStats>& uidIoStats,
-        [[maybe_unused]] const android::wp<ProcStat>& procStat,
-        [[maybe_unused]] const android::wp<ProcPidStat>& procPidStat) {
-    if (uidIoStats == nullptr) {
+        time_t time, SystemState systemState,
+        const android::wp<UidStatsCollectorInterface>& uidStatsCollector,
+        [[maybe_unused]] const android::wp<ProcStat>& procStat) {
+    android::sp<UidStatsCollectorInterface> uidStatsCollectorSp = uidStatsCollector.promote();
+    if (uidStatsCollectorSp == nullptr) {
         return Error() << "Per-UID I/O stats collector must not be null";
     }
 
@@ -191,36 +192,24 @@ Result<void> IoOveruseMonitor::onPeriodicCollection(
     mLastUserPackageIoMonitorTime = time;
     const auto [startTime, durationInSeconds] = calculateStartAndDuration(curGmt);
 
-    auto perUidIoUsage = uidIoStats.promote()->deltaStats();
-    /*
-     * TODO(b/185849350): Maybe move the packageInfo fetching logic into UidIoStats module.
-     *  This will also help avoid fetching package names in IoPerfCollection module.
-     */
-    std::vector<uid_t> seenUids;
-    for (auto it = perUidIoUsage.begin(); it != perUidIoUsage.end();) {
-        /*
-         * UidIoStats::deltaStats returns entries with zero write bytes because other metrics
-         * in these entries are non-zero.
-         */
-        if (it->second.ios.sumWriteBytes() == 0) {
-            it = perUidIoUsage.erase(it);
-            continue;
-        }
-        seenUids.push_back(it->first);
-        ++it;
-    }
-    if (perUidIoUsage.empty()) {
+    auto uidStats = uidStatsCollectorSp->deltaStats();
+    if (uidStats.empty()) {
         return {};
     }
-    const auto packageInfosByUid = mPackageInfoResolver->getPackageInfosForUids(seenUids);
     std::unordered_map<uid_t, IoOveruseStats> overusingNativeStats;
     bool isGarageModeActive = systemState == SystemState::GARAGE_MODE;
-    for (const auto& [uid, uidIoStats] : perUidIoUsage) {
-        const auto& packageInfo = packageInfosByUid.find(uid);
-        if (packageInfo == packageInfosByUid.end()) {
+    for (const auto& curUidStats : uidStats) {
+        if (curUidStats.ioStats.sumWriteBytes() == 0 || !curUidStats.hasPackageInfo()) {
+            /* 1. Ignore UIDs with zero written bytes since the last collection because they are
+             * either already accounted for or no writes made since system start.
+             *
+             * 2. UID stats without package info is not useful because the stats isn't attributed to
+             * any package/service.
+             */
             continue;
         }
-        UserPackageIoUsage curUsage(packageInfo->second, uidIoStats.ios, isGarageModeActive);
+        UserPackageIoUsage curUsage(curUidStats.packageInfo, curUidStats.ioStats,
+                                    isGarageModeActive);
         UserPackageIoUsage* dailyIoUsage;
         if (auto cachedUsage = mUserPackageDailyIoUsageById.find(curUsage.id());
             cachedUsage != mUserPackageDailyIoUsageById.end()) {
@@ -235,7 +224,7 @@ Result<void> IoOveruseMonitor::onPeriodicCollection(
         const auto threshold = mIoOveruseConfigs->fetchThreshold(dailyIoUsage->packageInfo);
 
         PackageIoOveruseStats stats;
-        stats.uid = uid;
+        stats.uid = curUidStats.packageInfo.packageIdentifier.uid;
         stats.shouldNotify = false;
         stats.ioOveruseStats.startTime = startTime;
         stats.ioOveruseStats.durationInSeconds = durationInSeconds;
@@ -272,7 +261,7 @@ Result<void> IoOveruseMonitor::onPeriodicCollection(
              */
             stats.shouldNotify = true;
             if (dailyIoUsage->packageInfo.uidType == UidType::NATIVE) {
-                overusingNativeStats[uid] = stats.ioOveruseStats;
+                overusingNativeStats[stats.uid] = stats.ioOveruseStats;
             }
             shouldSyncWatchdogService = true;
         } else if (dailyIoUsage->packageInfo.uidType != UidType::NATIVE &&
@@ -320,10 +309,10 @@ Result<void> IoOveruseMonitor::onPeriodicCollection(
 Result<void> IoOveruseMonitor::onCustomCollection(
         time_t time, SystemState systemState,
         [[maybe_unused]] const std::unordered_set<std::string>& filterPackages,
-        const android::wp<UidIoStats>& uidIoStats, const android::wp<ProcStat>& procStat,
-        const android::wp<ProcPidStat>& procPidStat) {
+        const android::wp<UidStatsCollectorInterface>& uidStatsCollector,
+        const android::wp<ProcStat>& procStat) {
     // Nothing special for custom collection.
-    return onPeriodicCollection(time, systemState, uidIoStats, procStat, procPidStat);
+    return onPeriodicCollection(time, systemState, uidStatsCollector, procStat);
 }
 
 Result<void> IoOveruseMonitor::onPeriodicMonitor(
@@ -380,13 +369,13 @@ Result<void> IoOveruseMonitor::onPeriodicMonitor(
     return {};
 }
 
-Result<void> IoOveruseMonitor::onDump([[maybe_unused]] int fd) {
+Result<void> IoOveruseMonitor::onDump([[maybe_unused]] int fd) const {
     // TODO(b/183436216): Dump the list of killed/disabled packages. Dump the list of packages that
     //  exceed xx% of their threshold.
     return {};
 }
 
-bool IoOveruseMonitor::dumpHelpText(int fd) {
+bool IoOveruseMonitor::dumpHelpText(int fd) const {
     return WriteStringToFd(StringPrintf(kHelpText, name().c_str(), kResetResourceOveruseStatsFlag),
                            fd);
 }
@@ -438,7 +427,7 @@ Result<void> IoOveruseMonitor::updateResourceOveruseConfigurations(
 }
 
 Result<void> IoOveruseMonitor::getResourceOveruseConfigurations(
-        std::vector<ResourceOveruseConfiguration>* configs) {
+        std::vector<ResourceOveruseConfiguration>* configs) const {
     std::shared_lock readLock(mRwMutex);
     if (!isInitializedLocked()) {
         return Error(Status::EX_ILLEGAL_STATE) << name() << " is not initialized";
@@ -496,7 +485,7 @@ Result<void> IoOveruseMonitor::removeIoOveruseListener(
     return {};
 }
 
-Result<void> IoOveruseMonitor::getIoOveruseStats(IoOveruseStats* ioOveruseStats) {
+Result<void> IoOveruseMonitor::getIoOveruseStats(IoOveruseStats* ioOveruseStats) const {
     if (!isInitialized()) {
         return Error(Status::EX_ILLEGAL_STATE) << "I/O overuse monitor is not initialized";
     }
@@ -580,14 +569,14 @@ bool IoOveruseMonitor::findListenerAndProcessLocked(const sp<IBinder>& binder,
 }
 
 IoOveruseMonitor::UserPackageIoUsage::UserPackageIoUsage(const PackageInfo& pkgInfo,
-                                                         const IoUsage& ioUsage,
+                                                         const UidIoStats& uidIoStats,
                                                          const bool isGarageModeActive) {
     packageInfo = pkgInfo;
     if (isGarageModeActive) {
-        writtenBytes.garageModeBytes = ioUsage.sumWriteBytes();
+        writtenBytes.garageModeBytes = uidIoStats.sumWriteBytes();
     } else {
-        writtenBytes.foregroundBytes = ioUsage.metrics[WRITE_BYTES][FOREGROUND];
-        writtenBytes.backgroundBytes = ioUsage.metrics[WRITE_BYTES][BACKGROUND];
+        writtenBytes.foregroundBytes = uidIoStats.metrics[WRITE_BYTES][FOREGROUND];
+        writtenBytes.backgroundBytes = uidIoStats.metrics[WRITE_BYTES][BACKGROUND];
     }
 }
 
diff --git a/cpp/watchdog/server/src/IoOveruseMonitor.h b/cpp/watchdog/server/src/IoOveruseMonitor.h
index 27e179fb6a..df7ff609a3 100644
--- a/cpp/watchdog/server/src/IoOveruseMonitor.h
+++ b/cpp/watchdog/server/src/IoOveruseMonitor.h
@@ -19,9 +19,8 @@
 
 #include "IoOveruseConfigs.h"
 #include "PackageInfoResolver.h"
-#include "ProcPidStat.h"
 #include "ProcStat.h"
-#include "UidIoStats.h"
+#include "UidStatsCollector.h"
 #include "WatchdogPerfService.h"
 
 #include <android-base/result.h>
@@ -62,17 +61,17 @@ class IoOveruseMonitorPeer;
 // Used only in tests.
 std::tuple<int64_t, int64_t> calculateStartAndDuration(const time_t& currentTime);
 
-/*
+/**
  * IIoOveruseMonitor interface defines the methods that the I/O overuse monitoring module
  * should implement.
  */
 class IIoOveruseMonitor : virtual public IDataProcessorInterface {
 public:
     // Returns whether or not the monitor is initialized.
-    virtual bool isInitialized() = 0;
+    virtual bool isInitialized() const = 0;
 
     // Dumps the help text.
-    virtual bool dumpHelpText(int fd) = 0;
+    virtual bool dumpHelpText(int fd) const = 0;
 
     // Below API is from internal/ICarWatchdog.aidl. Please refer to the AIDL for description.
     virtual android::base::Result<void> updateResourceOveruseConfigurations(
@@ -81,7 +80,7 @@ public:
                     configs) = 0;
     virtual android::base::Result<void> getResourceOveruseConfigurations(
             std::vector<android::automotive::watchdog::internal::ResourceOveruseConfiguration>*
-                    configs) = 0;
+                    configs) const = 0;
     virtual android::base::Result<void> actionTakenOnIoOveruse(
             const std::vector<
                     android::automotive::watchdog::internal::PackageResourceOveruseAction>&
@@ -94,7 +93,7 @@ public:
     virtual android::base::Result<void> removeIoOveruseListener(
             const sp<IResourceOveruseListener>& listener) = 0;
 
-    virtual android::base::Result<void> getIoOveruseStats(IoOveruseStats* ioOveruseStats) = 0;
+    virtual android::base::Result<void> getIoOveruseStats(IoOveruseStats* ioOveruseStats) const = 0;
 
     virtual android::base::Result<void> resetIoOveruseStats(
             const std::vector<std::string>& packageNames) = 0;
@@ -106,42 +105,42 @@ public:
 
     ~IoOveruseMonitor() { terminate(); }
 
-    bool isInitialized() {
+    bool isInitialized() const override {
         std::shared_lock readLock(mRwMutex);
         return isInitializedLocked();
     }
 
     // Below methods implement IDataProcessorInterface.
-    std::string name() { return "IoOveruseMonitor"; }
+    std::string name() const override { return "IoOveruseMonitor"; }
     friend std::ostream& operator<<(std::ostream& os, const IoOveruseMonitor& monitor);
     android::base::Result<void> onBoottimeCollection(
-            time_t /*time*/, const android::wp<UidIoStats>& /*uidIoStats*/,
-            const android::wp<ProcStat>& /*procStat*/,
-            const android::wp<ProcPidStat>& /*procPidStat*/) {
+            [[maybe_unused]] time_t time,
+            [[maybe_unused]] const android::wp<UidStatsCollectorInterface>& uidStatsCollector,
+            [[maybe_unused]] const android::wp<ProcStat>& procStat) override {
         // No I/O overuse monitoring during boot-time.
         return {};
     }
 
-    android::base::Result<void> onPeriodicCollection(time_t time, SystemState systemState,
-                                                     const android::wp<UidIoStats>& uidIoStats,
-                                                     const android::wp<ProcStat>& procStat,
-                                                     const android::wp<ProcPidStat>& procPidStat);
+    android::base::Result<void> onPeriodicCollection(
+            time_t time, SystemState systemState,
+            const android::wp<UidStatsCollectorInterface>& uidStatsCollector,
+            const android::wp<ProcStat>& procStat) override;
 
     android::base::Result<void> onCustomCollection(
             time_t time, SystemState systemState,
             const std::unordered_set<std::string>& filterPackages,
-            const android::wp<UidIoStats>& uidIoStats, const android::wp<ProcStat>& procStat,
-            const android::wp<ProcPidStat>& procPidStat);
+            const android::wp<UidStatsCollectorInterface>& uidStatsCollector,
+            const android::wp<ProcStat>& procStat) override;
 
     android::base::Result<void> onPeriodicMonitor(
             time_t time, const android::wp<IProcDiskStatsInterface>& procDiskStats,
-            const std::function<void()>& alertHandler);
+            const std::function<void()>& alertHandler) override;
 
-    android::base::Result<void> onDump(int fd);
+    android::base::Result<void> onDump(int fd) const override;
 
-    bool dumpHelpText(int fd);
+    bool dumpHelpText(int fd) const override;
 
-    android::base::Result<void> onCustomCollectionDump(int /*fd*/) {
+    android::base::Result<void> onCustomCollectionDump([[maybe_unused]] int fd) override {
         // No special processing for custom collection. Thus no custom collection dump.
         return {};
     }
@@ -149,26 +148,28 @@ public:
     // Below methods implement AIDL interfaces.
     android::base::Result<void> updateResourceOveruseConfigurations(
             const std::vector<
-                    android::automotive::watchdog::internal::ResourceOveruseConfiguration>&
-                    configs);
+                    android::automotive::watchdog::internal::ResourceOveruseConfiguration>& configs)
+            override;
 
     android::base::Result<void> getResourceOveruseConfigurations(
             std::vector<android::automotive::watchdog::internal::ResourceOveruseConfiguration>*
-                    configs);
+                    configs) const override;
 
     android::base::Result<void> actionTakenOnIoOveruse(
             const std::vector<
-                    android::automotive::watchdog::internal::PackageResourceOveruseAction>&
-                    actions);
+                    android::automotive::watchdog::internal::PackageResourceOveruseAction>& actions)
+            override;
 
-    android::base::Result<void> addIoOveruseListener(const sp<IResourceOveruseListener>& listener);
+    android::base::Result<void> addIoOveruseListener(
+            const sp<IResourceOveruseListener>& listener) override;
 
     android::base::Result<void> removeIoOveruseListener(
-            const sp<IResourceOveruseListener>& listener);
+            const sp<IResourceOveruseListener>& listener) override;
 
-    android::base::Result<void> getIoOveruseStats(IoOveruseStats* ioOveruseStats);
+    android::base::Result<void> getIoOveruseStats(IoOveruseStats* ioOveruseStats) const override;
 
-    android::base::Result<void> resetIoOveruseStats(const std::vector<std::string>& packageName);
+    android::base::Result<void> resetIoOveruseStats(
+            const std::vector<std::string>& packageName) override;
 
 protected:
     android::base::Result<void> init();
@@ -183,7 +184,7 @@ private:
 
     struct UserPackageIoUsage {
         UserPackageIoUsage(const android::automotive::watchdog::internal::PackageInfo& packageInfo,
-                           const IoUsage& IoUsage, const bool isGarageModeActive);
+                           const UidIoStats& uidIoStats, const bool isGarageModeActive);
         android::automotive::watchdog::internal::PackageInfo packageInfo = {};
         PerStateBytes writtenBytes = {};
         PerStateBytes forgivenWriteBytes = {};
@@ -211,7 +212,7 @@ private:
     };
 
 private:
-    bool isInitializedLocked() { return mIoOveruseConfigs != nullptr; }
+    bool isInitializedLocked() const { return mIoOveruseConfigs != nullptr; }
 
     void notifyNativePackagesLocked(const std::unordered_map<uid_t, IoOveruseStats>& statsByUid);
 
@@ -221,7 +222,7 @@ private:
     using Processor = std::function<void(ListenersByUidMap&, ListenersByUidMap::const_iterator)>;
     bool findListenerAndProcessLocked(const sp<IBinder>& binder, const Processor& processor);
 
-    /*
+    /**
      * Writes in-memory configs to disk asynchronously if configs are not written after latest
      * update.
      */
@@ -239,7 +240,7 @@ private:
     // Summary of configs available for all the components and system-wide overuse alert thresholds.
     sp<IIoOveruseConfigs> mIoOveruseConfigs GUARDED_BY(mRwMutex);
 
-    /*
+    /**
      * Delta of system-wide written kib across all disks from the last |mPeriodicMonitorBufferSize|
      * polls along with the polling duration.
      */
diff --git a/cpp/watchdog/server/tests/IoOveruseMonitorTest.cpp b/cpp/watchdog/server/tests/IoOveruseMonitorTest.cpp
index 7dcc2d4a4c..76f3cbbeaf 100644
--- a/cpp/watchdog/server/tests/IoOveruseMonitorTest.cpp
+++ b/cpp/watchdog/server/tests/IoOveruseMonitorTest.cpp
@@ -19,14 +19,17 @@
 #include "MockPackageInfoResolver.h"
 #include "MockProcDiskStats.h"
 #include "MockResourceOveruseListener.h"
-#include "MockUidIoStats.h"
+#include "MockUidStatsCollector.h"
 #include "MockWatchdogServiceHelper.h"
+#include "PackageInfoTestUtils.h"
 
 #include <binder/IPCThreadState.h>
 #include <binder/Status.h>
 #include <utils/RefBase.h>
 
 #include <functional>
+#include <tuple>
+#include <unordered_map>
 
 namespace android {
 namespace automotive {
@@ -47,6 +50,7 @@ using ::android::base::StringAppendF;
 using ::android::binder::Status;
 using ::testing::_;
 using ::testing::DoAll;
+using ::testing::Eq;
 using ::testing::Return;
 using ::testing::ReturnRef;
 using ::testing::SaveArg;
@@ -68,20 +72,11 @@ IoOveruseAlertThreshold toIoOveruseAlertThreshold(const int64_t durationInSecond
     return threshold;
 }
 
-PackageIdentifier constructPackageIdentifier(const char* packageName, const int32_t uid) {
-    PackageIdentifier packageIdentifier;
-    packageIdentifier.name = packageName;
-    packageIdentifier.uid = uid;
-    return packageIdentifier;
-}
-
-PackageInfo constructPackageInfo(const char* packageName, const int32_t uid,
-                                 const UidType uidType) {
+struct PackageWrittenBytes {
     PackageInfo packageInfo;
-    packageInfo.packageIdentifier = constructPackageIdentifier(packageName, uid);
-    packageInfo.uidType = uidType;
-    return packageInfo;
-}
+    int32_t foregroundBytes;
+    int32_t backgroundBytes;
+};
 
 PerStateBytes constructPerStateBytes(const int64_t fgBytes, const int64_t bgBytes,
                                      const int64_t gmBytes) {
@@ -196,42 +191,25 @@ protected:
         mMockWatchdogServiceHelper = sp<MockWatchdogServiceHelper>::make();
         mMockIoOveruseConfigs = sp<MockIoOveruseConfigs>::make();
         mMockPackageInfoResolver = sp<MockPackageInfoResolver>::make();
+        mMockUidStatsCollector = sp<MockUidStatsCollector>::make();
         mIoOveruseMonitor = sp<IoOveruseMonitor>::make(mMockWatchdogServiceHelper);
         mIoOveruseMonitorPeer = sp<internal::IoOveruseMonitorPeer>::make(mIoOveruseMonitor);
         mIoOveruseMonitorPeer->init(mMockIoOveruseConfigs, mMockPackageInfoResolver);
+        setUpPackagesAndConfigurations();
     }
 
     virtual void TearDown() {
         mMockWatchdogServiceHelper.clear();
         mMockIoOveruseConfigs.clear();
         mMockPackageInfoResolver.clear();
+        mMockUidStatsCollector.clear();
         mIoOveruseMonitor.clear();
         mIoOveruseMonitorPeer.clear();
     }
 
     void setUpPackagesAndConfigurations() {
-        std::unordered_map<uid_t, PackageInfo> packageInfoMapping =
-                {{1001000,
-                  constructPackageInfo(
-                          /*packageName=*/"system.daemon", /*uid=*/1001000, UidType::NATIVE)},
-                 {1112345,
-                  constructPackageInfo(
-                          /*packageName=*/"com.android.google.package", /*uid=*/1112345,
-                          UidType::APPLICATION)},
-                 {1113999,
-                  constructPackageInfo(
-                          /*packageName=*/"com.android.google.package", /*uid=*/1113999,
-                          UidType::APPLICATION)},
-                 {1212345,
-                  constructPackageInfo(
-                          /*packageName=*/"com.android.google.package", /*uid=*/1212345,
-                          UidType::APPLICATION)},
-                 {1312345,
-                  constructPackageInfo(
-                          /*packageName=*/"com.android.google.package", /*uid=*/1312345,
-                          UidType::APPLICATION)}};
         ON_CALL(*mMockPackageInfoResolver, getPackageInfosForUids(_))
-                .WillByDefault(Return(packageInfoMapping));
+                .WillByDefault(Return(kPackageInfosByUid));
         mMockIoOveruseConfigs->injectPackageConfigs({
                 {"system.daemon",
                  {constructPerStateBytes(/*fgBytes=*/80'000, /*bgBytes=*/40'000,
@@ -244,6 +222,24 @@ protected:
         });
     }
 
+    std::vector<UidStats> constructUidStats(
+            std::unordered_map<uid_t, std::tuple<int32_t, int32_t>> writtenBytesByUid) {
+        std::vector<UidStats> uidStats;
+        for (const auto& [uid, writtenBytes] : writtenBytesByUid) {
+            PackageInfo packageInfo;
+            if (kPackageInfosByUid.find(uid) != kPackageInfosByUid.end()) {
+                packageInfo = kPackageInfosByUid.at(uid);
+            }
+            uidStats.push_back(UidStats{.packageInfo = packageInfo,
+                                        .ioStats = {/*fgRdBytes=*/989'000,
+                                                    /*bgRdBytes=*/678'000,
+                                                    /*fgWrBytes=*/std::get<0>(writtenBytes),
+                                                    /*bgWrBytes=*/std::get<1>(writtenBytes),
+                                                    /*fgFsync=*/10'000, /*bgFsync=*/50'000}});
+        }
+        return uidStats;
+    }
+
     void executeAsUid(uid_t uid, std::function<void()> func) {
         sp<ScopedChangeCallingUid> scopedChangeCallingUid = sp<ScopedChangeCallingUid>::make(uid);
         ASSERT_NO_FATAL_FAILURE(func());
@@ -252,12 +248,36 @@ protected:
     sp<MockWatchdogServiceHelper> mMockWatchdogServiceHelper;
     sp<MockIoOveruseConfigs> mMockIoOveruseConfigs;
     sp<MockPackageInfoResolver> mMockPackageInfoResolver;
+    sp<MockUidStatsCollector> mMockUidStatsCollector;
     sp<IoOveruseMonitor> mIoOveruseMonitor;
     sp<internal::IoOveruseMonitorPeer> mIoOveruseMonitorPeer;
+
+    static const std::unordered_map<uid_t, PackageInfo> kPackageInfosByUid;
 };
 
+const std::unordered_map<uid_t, PackageInfo> IoOveruseMonitorTest::kPackageInfosByUid =
+        {{1001000,
+          constructPackageInfo(
+                  /*packageName=*/"system.daemon",
+                  /*uid=*/1001000, UidType::NATIVE)},
+         {1112345,
+          constructPackageInfo(
+                  /*packageName=*/"com.android.google.package",
+                  /*uid=*/1112345, UidType::APPLICATION)},
+         {1113999,
+          constructPackageInfo(
+                  /*packageName=*/"com.android.google.package",
+                  /*uid=*/1113999, UidType::APPLICATION)},
+         {1212345,
+          constructPackageInfo(
+                  /*packageName=*/"com.android.google.package",
+                  /*uid=*/1212345, UidType::APPLICATION)},
+         {1312345,
+          constructPackageInfo(
+                  /*packageName=*/"com.android.google.package",
+                  /*uid=*/1312345, UidType::APPLICATION)}};
+
 TEST_F(IoOveruseMonitorTest, TestOnPeriodicCollection) {
-    setUpPackagesAndConfigurations();
     sp<MockResourceOveruseListener> mockResourceOveruseListener =
             sp<MockResourceOveruseListener>::make();
     ASSERT_NO_FATAL_FAILURE(executeAsUid(1001000, [&]() {
@@ -268,11 +288,11 @@ TEST_F(IoOveruseMonitorTest, TestOnPeriodicCollection) {
      * Package "system.daemon" (UID: 1001000) exceeds warn threshold percentage of 80% but no
      * warning is issued as it is a native UID.
      */
-    sp<MockUidIoStats> mockUidIoStats = sp<MockUidIoStats>::make();
-    mockUidIoStats->expectDeltaStats(
-            {{1001000, IoUsage(0, 0, /*fgWrBytes=*/70'000, /*bgWrBytes=*/20'000, 0, 0)},
-             {1112345, IoUsage(0, 0, /*fgWrBytes=*/35'000, /*bgWrBytes=*/15'000, 0, 0)},
-             {1212345, IoUsage(0, 0, /*fgWrBytes=*/70'000, /*bgWrBytes=*/20'000, 0, 0)}});
+    EXPECT_CALL(*mMockUidStatsCollector, deltaStats())
+            .WillOnce(Return(
+                    constructUidStats({{1001000, {/*fgWrBytes=*/70'000, /*bgWrBytes=*/20'000}},
+                                       {1112345, {/*fgWrBytes=*/35'000, /*bgWrBytes=*/15'000}},
+                                       {1212345, {/*fgWrBytes=*/70'000, /*bgWrBytes=*/20'000}}})));
 
     std::vector<PackageIoOveruseStats> actualIoOveruseStats;
     EXPECT_CALL(*mMockWatchdogServiceHelper, latestIoOveruseStats(_))
@@ -282,7 +302,7 @@ TEST_F(IoOveruseMonitorTest, TestOnPeriodicCollection) {
     const auto [startTime, durationInSeconds] = calculateStartAndDuration(currentTime);
 
     ASSERT_RESULT_OK(mIoOveruseMonitor->onPeriodicCollection(currentTime, SystemState::NORMAL_MODE,
-                                                             mockUidIoStats, nullptr, nullptr));
+                                                             mMockUidStatsCollector, nullptr));
 
     std::vector<PackageIoOveruseStats> expectedIoOveruseStats =
             {constructPackageIoOveruseStats(/*uid*=*/1001000, /*shouldNotify=*/false,
@@ -308,10 +328,11 @@ TEST_F(IoOveruseMonitorTest, TestOnPeriodicCollection) {
 
     ResourceOveruseStats actualOverusingNativeStats;
     // Package "com.android.google.package" for user 11 changed uid from 1112345 to 1113999.
-    mockUidIoStats->expectDeltaStats(
-            {{1001000, IoUsage(0, 0, /*fgWrBytes=*/30'000, /*bgWrBytes=*/0, 0, 0)},
-             {1113999, IoUsage(0, 0, /*fgWrBytes=*/25'000, /*bgWrBytes=*/10'000, 0, 0)},
-             {1212345, IoUsage(0, 0, /*fgWrBytes=*/20'000, /*bgWrBytes=*/30'000, 0, 0)}});
+    EXPECT_CALL(*mMockUidStatsCollector, deltaStats())
+            .WillOnce(Return(
+                    constructUidStats({{1001000, {/*fgWrBytes=*/30'000, /*bgWrBytes=*/0}},
+                                       {1113999, {/*fgWrBytes=*/25'000, /*bgWrBytes=*/10'000}},
+                                       {1212345, {/*fgWrBytes=*/20'000, /*bgWrBytes=*/30'000}}})));
     actualIoOveruseStats.clear();
     EXPECT_CALL(*mockResourceOveruseListener, onOveruse(_))
             .WillOnce(DoAll(SaveArg<0>(&actualOverusingNativeStats), Return(Status::ok())));
@@ -319,14 +340,14 @@ TEST_F(IoOveruseMonitorTest, TestOnPeriodicCollection) {
             .WillOnce(DoAll(SaveArg<0>(&actualIoOveruseStats), Return(Status::ok())));
 
     ASSERT_RESULT_OK(mIoOveruseMonitor->onPeriodicCollection(currentTime, SystemState::NORMAL_MODE,
-                                                             mockUidIoStats, nullptr, nullptr));
+                                                             mMockUidStatsCollector, nullptr));
 
     const auto expectedOverusingNativeStats = constructResourceOveruseStats(
             constructIoOveruseStats(/*isKillable=*/false,
                                     /*remaining=*/constructPerStateBytes(0, 20'000, 100'000),
                                     /*written=*/constructPerStateBytes(100'000, 20'000, 0),
                                     /*totalOveruses=*/1, startTime, durationInSeconds));
-    EXPECT_THAT(actualOverusingNativeStats, expectedOverusingNativeStats)
+    EXPECT_THAT(actualOverusingNativeStats, Eq(expectedOverusingNativeStats))
             << "Expected: " << expectedOverusingNativeStats.toString()
             << "\nActual: " << actualOverusingNativeStats.toString();
 
@@ -360,17 +381,18 @@ TEST_F(IoOveruseMonitorTest, TestOnPeriodicCollection) {
      * Current date changed so the daily I/O usage stats should be reset and the latest I/O overuse
      * stats should not aggregate with the previous day's stats.
      */
-    mockUidIoStats->expectDeltaStats(
-            {{1001000, IoUsage(0, 0, /*fgWrBytes=*/78'000, /*bgWrBytes=*/38'000, 0, 0)},
-             {1113999, IoUsage(0, 0, /*fgWrBytes=*/55'000, /*bgWrBytes=*/23'000, 0, 0)},
-             {1212345, IoUsage(0, 0, /*fgWrBytes=*/55'000, /*bgWrBytes=*/23'000, 0, 0)}});
+    EXPECT_CALL(*mMockUidStatsCollector, deltaStats())
+            .WillOnce(Return(
+                    constructUidStats({{1001000, {/*fgWrBytes=*/78'000, /*bgWrBytes=*/38'000}},
+                                       {1113999, {/*fgWrBytes=*/55'000, /*bgWrBytes=*/23'000}},
+                                       {1212345, {/*fgWrBytes=*/55'000, /*bgWrBytes=*/23'000}}})));
     actualIoOveruseStats.clear();
     EXPECT_CALL(*mMockWatchdogServiceHelper, latestIoOveruseStats(_))
             .WillOnce(DoAll(SaveArg<0>(&actualIoOveruseStats), Return(Status::ok())));
 
     currentTime += (24 * 60 * 60);  // Change collection time to next day.
     ASSERT_RESULT_OK(mIoOveruseMonitor->onPeriodicCollection(currentTime, SystemState::NORMAL_MODE,
-                                                             mockUidIoStats, nullptr, nullptr));
+                                                             mMockUidStatsCollector, nullptr));
 
     const auto [nextDayStartTime, nextDayDuration] = calculateStartAndDuration(currentTime);
     expectedIoOveruseStats =
@@ -397,7 +419,6 @@ TEST_F(IoOveruseMonitorTest, TestOnPeriodicCollection) {
 }
 
 TEST_F(IoOveruseMonitorTest, TestOnPeriodicCollectionWithGarageMode) {
-    setUpPackagesAndConfigurations();
     sp<MockResourceOveruseListener> mockResourceOveruseListener =
             sp<MockResourceOveruseListener>::make();
     ASSERT_NO_FATAL_FAILURE(executeAsUid(1001000, [&]() {
@@ -408,11 +429,11 @@ TEST_F(IoOveruseMonitorTest, TestOnPeriodicCollectionWithGarageMode) {
      * Package "system.daemon" (UID: 1001000) exceeds warn threshold percentage of 80% but no
      * warning is issued as it is a native UID.
      */
-    sp<MockUidIoStats> mockUidIoStats = sp<MockUidIoStats>::make();
-    mockUidIoStats->expectDeltaStats(
-            {{1001000, IoUsage(0, 0, /*fgWrBytes=*/70'000, /*bgWrBytes=*/60'000, 0, 0)},
-             {1112345, IoUsage(0, 0, /*fgWrBytes=*/35'000, /*bgWrBytes=*/15'000, 0, 0)},
-             {1212345, IoUsage(0, 0, /*fgWrBytes=*/90'000, /*bgWrBytes=*/20'000, 0, 0)}});
+    EXPECT_CALL(*mMockUidStatsCollector, deltaStats())
+            .WillOnce(Return(
+                    constructUidStats({{1001000, {/*fgWrBytes=*/70'000, /*bgWrBytes=*/60'000}},
+                                       {1112345, {/*fgWrBytes=*/35'000, /*bgWrBytes=*/15'000}},
+                                       {1212345, {/*fgWrBytes=*/90'000, /*bgWrBytes=*/20'000}}})));
 
     ResourceOveruseStats actualOverusingNativeStats;
     EXPECT_CALL(*mockResourceOveruseListener, onOveruse(_))
@@ -425,14 +446,14 @@ TEST_F(IoOveruseMonitorTest, TestOnPeriodicCollectionWithGarageMode) {
     const auto [startTime, durationInSeconds] = calculateStartAndDuration(currentTime);
 
     ASSERT_RESULT_OK(mIoOveruseMonitor->onPeriodicCollection(currentTime, SystemState::GARAGE_MODE,
-                                                             mockUidIoStats, nullptr, nullptr));
+                                                             mMockUidStatsCollector, nullptr));
 
     const auto expectedOverusingNativeStats = constructResourceOveruseStats(
             constructIoOveruseStats(/*isKillable=*/false,
                                     /*remaining=*/constructPerStateBytes(80'000, 40'000, 0),
                                     /*written=*/constructPerStateBytes(0, 0, 130'000),
                                     /*totalOveruses=*/1, startTime, durationInSeconds));
-    EXPECT_THAT(actualOverusingNativeStats, expectedOverusingNativeStats)
+    EXPECT_THAT(actualOverusingNativeStats, Eq(expectedOverusingNativeStats))
             << "Expected: " << expectedOverusingNativeStats.toString()
             << "\nActual: " << actualOverusingNativeStats.toString();
 
@@ -460,11 +481,10 @@ TEST_F(IoOveruseMonitorTest, TestOnPeriodicCollectionWithGarageMode) {
 }
 
 TEST_F(IoOveruseMonitorTest, TestOnPeriodicCollectionWithZeroWriteBytes) {
-    sp<MockUidIoStats> mockUidIoStats = sp<MockUidIoStats>::make();
-    mockUidIoStats->expectDeltaStats(
-            {{1001000, IoUsage(10, 0, /*fgWrBytes=*/0, /*bgWrBytes=*/0, 1, 0)},
-             {1112345, IoUsage(0, 20, /*fgWrBytes=*/0, /*bgWrBytes=*/0, 0, 0)},
-             {1212345, IoUsage(0, 00, /*fgWrBytes=*/0, /*bgWrBytes=*/0, 0, 1)}});
+    EXPECT_CALL(*mMockUidStatsCollector, deltaStats())
+            .WillOnce(Return(constructUidStats({{1001000, {/*fgWrBytes=*/0, /*bgWrBytes=*/0}},
+                                                {1112345, {/*fgWrBytes=*/0, /*bgWrBytes=*/0}},
+                                                {1212345, {/*fgWrBytes=*/0, /*bgWrBytes=*/0}}})));
 
     EXPECT_CALL(*mMockPackageInfoResolver, getPackageInfosForUids(_)).Times(0);
     EXPECT_CALL(*mMockIoOveruseConfigs, fetchThreshold(_)).Times(0);
@@ -474,22 +494,21 @@ TEST_F(IoOveruseMonitorTest, TestOnPeriodicCollectionWithZeroWriteBytes) {
     ASSERT_RESULT_OK(
             mIoOveruseMonitor->onPeriodicCollection(std::chrono::system_clock::to_time_t(
                                                             std::chrono::system_clock::now()),
-                                                    SystemState::NORMAL_MODE, mockUidIoStats,
-                                                    nullptr, nullptr));
+                                                    SystemState::NORMAL_MODE,
+                                                    mMockUidStatsCollector, nullptr));
 }
 
 TEST_F(IoOveruseMonitorTest, TestOnPeriodicCollectionWithSmallWrittenBytes) {
-    setUpPackagesAndConfigurations();
-    sp<MockUidIoStats> mockUidIoStats = sp<MockUidIoStats>::make();
     /*
      * UID 1212345 current written bytes < |KTestMinSyncWrittenBytes| so the UID's stats are not
      * synced.
      */
-    mockUidIoStats->expectDeltaStats(
-            {{1001000, IoUsage(10, 0, /*fgWrBytes=*/59'200, /*bgWrBytes=*/0, 1, 0)},
-             {1112345, IoUsage(0, 20, /*fgWrBytes=*/0, /*bgWrBytes=*/25'200, 0, 0)},
-             {1212345, IoUsage(0, 00, /*fgWrBytes=*/300, /*bgWrBytes=*/600, 0, 1)},
-             {1312345, IoUsage(0, 00, /*fgWrBytes=*/51'200, /*bgWrBytes=*/0, 0, 1)}});
+    EXPECT_CALL(*mMockUidStatsCollector, deltaStats())
+            .WillOnce(Return(
+                    constructUidStats({{1001000, {/*fgWrBytes=*/59'200, /*bgWrBytes=*/0}},
+                                       {1112345, {/*fgWrBytes=*/0, /*bgWrBytes=*/25'200}},
+                                       {1212345, {/*fgWrBytes=*/300, /*bgWrBytes=*/600}},
+                                       {1312345, {/*fgWrBytes=*/51'200, /*bgWrBytes=*/0}}})));
 
     std::vector<PackageIoOveruseStats> actualIoOveruseStats;
     EXPECT_CALL(*mMockWatchdogServiceHelper, latestIoOveruseStats(_))
@@ -499,7 +518,7 @@ TEST_F(IoOveruseMonitorTest, TestOnPeriodicCollectionWithSmallWrittenBytes) {
     const auto [startTime, durationInSeconds] = calculateStartAndDuration(currentTime);
 
     ASSERT_RESULT_OK(mIoOveruseMonitor->onPeriodicCollection(currentTime, SystemState::NORMAL_MODE,
-                                                             mockUidIoStats, nullptr, nullptr));
+                                                             mMockUidStatsCollector, nullptr));
 
     std::vector<PackageIoOveruseStats> expectedIoOveruseStats =
             {constructPackageIoOveruseStats(/*uid*=*/1001000, /*shouldNotify=*/false,
@@ -537,19 +556,15 @@ TEST_F(IoOveruseMonitorTest, TestOnPeriodicCollectionWithSmallWrittenBytes) {
      * UID 1312345 current written bytes is < |kTestMinSyncWrittenBytes| but exceeds warn threshold
      * and killable so the UID's stat are synced.
      */
-    mockUidIoStats->expectDeltaStats(
-            {{1001000,
-              IoUsage(10, 0, /*fgWrBytes=*/KTestMinSyncWrittenBytes - 100, /*bgWrBytes=*/0, 1, 0)},
-             {1112345,
-              IoUsage(0, 20, /*fgWrBytes=*/0, /*bgWrBytes=*/KTestMinSyncWrittenBytes - 100, 0, 0)},
-             {1212345,
-              IoUsage(0, 00, /*fgWrBytes=*/KTestMinSyncWrittenBytes - 300, /*bgWrBytes=*/0, 0, 1)},
-             {1312345,
-              IoUsage(0, 00, /*fgWrBytes=*/KTestMinSyncWrittenBytes - 100, /*bgWrBytes=*/0, 0,
-                      1)}});
+    EXPECT_CALL(*mMockUidStatsCollector, deltaStats())
+            .WillOnce(Return(constructUidStats(
+                    {{1001000, {/*fgWrBytes=*/KTestMinSyncWrittenBytes - 100, /*bgWrBytes=*/0}},
+                     {1112345, {/*fgWrBytes=*/0, /*bgWrBytes=*/KTestMinSyncWrittenBytes - 100}},
+                     {1212345, {/*fgWrBytes=*/KTestMinSyncWrittenBytes - 300, /*bgWrBytes=*/0}},
+                     {1312345, {/*fgWrBytes=*/KTestMinSyncWrittenBytes - 100, /*bgWrBytes=*/0}}})));
 
     ASSERT_RESULT_OK(mIoOveruseMonitor->onPeriodicCollection(currentTime, SystemState::NORMAL_MODE,
-                                                             mockUidIoStats, nullptr, nullptr));
+                                                             mMockUidStatsCollector, nullptr));
 
     expectedIoOveruseStats =
             {constructPackageIoOveruseStats(/*uid*=*/1112345, /*shouldNotify=*/true,
@@ -573,14 +588,11 @@ TEST_F(IoOveruseMonitorTest, TestOnPeriodicCollectionWithSmallWrittenBytes) {
 }
 
 TEST_F(IoOveruseMonitorTest, TestOnPeriodicCollectionWithNoPackageInfo) {
-    sp<MockUidIoStats> mockUidIoStats = sp<MockUidIoStats>::make();
-    mockUidIoStats->expectDeltaStats(
-            {{1001000, IoUsage(0, 0, /*fgWrBytes=*/70'000, /*bgWrBytes=*/20'000, 0, 0)},
-             {1112345, IoUsage(0, 0, /*fgWrBytes=*/35'000, /*bgWrBytes=*/15'000, 0, 0)},
-             {1212345, IoUsage(0, 0, /*fgWrBytes=*/70'000, /*bgWrBytes=*/20'000, 0, 0)}});
-
-    ON_CALL(*mMockPackageInfoResolver, getPackageInfosForUids(_))
-            .WillByDefault(Return(std::unordered_map<uid_t, PackageInfo>{}));
+    EXPECT_CALL(*mMockUidStatsCollector, deltaStats())
+            .WillOnce(Return(
+                    constructUidStats({{2301000, {/*fgWrBytes=*/70'000, /*bgWrBytes=*/20'000}},
+                                       {2412345, {/*fgWrBytes=*/35'000, /*bgWrBytes=*/15'000}},
+                                       {2512345, {/*fgWrBytes=*/70'000, /*bgWrBytes=*/20'000}}})));
 
     EXPECT_CALL(*mMockIoOveruseConfigs, fetchThreshold(_)).Times(0);
     EXPECT_CALL(*mMockIoOveruseConfigs, isSafeToKill(_)).Times(0);
@@ -589,8 +601,8 @@ TEST_F(IoOveruseMonitorTest, TestOnPeriodicCollectionWithNoPackageInfo) {
     ASSERT_RESULT_OK(
             mIoOveruseMonitor->onPeriodicCollection(std::chrono::system_clock::to_time_t(
                                                             std::chrono::system_clock::now()),
-                                                    SystemState::NORMAL_MODE, mockUidIoStats,
-                                                    nullptr, nullptr));
+                                                    SystemState::NORMAL_MODE,
+                                                    mMockUidStatsCollector, nullptr));
 }
 
 TEST_F(IoOveruseMonitorTest, TestOnPeriodicMonitor) {
@@ -712,16 +724,15 @@ TEST_F(IoOveruseMonitorTest, TestUnaddIoOveruseListenerOnUnlinkToDeathError) {
 }
 
 TEST_F(IoOveruseMonitorTest, TestGetIoOveruseStats) {
-    setUpPackagesAndConfigurations();
-    sp<MockUidIoStats> mockUidIoStats = sp<MockUidIoStats>::make();
-    mockUidIoStats->expectDeltaStats(
-            {{1001000, IoUsage(0, 0, /*fgWrBytes=*/90'000, /*bgWrBytes=*/20'000, 0, 0)}});
+    EXPECT_CALL(*mMockUidStatsCollector, deltaStats())
+            .WillOnce(Return(
+                    constructUidStats({{1001000, {/*fgWrBytes=*/90'000, /*bgWrBytes=*/20'000}}})));
 
     time_t currentTime = std::chrono::system_clock::to_time_t(std::chrono::system_clock::now());
     const auto [startTime, durationInSeconds] = calculateStartAndDuration(currentTime);
 
     ASSERT_RESULT_OK(mIoOveruseMonitor->onPeriodicCollection(currentTime, SystemState::NORMAL_MODE,
-                                                             mockUidIoStats, nullptr, nullptr));
+                                                             mMockUidStatsCollector, nullptr));
 
     const auto expected =
             constructIoOveruseStats(/*isKillable=*/false,
@@ -739,16 +750,15 @@ TEST_F(IoOveruseMonitorTest, TestGetIoOveruseStats) {
 }
 
 TEST_F(IoOveruseMonitorTest, TestResetIoOveruseStats) {
-    setUpPackagesAndConfigurations();
-    sp<MockUidIoStats> mockUidIoStats = sp<MockUidIoStats>::make();
-    mockUidIoStats->expectDeltaStats(
-            {{1001000, IoUsage(0, 0, /*fgWrBytes=*/90'000, /*bgWrBytes=*/20'000, 0, 0)}});
+    EXPECT_CALL(*mMockUidStatsCollector, deltaStats())
+            .WillOnce(Return(
+                    constructUidStats({{1001000, {/*fgWrBytes=*/90'000, /*bgWrBytes=*/20'000}}})));
 
     ASSERT_RESULT_OK(
             mIoOveruseMonitor->onPeriodicCollection(std::chrono::system_clock::to_time_t(
                                                             std::chrono::system_clock::now()),
-                                                    SystemState::NORMAL_MODE, mockUidIoStats,
-                                                    nullptr, nullptr));
+                                                    SystemState::NORMAL_MODE,
+                                                    mMockUidStatsCollector, nullptr));
 
     IoOveruseStats actual;
     ASSERT_NO_FATAL_FAILURE(executeAsUid(1001000, [&]() {
diff --git a/cpp/watchdog/server/tests/MockIoOveruseMonitor.h b/cpp/watchdog/server/tests/MockIoOveruseMonitor.h
index 79ba932f40..a9d677e42d 100644
--- a/cpp/watchdog/server/tests/MockIoOveruseMonitor.h
+++ b/cpp/watchdog/server/tests/MockIoOveruseMonitor.h
@@ -35,8 +35,8 @@ public:
         ON_CALL(*this, name()).WillByDefault(::testing::Return("MockIoOveruseMonitor"));
     }
     ~MockIoOveruseMonitor() {}
-    MOCK_METHOD(bool, isInitialized, (), (override));
-    MOCK_METHOD(bool, dumpHelpText, (int), (override));
+    MOCK_METHOD(bool, isInitialized, (), (const, override));
+    MOCK_METHOD(bool, dumpHelpText, (int), (const, override));
     MOCK_METHOD(android::base::Result<void>, updateResourceOveruseConfigurations,
                 (const std::vector<
                         android::automotive::watchdog::internal::ResourceOveruseConfiguration>&),
@@ -44,7 +44,7 @@ public:
     MOCK_METHOD(
             android::base::Result<void>, getResourceOveruseConfigurations,
             (std::vector<android::automotive::watchdog::internal::ResourceOveruseConfiguration>*),
-            (override));
+            (const, override));
     MOCK_METHOD(android::base::Result<void>, actionTakenOnIoOveruse,
                 (const std::vector<
                         android::automotive::watchdog::internal::PackageResourceOveruseAction>&
@@ -54,7 +54,8 @@ public:
                 (const sp<IResourceOveruseListener>&), (override));
     MOCK_METHOD(android::base::Result<void>, removeIoOveruseListener,
                 (const sp<IResourceOveruseListener>&), (override));
-    MOCK_METHOD(android::base::Result<void>, getIoOveruseStats, (IoOveruseStats*), (override));
+    MOCK_METHOD(android::base::Result<void>, getIoOveruseStats, (IoOveruseStats*),
+                (const, override));
     MOCK_METHOD(android::base::Result<void>, resetIoOveruseStats, (const std::vector<std::string>&),
                 (override));
 };