Improve the performance of AppDataUsagePreference

We only need aggregated summary, but NetworkCycleDataForUidLoader loads
usage for all cycles.

Create AppDataUsageSummaryRepository to reduce the system api call to
only once.

Fix: 304421722
Test: manual - on AppInfoSettings
Test: unit test
Change-Id: I115dfb51dbf77ed3fdde985aa1a968ff7462bebc

src/com/android/settings/datausage/AppDataUsage.java

@@ -27,7 +27,6 @@ import android.content.pm.PackageManager;
 import android.graphics.drawable.Drawable;
 import android.net.NetworkTemplate;
 import android.os.Bundle;
-import android.os.Process;
 import android.os.UserHandle;
 import android.util.ArraySet;
 import android.util.IconDrawableFactory;
@@ -139,14 +138,6 @@ public class AppDataUsage extends DataUsageBaseFragment implements OnPreferenceC
             }
         }
 
-        if (mAppItem.key > 0 && UserHandle.isApp(mAppItem.key)) {
-            // In case we've been asked data usage for an app, automatically
-            // include data usage of the corresponding SDK sandbox
-            final int appSandboxUid = Process.toSdkSandboxUid(mAppItem.key);
-            if (!mAppItem.uids.get(appSandboxUid)) {
-                mAppItem.addUid(appSandboxUid);
-            }
-        }
         mTotalUsage = findPreference(KEY_TOTAL_USAGE);
         mForegroundUsage = findPreference(KEY_FOREGROUND_USAGE);
         mBackgroundUsage = findPreference(KEY_BACKGROUND_USAGE);

src/com/android/settings/datausage/lib/AppDataUsageDetailsRepository.kt

@@ -17,12 +17,11 @@
 package com.android.settings.datausage.lib
 
 import android.app.usage.NetworkStats
-import android.app.usage.NetworkStatsManager
 import android.content.Context
 import android.net.NetworkTemplate
-import android.util.Log
 import android.util.Range
 import androidx.annotation.VisibleForTesting
+import com.android.settings.datausage.lib.AppDataUsageRepository.Companion.withSdkSandboxUids
 import kotlinx.coroutines.async
 import kotlinx.coroutines.awaitAll
 import kotlinx.coroutines.coroutineScope
@@ -35,11 +34,12 @@ class AppDataUsageDetailsRepository @JvmOverloads constructor(
     context: Context,
     private val template: NetworkTemplate,
     private val cycles: List<Long>?,
-    private val uids: List<Int>,
+    uids: List<Int>,
     private val networkCycleDataRepository: INetworkCycleDataRepository =
-        NetworkCycleDataRepository(context, template)
+        NetworkCycleDataRepository(context, template),
 ) : IAppDataUsageDetailsRepository {
-    private val networkStatsManager = context.getSystemService(NetworkStatsManager::class.java)!!
+    private val withSdkSandboxUids = withSdkSandboxUids(uids)
+    private val networkStatsRepository = NetworkStatsRepository(context, template)
 
     override suspend fun queryDetailsForCycles(): List<NetworkUsageDetailsData> = coroutineScope {
         getCycles().map {
@@ -56,12 +56,11 @@ class AppDataUsageDetailsRepository @JvmOverloads constructor(
     private fun queryDetails(range: Range<Long>): NetworkUsageDetailsData {
         var totalUsage = 0L
         var foregroundUsage = 0L
-        for (uid in uids) {
+        for (uid in withSdkSandboxUids) {
             val usage = getUsage(range, uid, NetworkStats.Bucket.STATE_ALL)
             if (usage > 0L) {
                 totalUsage += usage
-                foregroundUsage +=
+                foregroundUsage += getUsage(range, uid, NetworkStats.Bucket.STATE_FOREGROUND)
-                    getUsage(range, uid, NetworkStats.Bucket.STATE_FOREGROUND)
             }
         }
         return NetworkUsageDetailsData(
@@ -73,25 +72,6 @@ class AppDataUsageDetailsRepository @JvmOverloads constructor(
     }
 
     @VisibleForTesting
-    fun getUsage(range: Range<Long>, uid: Int, state: Int): Long = try {
+    fun getUsage(range: Range<Long>, uid: Int, state: Int): Long =
-        networkStatsManager.queryDetailsForUidTagState(
+        networkStatsRepository.queryAggregateForUid(range, uid, state)?.usage ?: 0
-            template, range.lower, range.upper, uid, NetworkStats.Bucket.TAG_NONE, state,
-        ).getTotalUsage()
-    } catch (e: Exception) {
-        Log.e(TAG, "Exception querying network detail.", e)
-        0
-    }
-
-    private fun NetworkStats.getTotalUsage(): Long = use {
-        var bytes = 0L
-        val bucket = NetworkStats.Bucket()
-        while (getNextBucket(bucket)) {
-            bytes += bucket.rxBytes + bucket.txBytes
-        }
-        return bytes
-    }
-
-    private companion object {
-        private const val TAG = "AppDataUsageDetailsRepo"
-    }
 }

src/com/android/settings/datausage/lib/AppDataUsageRepository.kt

@@ -17,18 +17,17 @@
 package com.android.settings.datausage.lib
 
 import android.app.usage.NetworkStats
-import android.app.usage.NetworkStatsManager
 import android.content.Context
 import android.net.NetworkPolicyManager
 import android.net.NetworkTemplate
 import android.os.Process
 import android.os.UserHandle
-import android.util.Log
 import android.util.SparseArray
 import android.util.SparseBooleanArray
 import androidx.annotation.VisibleForTesting
 import androidx.core.util.keyIterator
 import com.android.settings.R
+import com.android.settings.datausage.lib.NetworkStatsRepository.Companion.Bucket
 import com.android.settingslib.AppItem
 import com.android.settingslib.net.UidDetailProvider
 import com.android.settingslib.spaprivileged.framework.common.userManager
@@ -36,14 +35,14 @@ import com.android.settingslib.spaprivileged.framework.common.userManager
 class AppDataUsageRepository(
     private val context: Context,
     private val currentUserId: Int,
-    private val template: NetworkTemplate,
+    template: NetworkTemplate,
     private val getPackageName: (AppItem) -> String?,
 ) {
-    private val networkStatsManager = context.getSystemService(NetworkStatsManager::class.java)!!
+    private val networkStatsRepository = NetworkStatsRepository(context, template)
 
     fun getAppPercent(carrierId: Int?, startTime: Long, endTime: Long): List<Pair<AppItem, Int>> {
-        val networkStats = querySummary(startTime, endTime) ?: return emptyList()
+        val buckets = networkStatsRepository.querySummary(startTime, endTime)
-        return getAppPercent(carrierId, convertToBuckets(networkStats))
+        return getAppPercent(carrierId, buckets)
     }
 
     @VisibleForTesting
@@ -78,13 +77,6 @@ class AppDataUsageRepository(
         }
     }
 
-    private fun querySummary(startTime: Long, endTime: Long): NetworkStats? = try {
-        networkStatsManager.querySummary(template, startTime, endTime)
-    } catch (e: RuntimeException) {
-        Log.e(TAG, "Exception querying network detail.", e)
-        null
-    }
-
     private fun filterItems(carrierId: Int?, items: List<AppItem>): List<AppItem> {
         // When there is no specified SubscriptionInfo, Wi-Fi data usage will be displayed.
         // In this case, the carrier service package also needs to be hidden.
@@ -189,14 +181,6 @@ class AppDataUsageRepository(
     }
 
     companion object {
-        private const val TAG = "AppDataUsageRepository"
-
-        @VisibleForTesting
-        data class Bucket(
-            val uid: Int,
-            val bytes: Long,
-        )
-
         @JvmStatic
         fun getAppUidList(uids: SparseBooleanArray) =
             uids.keyIterator().asSequence().map { getAppUid(it) }.distinct().toList()
@@ -210,15 +194,20 @@ class AppDataUsageRepository(
             return uid
         }
 
-        private fun convertToBuckets(stats: NetworkStats): List<Bucket> {
+        /**
-            val buckets = mutableListOf<Bucket>()
+         * Gets the apps' uids, also add the apps' SDK sandboxes' uids.
-            stats.use {
+         *
-                val bucket = NetworkStats.Bucket()
+         * In case we've been asked data usage for an app, include data usage of the corresponding
-                while (it.getNextBucket(bucket)) {
+         * SDK sandbox.
-                    buckets += Bucket(uid = bucket.uid, bytes = bucket.rxBytes + bucket.txBytes)
+         */
+        fun withSdkSandboxUids(uids: List<Int>): List<Int> {
+            val set = uids.toMutableSet()
+            for (uid in uids) {
+                if (Process.isApplicationUid(uid)) {
+                    set += Process.toSdkSandboxUid(uid)
                 }
             }
-            return buckets
+            return set.toList()
         }
     }
 }

src/com/android/settings/datausage/lib/AppDataUsageSummaryRepository.kt

@@ -0,0 +1,45 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.datausage.lib
+
+import android.content.Context
+import android.net.NetworkTemplate
+import com.android.settings.datausage.lib.AppDataUsageRepository.Companion.withSdkSandboxUids
+import com.android.settings.datausage.lib.NetworkStatsRepository.Companion.AllTimeRange
+import kotlinx.coroutines.async
+import kotlinx.coroutines.awaitAll
+import kotlinx.coroutines.coroutineScope
+
+interface IAppDataUsageSummaryRepository {
+    suspend fun querySummary(uid: Int): NetworkUsageData?
+}
+
+class AppDataUsageSummaryRepository(
+    context: Context,
+    template: NetworkTemplate,
+    private val networkStatsRepository: NetworkStatsRepository =
+        NetworkStatsRepository(context, template),
+) : IAppDataUsageSummaryRepository {
+
+    override suspend fun querySummary(uid: Int): NetworkUsageData? = coroutineScope {
+        withSdkSandboxUids(listOf(uid)).map { uid ->
+            async {
+                networkStatsRepository.queryAggregateForUid(range = AllTimeRange, uid = uid)
+            }
+        }.awaitAll().filterNotNull().aggregate()
+    }
+}

src/com/android/settings/datausage/lib/NetworkCycleDataRepository.kt

@@ -16,16 +16,12 @@
 
 package com.android.settings.datausage.lib
 
-import android.app.usage.NetworkStats
-import android.app.usage.NetworkStatsManager
 import android.content.Context
 import android.net.NetworkPolicy
 import android.net.NetworkPolicyManager
 import android.net.NetworkTemplate
 import android.text.format.DateUtils
-import android.util.Log
 import android.util.Range
-import androidx.annotation.VisibleForTesting
 import com.android.settingslib.NetworkPolicyEditor
 import kotlinx.coroutines.async
 import kotlinx.coroutines.awaitAll
@@ -41,8 +37,9 @@ interface INetworkCycleDataRepository {
 class NetworkCycleDataRepository(
     context: Context,
     private val networkTemplate: NetworkTemplate,
+    private val networkStatsRepository: NetworkStatsRepository =
+        NetworkStatsRepository(context, networkTemplate),
 ) : INetworkCycleDataRepository {
-    private val networkStatsManager = context.getSystemService(NetworkStatsManager::class.java)!!
 
     private val policyManager = context.getSystemService(NetworkPolicyManager::class.java)!!
 
@@ -57,7 +54,7 @@ class NetworkCycleDataRepository(
     }
 
     private fun queryCyclesAsFourWeeks(): List<Range<Long>> {
-        val timeRange = getTimeRange()
+        val timeRange = networkStatsRepository.getTimeRange() ?: return emptyList()
         return reverseBucketRange(
             startTime = timeRange.lower,
             endTime = timeRange.upper,
@@ -65,22 +62,6 @@ class NetworkCycleDataRepository(
         )
     }
 
-    @VisibleForTesting
-    fun getTimeRange(): Range<Long> = getTimeRangeOf(
-        networkStatsManager.queryDetailsForDevice(networkTemplate, Long.MIN_VALUE, Long.MAX_VALUE)
-    )
-
-    private fun getTimeRangeOf(stats: NetworkStats): Range<Long> {
-        var start = Long.MAX_VALUE
-        var end = Long.MIN_VALUE
-        val bucket = NetworkStats.Bucket()
-        while (stats.getNextBucket(bucket)) {
-            start = start.coerceAtMost(bucket.startTimeStamp)
-            end = end.coerceAtLeast(bucket.endTimeStamp)
-        }
-        return Range(start, end)
-    }
-
     override fun getPolicy(): NetworkPolicy? =
         with(NetworkPolicyEditor(policyManager)) {
             read()
@@ -88,7 +69,7 @@ class NetworkCycleDataRepository(
         }
 
     override suspend fun querySummary(startTime: Long, endTime: Long): NetworkCycleChartData? {
-        val usage = getUsage(startTime, endTime)
+        val usage = networkStatsRepository.querySummaryForDevice(startTime, endTime)
         if (usage > 0L) {
             return NetworkCycleChartData(
                 total = NetworkUsageData(startTime, endTime, usage),
@@ -108,7 +89,7 @@ class NetworkCycleDataRepository(
                 NetworkUsageData(
                     startTime = range.lower,
                     endTime = range.upper,
-                    usage = getUsage(range.lower, range.upper),
+                    usage = networkStatsRepository.querySummaryForDevice(range.lower, range.upper),
                 )
             }
         }.awaitAll()
@@ -139,17 +120,4 @@ class NetworkCycleDataRepository(
         }
         return buckets
     }
-
-    private fun getUsage(start: Long, end: Long): Long = try {
-        networkStatsManager.querySummaryForDevice(networkTemplate, start, end).let {
-            it.rxBytes + it.txBytes
-        }
-    } catch (e: Exception) {
-        Log.e(TAG, "Exception querying network detail.", e)
-        0
-    }
-
-    companion object {
-        private const val TAG = "NetworkCycleDataRepository"
-    }
 }

src/com/android/settings/datausage/lib/NetworkStatsRepository.kt

@@ -0,0 +1,102 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.datausage.lib
+
+import android.app.usage.NetworkStats
+import android.app.usage.NetworkStatsManager
+import android.content.Context
+import android.net.NetworkTemplate
+import android.util.Log
+import android.util.Range
+
+class NetworkStatsRepository(context: Context, private val template: NetworkTemplate) {
+    private val networkStatsManager = context.getSystemService(NetworkStatsManager::class.java)!!
+
+    fun queryAggregateForUid(
+        range: Range<Long>,
+        uid: Int,
+        state: Int = NetworkStats.Bucket.STATE_ALL,
+    ): NetworkUsageData? = try {
+        networkStatsManager.queryDetailsForUidTagState(
+            template, range.lower, range.upper, uid, NetworkStats.Bucket.TAG_NONE, state,
+        ).aggregate()
+    } catch (e: Exception) {
+        Log.e(TAG, "Exception queryDetailsForUidTagState", e)
+        null
+    }
+
+    fun getTimeRange(): Range<Long>? = try {
+        networkStatsManager.queryDetailsForDevice(template, Long.MIN_VALUE, Long.MAX_VALUE)
+            .aggregate()?.timeRange
+    } catch (e: Exception) {
+        Log.e(TAG, "Exception queryDetailsForDevice", e)
+        null
+    }
+
+    fun querySummaryForDevice(startTime: Long, endTime: Long): Long = try {
+        networkStatsManager.querySummaryForDevice(template, startTime, endTime).bytes
+    } catch (e: Exception) {
+        Log.e(TAG, "Exception querySummaryForDevice", e)
+        0
+    }
+
+    fun querySummary(startTime: Long, endTime: Long): List<Bucket> = try {
+        networkStatsManager.querySummary(template, startTime, endTime).convertToBuckets()
+    } catch (e: Exception) {
+        Log.e(TAG, "Exception querySummary", e)
+        emptyList()
+    }
+
+    companion object {
+        private const val TAG = "NetworkStatsRepository"
+
+        val AllTimeRange = Range(Long.MIN_VALUE, Long.MAX_VALUE)
+
+        data class Bucket(
+            val uid: Int,
+            val bytes: Long,
+        )
+
+        private fun NetworkStats.convertToBuckets(): List<Bucket> = use {
+            val buckets = mutableListOf<Bucket>()
+            val bucket = NetworkStats.Bucket()
+            while (getNextBucket(bucket)) {
+                buckets += Bucket(uid = bucket.uid, bytes = bucket.bytes)
+            }
+            buckets
+        }
+
+        private fun NetworkStats.aggregate(): NetworkUsageData? = use {
+            var startTime = Long.MAX_VALUE
+            var endTime = Long.MIN_VALUE
+            var usage = 0L
+            val bucket = NetworkStats.Bucket()
+            while (getNextBucket(bucket)) {
+                startTime = startTime.coerceAtMost(bucket.startTimeStamp)
+                endTime = endTime.coerceAtLeast(bucket.endTimeStamp)
+                usage += bucket.bytes
+            }
+            when {
+                startTime > endTime -> null
+                else -> NetworkUsageData(startTime, endTime, usage)
+            }
+        }
+
+        private val NetworkStats.Bucket.bytes: Long
+            get() = rxBytes + txBytes
+    }
+}

src/com/android/settings/datausage/lib/NetworkUsageData.kt

@@ -16,6 +16,8 @@
 
 package com.android.settings.datausage.lib
 
+import android.util.Range
+
 /**
  * Base data structure representing usage data in a period.
  */
@@ -23,4 +25,15 @@ data class NetworkUsageData(
     val startTime: Long,
     val endTime: Long,
     val usage: Long,
+) {
+    val timeRange = Range(startTime, endTime)
+}
+
+fun List<NetworkUsageData>.aggregate(): NetworkUsageData? = when {
+    isEmpty() -> null
+    else -> NetworkUsageData(
+        startTime = minOf { it.startTime },
+        endTime = maxOf { it.endTime },
+        usage = sumOf { it.usage },
     )
+}

src/com/android/settings/spa/app/appinfo/AppDataUsagePreference.kt

@@ -19,7 +19,6 @@ package com.android.settings.spa.app.appinfo
 import android.content.Context
 import android.content.pm.ApplicationInfo
 import android.net.NetworkTemplate
-import android.os.Process
 import android.text.format.DateUtils
 import android.text.format.Formatter
 import androidx.compose.runtime.Composable
@@ -32,11 +31,11 @@ import com.android.settings.R
 import com.android.settings.Utils
 import com.android.settings.applications.appinfo.AppInfoDashboardFragment
 import com.android.settings.datausage.AppDataUsage
+import com.android.settings.datausage.lib.AppDataUsageSummaryRepository
+import com.android.settings.datausage.lib.IAppDataUsageSummaryRepository
 import com.android.settings.datausage.lib.INetworkTemplates
 import com.android.settings.datausage.lib.NetworkTemplates
 import com.android.settings.datausage.lib.NetworkTemplates.getTitleResId
-import com.android.settingslib.net.NetworkCycleDataForUid
-import com.android.settingslib.net.NetworkCycleDataForUidLoader
 import com.android.settingslib.spa.framework.compose.toState
 import com.android.settingslib.spa.widget.preference.Preference
 import com.android.settingslib.spa.widget.preference.PreferenceModel
@@ -53,11 +52,17 @@ import kotlinx.coroutines.withContext
 fun AppDataUsagePreference(
     app: ApplicationInfo,
     networkTemplates: INetworkTemplates = NetworkTemplates,
+    repositoryFactory: (
+        context: Context,
+        networkTemplate: NetworkTemplate,
+    ) -> IAppDataUsageSummaryRepository = { context, networkTemplate ->
+        AppDataUsageSummaryRepository(context, networkTemplate)
+    }
 ) {
     val context = LocalContext.current
     val coroutineScope = rememberCoroutineScope()
     val presenter = remember {
-        AppDataUsagePresenter(context, app, coroutineScope, networkTemplates)
+        AppDataUsagePresenter(context, app, coroutineScope, networkTemplates, repositoryFactory)
     }
     if (!presenter.isAvailableFlow.collectAsStateWithLifecycle(initialValue = false).value) return
 
@@ -80,6 +85,10 @@ private class AppDataUsagePresenter(
     private val app: ApplicationInfo,
     coroutineScope: CoroutineScope,
     networkTemplates: INetworkTemplates,
+    private val repositoryFactory: (
+        context: Context,
+        networkTemplate: NetworkTemplate,
+    ) -> IAppDataUsageSummaryRepository,
 ) {
     val isAvailableFlow = flow { emit(isAvailable()) }
 
@@ -99,36 +108,18 @@ private class AppDataUsagePresenter(
     val summaryFlow = templateFlow.map { getSummary(it) }
 
     private suspend fun getSummary(template: NetworkTemplate) = withContext(Dispatchers.IO) {
-        val appUsageData = getAppUsageData(template)
+        val appUsageData = repositoryFactory(context, template).querySummary(app.uid)
-        val totalBytes = appUsageData.sumOf { it.totalUsage }
+        if (appUsageData == null || appUsageData.usage == 0L) {
-        if (totalBytes == 0L) {
             context.getString(R.string.no_data_usage)
         } else {
-            val startTime = appUsageData.minOfOrNull { it.startTime } ?: System.currentTimeMillis()
             context.getString(
                 R.string.data_summary_format,
-                Formatter.formatFileSize(context, totalBytes, Formatter.FLAG_IEC_UNITS),
+                Formatter.formatFileSize(context, appUsageData.usage, Formatter.FLAG_IEC_UNITS),
-                DateUtils.formatDateTime(context, startTime, DATE_FORMAT),
+                DateUtils.formatDateTime(context, appUsageData.startTime, DATE_FORMAT),
             )
         }
     }
 
-    private suspend fun getAppUsageData(template: NetworkTemplate): List<NetworkCycleDataForUid> =
-        withContext(Dispatchers.IO) {
-            createLoader(template).loadInBackground() ?: emptyList()
-        }
-
-    private fun createLoader(template: NetworkTemplate): NetworkCycleDataForUidLoader =
-        NetworkCycleDataForUidLoader.builder(context).apply {
-            setRetrieveDetail(false)
-            setNetworkTemplate(template)
-            addUid(app.uid)
-            if (Process.isApplicationUid(app.uid)) {
-                // Also add in network usage for the app's SDK sandbox
-                addUid(Process.toSdkSandboxUid(app.uid))
-            }
-        }.build()
-
     fun startActivity() {
         AppInfoDashboardFragment.startAppInfoFragment(
             AppDataUsage::class.java,

tests/spa_unit/src/com/android/settings/datausage/lib/AppDataUsageRepositoryTest.kt

@@ -26,7 +26,7 @@ import android.os.UserManager
 import androidx.test.core.app.ApplicationProvider
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.settings.R
-import com.android.settings.datausage.lib.AppDataUsageRepository.Companion.Bucket
+import com.android.settings.datausage.lib.NetworkStatsRepository.Companion.Bucket
 import com.android.settingslib.AppItem
 import com.android.settingslib.spaprivileged.framework.common.userManager
 import com.google.common.truth.Truth.assertThat

tests/spa_unit/src/com/android/settings/datausage/lib/AppDataUsageSummaryRepositoryTest.kt

@@ -0,0 +1,74 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.datausage.lib
+
+import android.content.Context
+import android.net.NetworkTemplate
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.runBlocking
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.mock
+
+@RunWith(AndroidJUnit4::class)
+class AppDataUsageSummaryRepositoryTest {
+
+    private val context: Context = ApplicationProvider.getApplicationContext()
+
+    private val template = mock<NetworkTemplate>()
+
+    private val networkStatsRepository = mock<NetworkStatsRepository> {
+        on {
+            queryAggregateForUid(range = NetworkStatsRepository.AllTimeRange, uid = APP_UID)
+        } doReturn NetworkUsageData(APP_START_TIME, APP_END_TIME, APP_USAGE)
+
+        on {
+            queryAggregateForUid(range = NetworkStatsRepository.AllTimeRange, uid = SDK_SANDBOX_UID)
+        } doReturn NetworkUsageData(SDK_SANDBOX_START_TIME, SDK_SANDBOX_END_TIME, SDK_SANDBOX_USAGE)
+    }
+
+    private val repository =
+        AppDataUsageSummaryRepository(context, template, networkStatsRepository)
+
+    @Test
+    fun querySummary(): Unit = runBlocking {
+        val networkUsageData = repository.querySummary(APP_UID)
+
+        assertThat(networkUsageData).isEqualTo(
+            NetworkUsageData(
+                startTime = APP_START_TIME,
+                endTime = SDK_SANDBOX_END_TIME,
+                usage = APP_USAGE + SDK_SANDBOX_USAGE,
+            )
+        )
+    }
+
+    private companion object {
+        const val APP_UID = 10000
+        const val APP_START_TIME = 10L
+        const val APP_END_TIME = 30L
+        const val APP_USAGE = 3L
+
+        const val SDK_SANDBOX_UID = 20000
+        const val SDK_SANDBOX_START_TIME = 20L
+        const val SDK_SANDBOX_END_TIME = 40L
+        const val SDK_SANDBOX_USAGE = 5L
+    }
+}

tests/spa_unit/src/com/android/settings/datausage/lib/NetworkCycleDataRepositoryTest.kt

@@ -16,8 +16,6 @@
 
 package com.android.settings.datausage.lib
 
-import android.app.usage.NetworkStats.Bucket
-import android.app.usage.NetworkStatsManager
 import android.content.Context
 import android.net.NetworkPolicy
 import android.net.NetworkTemplate
@@ -32,44 +30,37 @@ import java.time.ZonedDateTime
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.kotlin.any
 import org.mockito.kotlin.doReturn
-import org.mockito.kotlin.eq
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.spy
+import org.mockito.kotlin.stub
 import org.mockito.kotlin.whenever
 
 @RunWith(AndroidJUnit4::class)
 class NetworkCycleDataRepositoryTest {
-    private val mockNetworkStatsManager = mock<NetworkStatsManager> {
+    private val mockNetworkStatsRepository = mock<NetworkStatsRepository> {
-        on { querySummaryForDevice(any(), eq(CYCLE1_START_TIME), eq(CYCLE1_END_TIME)) } doReturn
+        on { querySummaryForDevice(CYCLE1_START_TIME, CYCLE1_END_TIME) } doReturn CYCLE1_BYTES
-            CYCLE1_BUCKET
 
         on {
             querySummaryForDevice(
-                any(),
+                startTime = CYCLE2_END_TIME - DateUtils.WEEK_IN_MILLIS * 4,
-                eq(CYCLE2_END_TIME - DateUtils.WEEK_IN_MILLIS * 4),
+                endTime = CYCLE2_END_TIME,
-                eq(CYCLE2_END_TIME),
             )
-        } doReturn CYCLE2_BUCKET
+        } doReturn CYCLE2_BYTES
 
-        on { querySummaryForDevice(any(), eq(CYCLE3_START_TIME), eq(CYCLE4_END_TIME)) } doReturn
+        on { querySummaryForDevice(CYCLE3_START_TIME, CYCLE4_END_TIME) } doReturn
-            CYCLE3_AND_4_BUCKET
+            CYCLE3_BYTES + CYCLE4_BYTES
 
-        on { querySummaryForDevice(any(), eq(CYCLE3_START_TIME), eq(CYCLE3_END_TIME)) } doReturn
+        on { querySummaryForDevice(CYCLE3_START_TIME, CYCLE3_END_TIME) } doReturn CYCLE3_BYTES
-            CYCLE3_BUCKET
+        on { querySummaryForDevice(CYCLE4_START_TIME, CYCLE4_END_TIME) } doReturn CYCLE4_BYTES
-
-        on { querySummaryForDevice(any(), eq(CYCLE4_START_TIME), eq(CYCLE4_END_TIME)) } doReturn
-            CYCLE4_BUCKET
     }
 
-    private val context: Context = spy(ApplicationProvider.getApplicationContext()) {
+    private val context: Context = ApplicationProvider.getApplicationContext()
-        on { getSystemService(NetworkStatsManager::class.java) } doReturn mockNetworkStatsManager
-    }
 
     private val template = mock<NetworkTemplate>()
 
-    private val repository = spy(NetworkCycleDataRepository(context, template))
+    private val repository =
+        spy(NetworkCycleDataRepository(context, template, mockNetworkStatsRepository))
 
     @Test
     fun loadCycles_byPolicy() = runTest {
@@ -82,13 +73,17 @@ class NetworkCycleDataRepositoryTest {
 
         val cycles = repository.loadCycles()
 
-        assertThat(cycles).containsExactly(NetworkUsageData(startTime = 1, endTime = 2, usage = 11))
+        assertThat(cycles).containsExactly(
+            NetworkUsageData(startTime = 1, endTime = 2, usage = CYCLE1_BYTES),
+        )
     }
 
     @Test
     fun loadCycles_asFourWeeks() = runTest {
         doReturn(null).whenever(repository).getPolicy()
-        doReturn(Range(CYCLE2_START_TIME, CYCLE2_END_TIME)).whenever(repository).getTimeRange()
+        mockNetworkStatsRepository.stub {
+            on { getTimeRange() } doReturn Range(CYCLE2_START_TIME, CYCLE2_END_TIME)
+        }
 
         val cycles = repository.loadCycles()
 
@@ -96,7 +91,7 @@ class NetworkCycleDataRepositoryTest {
             NetworkUsageData(
                 startTime = CYCLE2_END_TIME - DateUtils.WEEK_IN_MILLIS * 4,
                 endTime = CYCLE2_END_TIME,
-                usage = 22,
+                usage = CYCLE2_BYTES,
             ),
         )
     }
@@ -110,18 +105,18 @@ class NetworkCycleDataRepositoryTest {
                 total = NetworkUsageData(
                     startTime = CYCLE3_START_TIME,
                     endTime = CYCLE4_END_TIME,
-                    usage = 77,
+                    usage = CYCLE3_BYTES + CYCLE4_BYTES,
                 ),
                 dailyUsage = listOf(
                     NetworkUsageData(
                         startTime = CYCLE3_START_TIME,
                         endTime = CYCLE3_END_TIME,
-                        usage = 33,
+                        usage = CYCLE3_BYTES,
                     ),
                     NetworkUsageData(
                         startTime = CYCLE4_START_TIME,
                         endTime = CYCLE4_END_TIME,
-                        usage = 44,
+                        usage = CYCLE4_BYTES,
                     ),
                 ),
             )
@@ -134,35 +129,18 @@ class NetworkCycleDataRepositoryTest {
     private companion object {
         const val CYCLE1_START_TIME = 1L
         const val CYCLE1_END_TIME = 2L
-        val CYCLE1_BUCKET = mock<Bucket> {
+        const val CYCLE1_BYTES = 11L
-            on { rxBytes } doReturn 1
-            on { txBytes } doReturn 10
-        }
 
         const val CYCLE2_START_TIME = 1695555555000L
         const val CYCLE2_END_TIME = 1695566666000L
-        val CYCLE2_BUCKET = mock<Bucket> {
+        const val CYCLE2_BYTES = 22L
-            on { rxBytes } doReturn 2
-            on { txBytes } doReturn 20
-        }
 
         const val CYCLE3_START_TIME = 1695555555000L
         const val CYCLE3_END_TIME = CYCLE3_START_TIME + DateUtils.DAY_IN_MILLIS
-        val CYCLE3_BUCKET = mock<Bucket> {
+        const val CYCLE3_BYTES = 33L
-            on { rxBytes } doReturn 3
-            on { txBytes } doReturn 30
-        }
 
         const val CYCLE4_START_TIME = CYCLE3_END_TIME
         const val CYCLE4_END_TIME = CYCLE4_START_TIME + DateUtils.DAY_IN_MILLIS
-        val CYCLE4_BUCKET = mock<Bucket> {
+        const val CYCLE4_BYTES = 44L
-            on { rxBytes } doReturn 4
-            on { txBytes } doReturn 40
-        }
-
-        val CYCLE3_AND_4_BUCKET = mock<Bucket> {
-            on { rxBytes } doReturn 7
-            on { txBytes } doReturn 70
-        }
     }
 }

tests/spa_unit/src/com/android/settings/datausage/lib/NetworkStatsRepositoryTest.kt

@@ -0,0 +1,62 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.datausage.lib
+
+import android.app.usage.NetworkStats
+import android.app.usage.NetworkStatsManager
+import android.content.Context
+import android.net.NetworkTemplate
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.spy
+
+@RunWith(AndroidJUnit4::class)
+class NetworkStatsRepositoryTest {
+    private val template = mock<NetworkTemplate>()
+
+    private val mockNetworkStatsManager = mock<NetworkStatsManager> {
+        on { querySummaryForDevice(template, START_TIME, END_TIME) } doReturn BUCKET
+    }
+
+    private val context: Context = spy(ApplicationProvider.getApplicationContext()) {
+        on { getSystemService(NetworkStatsManager::class.java) } doReturn mockNetworkStatsManager
+    }
+
+    private val repository = NetworkStatsRepository(context, template)
+
+    @Test
+    fun querySummaryForDevice() {
+        val bytes = repository.querySummaryForDevice(START_TIME, END_TIME)
+
+        assertThat(bytes).isEqualTo(11)
+    }
+
+    private companion object {
+        const val START_TIME = 1L
+        const val END_TIME = 2L
+
+        val BUCKET = mock<NetworkStats.Bucket> {
+            on { rxBytes } doReturn 1
+            on { txBytes } doReturn 10
+        }
+    }
+}

tests/spa_unit/src/com/android/settings/spa/app/appinfo/AppDataUsagePreferenceTest.kt

@@ -38,9 +38,9 @@ import com.android.settings.R
 import com.android.settings.Utils
 import com.android.settings.applications.appinfo.AppInfoDashboardFragment
 import com.android.settings.datausage.AppDataUsage
+import com.android.settings.datausage.lib.IAppDataUsageSummaryRepository
 import com.android.settings.datausage.lib.INetworkTemplates
-import com.android.settingslib.net.NetworkCycleDataForUid
+import com.android.settings.datausage.lib.NetworkUsageData
-import com.android.settingslib.net.NetworkCycleDataForUidLoader
 import com.android.settingslib.spa.testutils.delay
 import com.android.settingslib.spa.testutils.waitUntilExists
 import org.junit.After
@@ -48,11 +48,7 @@ import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mock
-import org.mockito.Mockito.mock
-import org.mockito.Mockito.verify
 import org.mockito.MockitoSession
-import org.mockito.Spy
 import org.mockito.quality.Strictness
 import org.mockito.Mockito.`when` as whenever
 
@@ -63,28 +59,26 @@ class AppDataUsagePreferenceTest {
 
     private lateinit var mockSession: MockitoSession
 
-    @Spy
     private val context: Context = ApplicationProvider.getApplicationContext()
 
-    @Mock
+    private var networkUsageData: NetworkUsageData? = null
-    private lateinit var builder: NetworkCycleDataForUidLoader.Builder<NetworkCycleDataForUidLoader>
 
-    @Mock
+    private inner class TestRepository : IAppDataUsageSummaryRepository {
-    private lateinit var loader: NetworkCycleDataForUidLoader
+        override suspend fun querySummary(uid: Int): NetworkUsageData? = when (uid) {
+            UID -> networkUsageData
+            else -> null
+        }
+    }
 
     @Before
     fun setUp() {
         mockSession = mockitoSession()
             .initMocks(this)
             .mockStatic(Utils::class.java)
-            .mockStatic(NetworkCycleDataForUidLoader::class.java)
-            .mockStatic(NetworkTemplate::class.java)
             .mockStatic(AppInfoDashboardFragment::class.java)
             .strictness(Strictness.LENIENT)
             .startMocking()
         whenever(Utils.isBandwidthControlEnabled()).thenReturn(true)
-        whenever(NetworkCycleDataForUidLoader.builder(context)).thenReturn(builder)
-        whenever(builder.build()).thenReturn(loader)
     }
 
     @After
@@ -121,17 +115,9 @@ class AppDataUsagePreferenceTest {
             .assertIsEnabled()
     }
 
-    @Test
-    fun setCorrectValuesForBuilder() {
-        setContent()
-
-        verify(builder).setRetrieveDetail(false)
-        verify(builder).addUid(UID)
-    }
-
     @Test
     fun whenNoDataUsage() {
-        whenever(loader.loadInBackground()).thenReturn(emptyList())
+        networkUsageData = null
 
         setContent()
 
@@ -140,10 +126,11 @@ class AppDataUsagePreferenceTest {
 
     @Test
     fun whenHasDataUsage() {
-        val cycleData = mock(NetworkCycleDataForUid::class.java)
+        networkUsageData = NetworkUsageData(
-        whenever(cycleData.totalUsage).thenReturn(123)
+            startTime = 1666666666666L,
-        whenever(cycleData.startTime).thenReturn(1666666666666)
+            endTime = 1666666666666L,
-        whenever(loader.loadInBackground()).thenReturn(listOf(cycleData))
+            usage = 123L,
+        )
 
         setContent()
 
@@ -152,8 +139,6 @@ class AppDataUsagePreferenceTest {
 
     @Test
     fun whenClick_startActivity() {
-        whenever(loader.loadInBackground()).thenReturn(emptyList())
-
         setContent()
         composeTestRule.onRoot().performClick()
 
@@ -170,7 +155,9 @@ class AppDataUsagePreferenceTest {
     private fun setContent(app: ApplicationInfo = APP) {
         composeTestRule.setContent {
             CompositionLocalProvider(LocalContext provides context) {
-                AppDataUsagePreference(app, TestNetworkTemplates)
+                AppDataUsagePreference(app, TestNetworkTemplates) { _, _ ->
+                    TestRepository()
+                }
             }
         }
         composeTestRule.delay()