Merge "Use single queryDetailsForDevice if possible" into main

2023-12-18 04:44:43 +00:00
parent acce4b46ae 51fc4bef49
commit 7429837eab
14 changed files with 405 additions and 196 deletions
--- a/src/com/android/settings/datausage/ChartDataUsagePreferenceController.kt
+++ b/src/com/android/settings/datausage/ChartDataUsagePreferenceController.kt
@@ -18,39 +18,22 @@ package com.android.settings.datausage

 import android.content.Context
 import android.net.NetworkTemplate
-import androidx.annotation.OpenForTesting
-import androidx.annotation.VisibleForTesting
-import androidx.lifecycle.LifecycleCoroutineScope
-import androidx.lifecycle.LifecycleOwner
-import androidx.lifecycle.lifecycleScope
 import androidx.preference.PreferenceScreen
 import com.android.settings.core.BasePreferenceController
 import com.android.settings.datausage.lib.INetworkCycleDataRepository
 import com.android.settings.datausage.lib.NetworkCycleChartData
 import com.android.settings.datausage.lib.NetworkCycleDataRepository
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.launch
-import kotlinx.coroutines.withContext

-@OpenForTesting
-open class ChartDataUsagePreferenceController(context: Context, preferenceKey: String) :
+class ChartDataUsagePreferenceController(context: Context, preferenceKey: String) :
    BasePreferenceController(context, preferenceKey) {

    private lateinit var repository: INetworkCycleDataRepository
    private lateinit var preference: ChartDataUsagePreference
-    private lateinit var lifecycleScope: LifecycleCoroutineScope
-    private var lastStartTime: Long? = null
-    private var lastEndTime: Long? = null

-    open fun init(template: NetworkTemplate) {
+    fun init(template: NetworkTemplate) {
        this.repository = NetworkCycleDataRepository(mContext, template)
    }

-    @VisibleForTesting
-    fun init(repository: INetworkCycleDataRepository) {
-        this.repository = repository
-    }
-
    override fun getAvailabilityStatus() = AVAILABLE

    override fun displayPreference(screen: PreferenceScreen) {
@@ -58,33 +41,20 @@ open class ChartDataUsagePreferenceController(context: Context, preferenceKey: S
        preference = screen.findPreference(preferenceKey)!!
    }

-    override fun onViewCreated(viewLifecycleOwner: LifecycleOwner) {
-        lifecycleScope = viewLifecycleOwner.lifecycleScope
-    }
-
    /**
     * Sets whether billing cycle modifiable.
     *
     * Don't bind warning / limit sweeps if not modifiable.
     */
-    open fun setBillingCycleModifiable(isModifiable: Boolean) {
+    fun setBillingCycleModifiable(isModifiable: Boolean) {
        preference.setNetworkPolicy(
            if (isModifiable) repository.getPolicy() else null
        )
    }

-    fun update(startTime: Long, endTime: Long) {
-        if (lastStartTime == startTime && lastEndTime == endTime) return
-        lastStartTime = startTime
-        lastEndTime = endTime
-
-        preference.setTime(startTime, endTime)
-        preference.setNetworkCycleData(NetworkCycleChartData.AllZero)
-        lifecycleScope.launch {
-            val chartData = withContext(Dispatchers.Default) {
-                repository.queryChartData(startTime, endTime)
-            }
+    /** Updates chart to show selected cycle. */
+    fun update(chartData: NetworkCycleChartData) {
+        preference.setTime(chartData.total.startTime, chartData.total.endTime)
        preference.setNetworkCycleData(chartData)
    }
-    }
 }
--- a/src/com/android/settings/datausage/DataUsageList.kt
+++ b/src/com/android/settings/datausage/DataUsageList.kt
@@ -27,6 +27,7 @@ import android.util.Log
 import android.view.View
 import androidx.annotation.OpenForTesting
 import androidx.annotation.VisibleForTesting
+import androidx.fragment.app.viewModels
 import androidx.preference.Preference
 import com.android.settings.R
 import com.android.settings.datausage.lib.BillingCycleRepository
@@ -59,6 +60,8 @@ open class DataUsageList : DataUsageBaseFragment() {
    private lateinit var chartDataUsagePreferenceController: ChartDataUsagePreferenceController
    private lateinit var billingCycleRepository: BillingCycleRepository

+    private val viewModel: DataUsageListViewModel by viewModels()
+
    @VisibleForTesting
    var dataUsageListHeaderController: DataUsageListHeaderController? = null

@@ -104,14 +107,21 @@ open class DataUsageList : DataUsageBaseFragment() {
            .collectLatestWithLifecycle(viewLifecycleOwner) { updatePolicy() }

        val template = template ?: return
+        viewModel.templateFlow.value = template
        dataUsageListHeaderController = DataUsageListHeaderController(
            setPinnedHeaderView(R.layout.apps_filter_spinner),
            template,
            metricsCategory,
            viewLifecycleOwner,
-            ::onCyclesLoad,
+            viewModel.cyclesFlow,
            ::updateSelectedCycle,
        )
+        viewModel.cyclesFlow.collectLatestWithLifecycle(viewLifecycleOwner) { cycles ->
+            dataUsageListAppsController.updateCycles(cycles)
+        }
+        viewModel.chartDataFlow.collectLatestWithLifecycle(viewLifecycleOwner) { chartData ->
+            chartDataUsagePreferenceController.update(chartData)
+        }
    }

    override fun getPreferenceScreenResId() = R.xml.data_usage_list
@@ -158,10 +168,6 @@ open class DataUsageList : DataUsageBaseFragment() {
                .getActiveSubscriptionInfo(subId) != null)
    }

-    private fun onCyclesLoad(networkUsageData: List<NetworkUsageData>) {
-        dataUsageListAppsController.updateCycles(networkUsageData)
-    }
-
    /**
     * Updates the chart and detail data when initial loaded or selected cycle changed.
     */
@@ -169,19 +175,11 @@ open class DataUsageList : DataUsageBaseFragment() {
        Log.d(TAG, "showing cycle $usageData")

        usageAmount.title = usageData.getDataUsedString(requireContext())
+        viewModel.selectedCycleFlow.value = usageData

-        updateChart(usageData)
        updateApps(usageData)
    }

-    /** Updates chart to show selected cycle. */
-    private fun updateChart(usageData: NetworkUsageData) {
-        chartDataUsagePreferenceController.update(
-            startTime = usageData.startTime,
-            endTime = usageData.endTime,
-        )
-    }
-
    /** Updates applications data usage. */
    private fun updateApps(usageData: NetworkUsageData) {
        dataUsageListAppsController.update(
--- a/src/com/android/settings/datausage/DataUsageListHeaderController.kt
+++ b/src/com/android/settings/datausage/DataUsageListHeaderController.kt
@@ -24,19 +24,13 @@ import android.view.accessibility.AccessibilityEvent
 import android.widget.AdapterView
 import android.widget.Spinner
 import androidx.annotation.OpenForTesting
-import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleOwner
-import androidx.lifecycle.lifecycleScope
-import androidx.lifecycle.repeatOnLifecycle
 import com.android.settings.R
 import com.android.settings.core.SubSettingLauncher
 import com.android.settings.datausage.CycleAdapter.SpinnerInterface
-import com.android.settings.datausage.lib.INetworkCycleDataRepository
-import com.android.settings.datausage.lib.NetworkCycleDataRepository
 import com.android.settings.datausage.lib.NetworkUsageData
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.launch
-import kotlinx.coroutines.withContext
+import com.android.settingslib.spa.framework.util.collectLatestWithLifecycle
+import kotlinx.coroutines.flow.Flow

@OpenForTesting
 open class DataUsageListHeaderController(
@@ -44,10 +38,8 @@ open class DataUsageListHeaderController(
    template: NetworkTemplate,
    sourceMetricsCategory: Int,
    viewLifecycleOwner: LifecycleOwner,
-    private val onCyclesLoad: (usageDataList: List<NetworkUsageData>) -> Unit,
+    cyclesFlow: Flow<List<NetworkUsageData>>,
    private val updateSelectedCycle: (usageData: NetworkUsageData) -> Unit,
-    private val repository: INetworkCycleDataRepository =
-        NetworkCycleDataRepository(header.context, template),
 ) {
    private val context = header.context

@@ -104,15 +96,11 @@ open class DataUsageListHeaderController(
            }
        }

-        viewLifecycleOwner.lifecycleScope.launch {
-            viewLifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
-                cycles = withContext(Dispatchers.Default) {
-                    repository.loadCycles()
-                }
+        cyclesFlow.collectLatestWithLifecycle(viewLifecycleOwner) {
+            cycles = it
            updateCycleData()
        }
    }
-    }

    open fun setConfigButtonVisible(visible: Boolean) {
        configureButton.visibility = if (visible) View.VISIBLE else View.GONE
@@ -121,7 +109,6 @@ open class DataUsageListHeaderController(
    private fun updateCycleData() {
        cycleAdapter.updateCycleList(cycles.map { Range(it.startTime, it.endTime) })
        cycleSpinner.visibility = View.VISIBLE
-        onCyclesLoad(cycles)
    }

    private fun setSelectedCycle(position: Int) {
--- a/src/com/android/settings/datausage/DataUsageListViewModel.kt
+++ b/src/com/android/settings/datausage/DataUsageListViewModel.kt
@@ -0,0 +1,71 @@
+/*
+ * 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
+
+import android.app.Application
+import android.net.NetworkTemplate
+import androidx.lifecycle.AndroidViewModel
+import androidx.lifecycle.viewModelScope
+import com.android.settings.datausage.lib.NetworkCycleBucketRepository
+import com.android.settings.datausage.lib.NetworkStatsRepository
+import com.android.settings.datausage.lib.NetworkStatsRepository.Companion.Bucket
+import com.android.settings.datausage.lib.NetworkStatsRepository.Companion.filterTime
+import com.android.settings.datausage.lib.NetworkUsageData
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.filterNotNull
+import kotlinx.coroutines.flow.flowOn
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.plus
+
+data class SelectedBuckets(
+    val selectedCycle: NetworkUsageData,
+    val buckets: List<Bucket>,
+)
+
+class DataUsageListViewModel(application: Application) : AndroidViewModel(application) {
+    private val scope = viewModelScope + Dispatchers.Default
+
+    val templateFlow = MutableStateFlow<NetworkTemplate?>(null)
+
+    private val bucketsFlow = templateFlow.filterNotNull().map { template ->
+        NetworkStatsRepository(getApplication(), template).queryDetailsForDevice()
+    }.stateIn(scope, SharingStarted.WhileSubscribed(), emptyList())
+
+    val cyclesFlow = combine(templateFlow.filterNotNull(), bucketsFlow) { template, buckets ->
+        NetworkCycleBucketRepository(application, template, buckets).loadCycles()
+    }.flowOn(Dispatchers.Default)
+
+    val selectedCycleFlow = MutableStateFlow<NetworkUsageData?>(null)
+
+    private val selectedBucketsFlow =
+        combine(selectedCycleFlow.filterNotNull(), bucketsFlow) { selectedCycle, buckets ->
+            SelectedBuckets(
+                selectedCycle = selectedCycle,
+                buckets = buckets.filterTime(selectedCycle.startTime, selectedCycle.endTime),
+            )
+        }.flowOn(Dispatchers.Default)
+
+    val chartDataFlow =
+        combine(templateFlow.filterNotNull(), selectedBucketsFlow) { template, selectedBuckets ->
+            NetworkCycleBucketRepository(application, template, selectedBuckets.buckets)
+                .queryChartData(selectedBuckets.selectedCycle)
+        }.flowOn(Dispatchers.Default)
+}
--- a/src/com/android/settings/datausage/lib/NetworkCycleBucketRepository.kt
+++ b/src/com/android/settings/datausage/lib/NetworkCycleBucketRepository.kt
@@ -0,0 +1,70 @@
+/*
+ * 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 android.text.format.DateUtils
+import android.util.Range
+import com.android.settings.datausage.lib.NetworkCycleDataRepository.Companion.bucketRange
+import com.android.settings.datausage.lib.NetworkCycleDataRepository.Companion.reverseBucketRange
+import com.android.settings.datausage.lib.NetworkStatsRepository.Companion.Bucket
+import com.android.settings.datausage.lib.NetworkStatsRepository.Companion.aggregate
+import com.android.settings.datausage.lib.NetworkStatsRepository.Companion.filterTime
+
+class NetworkCycleBucketRepository(
+    context: Context,
+    networkTemplate: NetworkTemplate,
+    private val buckets: List<Bucket>,
+    private val networkCycleDataRepository: NetworkCycleDataRepository =
+        NetworkCycleDataRepository(context, networkTemplate)
+) {
+
+    fun loadCycles(): List<NetworkUsageData> =
+        getCycles().map { aggregateUsage(it) }.filter { it.usage > 0 }
+
+    private fun getCycles(): List<Range<Long>> {
+        val policy = networkCycleDataRepository.getPolicy() ?: return queryCyclesAsFourWeeks()
+        return policy.cycleIterator().asSequence().map {
+            Range(it.lower.toInstant().toEpochMilli(), it.upper.toInstant().toEpochMilli())
+        }.toList()
+    }
+
+    private fun queryCyclesAsFourWeeks(): List<Range<Long>> {
+        val timeRange = buckets.aggregate()?.timeRange ?: return emptyList()
+        return reverseBucketRange(
+            startTime = timeRange.lower,
+            endTime = timeRange.upper,
+            step = DateUtils.WEEK_IN_MILLIS * 4,
+        )
+    }
+
+    fun queryChartData(usageData: NetworkUsageData) = NetworkCycleChartData(
+        total = usageData,
+        dailyUsage = bucketRange(
+            startTime = usageData.startTime,
+            endTime = usageData.endTime,
+            step = NetworkCycleChartData.BUCKET_DURATION.inWholeMilliseconds,
+        ).map { aggregateUsage(it) },
+    )
+
+    private fun aggregateUsage(range: Range<Long>) = NetworkUsageData(
+        startTime = range.lower,
+        endTime = range.upper,
+        usage = buckets.filterTime(range.lower, range.upper).aggregate()?.usage ?: 0,
+    )
+}
--- a/src/com/android/settings/datausage/lib/NetworkCycleChartData.kt
+++ b/src/com/android/settings/datausage/lib/NetworkCycleChartData.kt
@@ -26,11 +26,6 @@ data class NetworkCycleChartData(
    val dailyUsage: List<NetworkUsageData>,
 ) {
    companion object {
-        val AllZero = NetworkCycleChartData(
-            total = NetworkUsageData.AllZero,
-            dailyUsage = emptyList(),
-        )
-
        val BUCKET_DURATION = 1.days
    }
 }
--- a/src/com/android/settings/datausage/lib/NetworkCycleDataRepository.kt
+++ b/src/com/android/settings/datausage/lib/NetworkCycleDataRepository.kt
@@ -23,13 +23,10 @@ import android.net.NetworkTemplate
 import android.text.format.DateUtils
 import android.util.Range
 import com.android.settingslib.NetworkPolicyEditor
-import com.android.settingslib.spa.framework.util.asyncMap

 interface INetworkCycleDataRepository {
-    suspend fun loadCycles(): List<NetworkUsageData>
    fun getCycles(): List<Range<Long>>
    fun getPolicy(): NetworkPolicy?
-    suspend fun queryChartData(startTime: Long, endTime: Long): NetworkCycleChartData?
 }

 class NetworkCycleDataRepository(
@@ -41,9 +38,6 @@ class NetworkCycleDataRepository(

    private val policyManager = context.getSystemService(NetworkPolicyManager::class.java)!!

-    override suspend fun loadCycles(): List<NetworkUsageData> =
-        getCycles().queryUsage().filter { it.usage > 0 }
-
    fun loadFirstCycle(): NetworkUsageData? = getCycles().firstOrNull()?.let { queryUsage(it) }

    override fun getCycles(): List<Range<Long>> {
@@ -68,23 +62,6 @@ class NetworkCycleDataRepository(
            getPolicy(networkTemplate)
        }

-    override suspend fun queryChartData(startTime: Long, endTime: Long): NetworkCycleChartData? {
-        val usage = networkStatsRepository.querySummaryForDevice(startTime, endTime)
-        if (usage > 0L) {
-            return NetworkCycleChartData(
-                total = NetworkUsageData(startTime, endTime, usage),
-                dailyUsage = bucketRange(
-                    startTime = startTime,
-                    endTime = endTime,
-                    step = NetworkCycleChartData.BUCKET_DURATION.inWholeMilliseconds,
-                ).queryUsage(),
-            )
-        }
-        return null
-    }
-
-    private suspend fun List<Range<Long>>.queryUsage(): List<NetworkUsageData> =
-        asyncMap { queryUsage(it) }

    fun queryUsage(range: Range<Long>) = NetworkUsageData(
        startTime = range.lower,
@@ -92,10 +69,12 @@ class NetworkCycleDataRepository(
        usage = networkStatsRepository.querySummaryForDevice(range.lower, range.upper),
    )

-    private fun bucketRange(startTime: Long, endTime: Long, step: Long): List<Range<Long>> =
+    companion object {
+        fun bucketRange(startTime: Long, endTime: Long, step: Long): List<Range<Long>> =
            (startTime..endTime step step).zipWithNext(::Range)

-    private fun reverseBucketRange(startTime: Long, endTime: Long, step: Long): List<Range<Long>> =
+        fun reverseBucketRange(startTime: Long, endTime: Long, step: Long): List<Range<Long>> =
            (endTime downTo (startTime - step + 1) step step)
                .zipWithNext { bucketEnd, bucketStart -> Range(bucketStart, bucketEnd) }
+    }
 }
--- a/src/com/android/settings/datausage/lib/NetworkStatsRepository.kt
+++ b/src/com/android/settings/datausage/lib/NetworkStatsRepository.kt
@@ -33,20 +33,22 @@ class NetworkStatsRepository(context: Context, private val template: NetworkTemp
    ): NetworkUsageData? = try {
        networkStatsManager.queryDetailsForUidTagState(
            template, range.lower, range.upper, uid, NetworkStats.Bucket.TAG_NONE, state,
-        ).aggregate()
+        ).convertToBuckets().aggregate()
    } catch (e: Exception) {
        Log.e(TAG, "Exception queryDetailsForUidTagState", e)
        null
    }

-    fun getTimeRange(): Range<Long>? = try {
+    fun queryDetailsForDevice(): List<Bucket> = try {
        networkStatsManager.queryDetailsForDevice(template, Long.MIN_VALUE, Long.MAX_VALUE)
-            .aggregate()?.timeRange
+            .convertToBuckets()
    } catch (e: Exception) {
        Log.e(TAG, "Exception queryDetailsForDevice", e)
-        null
+        emptyList()
    }

+    fun getTimeRange(): Range<Long>? = queryDetailsForDevice().aggregate()?.timeRange
+
    fun querySummaryForDevice(startTime: Long, endTime: Long): Long = try {
        networkStatsManager.querySummaryForDevice(template, startTime, endTime).bytes
    } catch (e: Exception) {
@@ -70,33 +72,38 @@ class NetworkStatsRepository(context: Context, private val template: NetworkTemp
            val uid: Int,
            val bytes: Long,
            val state: Int = NetworkStats.Bucket.STATE_ALL,
+            val startTimeStamp: Long,
+            val endTimeStamp: Long,
        )

+        fun List<Bucket>.aggregate(): NetworkUsageData? = when {
+            isEmpty() -> null
+            else -> NetworkUsageData(
+                startTime = minOf { it.startTimeStamp },
+                endTime = maxOf { it.endTimeStamp },
+                usage = sumOf { it.bytes },
+            )
+        }
+
+        fun List<Bucket>.filterTime(startTime: Long, endTime: Long): List<Bucket> = filter {
+            it.startTimeStamp >= startTime && it.endTimeStamp <= endTime
+        }
+
        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, state = bucket.state)
+                buckets += Bucket(
+                    uid = bucket.uid,
+                    bytes = bucket.bytes,
+                    state = bucket.state,
+                    startTimeStamp = bucket.startTimeStamp,
+                    endTimeStamp = bucket.endTimeStamp,
+                )
            }
            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
    }
--- a/tests/spa_unit/src/com/android/settings/datausage/ChartDataUsagePreferenceControllerTest.kt
+++ b/tests/spa_unit/src/com/android/settings/datausage/ChartDataUsagePreferenceControllerTest.kt
@@ -17,16 +17,12 @@
 package com.android.settings.datausage

 import android.content.Context
-import android.util.Range
 import androidx.lifecycle.testing.TestLifecycleOwner
 import androidx.preference.PreferenceScreen
 import androidx.test.core.app.ApplicationProvider
 import androidx.test.ext.junit.runners.AndroidJUnit4
-import com.android.settings.datausage.lib.INetworkCycleDataRepository
 import com.android.settings.datausage.lib.NetworkCycleChartData
 import com.android.settings.datausage.lib.NetworkUsageData
-import kotlinx.coroutines.delay
-import kotlinx.coroutines.runBlocking
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -38,17 +34,6 @@ import org.mockito.kotlin.verify
 class ChartDataUsagePreferenceControllerTest {
    private val context: Context = ApplicationProvider.getApplicationContext()

-    private val repository = object : INetworkCycleDataRepository {
-        override suspend fun loadCycles() = emptyList<NetworkUsageData>()
-        override fun getCycles() = emptyList<Range<Long>>()
-        override fun getPolicy() = null
-
-        override suspend fun queryChartData(startTime: Long, endTime: Long) = when {
-            startTime == START_TIME && endTime == END_TIME -> CycleChartDate
-            else -> null
-        }
-    }
-
    private val preference = mock<ChartDataUsagePreference>()
    private val preferenceScreen = mock<PreferenceScreen> {
        onGeneric { findPreference(KEY) } doReturn preference
@@ -58,15 +43,13 @@ class ChartDataUsagePreferenceControllerTest {

    @Before
    fun setUp() {
-        controller.init(repository)
        controller.displayPreference(preferenceScreen)
        controller.onViewCreated(TestLifecycleOwner())
    }

    @Test
-    fun update() = runBlocking {
-        controller.update(START_TIME, END_TIME)
-        delay(100L)
+    fun update() {
+        controller.update(CycleChartDate)

        verify(preference).setTime(START_TIME, END_TIME)
        verify(preference).setNetworkCycleData(CycleChartDate)
--- a/tests/spa_unit/src/com/android/settings/datausage/DataUsageListHeaderControllerTest.kt
+++ b/tests/spa_unit/src/com/android/settings/datausage/DataUsageListHeaderControllerTest.kt
@@ -18,7 +18,6 @@ package com.android.settings.datausage

 import android.content.Context
 import android.net.NetworkTemplate
-import android.util.Range
 import android.view.LayoutInflater
 import android.view.View
 import android.widget.Spinner
@@ -27,10 +26,9 @@ import androidx.lifecycle.testing.TestLifecycleOwner
 import androidx.test.core.app.ApplicationProvider
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.settings.R
-import com.android.settings.datausage.lib.INetworkCycleDataRepository
-import com.android.settings.datausage.lib.NetworkUsageData
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.delay
+import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.runBlocking
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -47,13 +45,6 @@ class DataUsageListHeaderControllerTest {
        doNothing().whenever(mock).startActivity(any())
    }

-    private val repository = object : INetworkCycleDataRepository {
-        override suspend fun loadCycles() = emptyList<NetworkUsageData>()
-        override fun getCycles() = emptyList<Range<Long>>()
-        override fun getPolicy() = null
-        override suspend fun queryChartData(startTime: Long, endTime: Long) = null
-    }
-
    private val header =
        LayoutInflater.from(context).inflate(R.layout.apps_filter_spinner, null, false)

@@ -68,9 +59,8 @@ class DataUsageListHeaderControllerTest {
        template = mock<NetworkTemplate>(),
        sourceMetricsCategory = 0,
        viewLifecycleOwner = testLifecycleOwner,
-        onCyclesLoad = {},
+        cyclesFlow = flowOf(emptyList()),
        updateSelectedCycle = {},
-        repository = repository,
    )

    @Test
--- a/tests/spa_unit/src/com/android/settings/datausage/lib/AppDataUsageDetailsRepositoryTest.kt
+++ b/tests/spa_unit/src/com/android/settings/datausage/lib/AppDataUsageDetailsRepositoryTest.kt
@@ -49,13 +49,17 @@ class AppDataUsageDetailsRepositoryTest {
            on { queryBuckets(CYCLE1_START_TIME, CYCLE1_END_TIME) } doReturn listOf(
                Bucket(
                    uid = UID,
-                    state = NetworkStats.Bucket.STATE_DEFAULT,
                    bytes = BACKGROUND_USAGE,
+                    state = NetworkStats.Bucket.STATE_DEFAULT,
+                    startTimeStamp = 0L,
+                    endTimeStamp = 0L,
                ),
                Bucket(
                    uid = UID,
-                    state = NetworkStats.Bucket.STATE_FOREGROUND,
                    bytes = FOREGROUND_USAGE,
+                    state = NetworkStats.Bucket.STATE_FOREGROUND,
+                    startTimeStamp = 0L,
+                    endTimeStamp = 0L,
                ),
            )
        }
@@ -86,13 +90,17 @@ class AppDataUsageDetailsRepositoryTest {
            on { queryBuckets(CYCLE1_END_TIME, CYCLE2_END_TIME) } doReturn listOf(
                Bucket(
                    uid = UID,
-                    state = NetworkStats.Bucket.STATE_DEFAULT,
                    bytes = BACKGROUND_USAGE,
+                    state = NetworkStats.Bucket.STATE_DEFAULT,
+                    startTimeStamp = 0L,
+                    endTimeStamp = 0L,
                ),
                Bucket(
                    uid = UID,
-                    state = NetworkStats.Bucket.STATE_FOREGROUND,
                    bytes = FOREGROUND_USAGE,
+                    state = NetworkStats.Bucket.STATE_FOREGROUND,
+                    startTimeStamp = 0L,
+                    endTimeStamp = 0L,
                ),
            )
        }
--- a/tests/spa_unit/src/com/android/settings/datausage/lib/AppDataUsageRepositoryTest.kt
+++ b/tests/spa_unit/src/com/android/settings/datausage/lib/AppDataUsageRepositoryTest.kt
@@ -77,8 +77,8 @@ class AppDataUsageRepositoryTest {
            getPackageName = { null },
        )
        val buckets = listOf(
-            Bucket(uid = APP_ID_1, bytes = 1),
-            Bucket(uid = APP_ID_2, bytes = 2),
+            Bucket(uid = APP_ID_1, bytes = 1, startTimeStamp = 0, endTimeStamp = 0),
+            Bucket(uid = APP_ID_2, bytes = 2, startTimeStamp = 0, endTimeStamp = 0),
        )

        val appPercentList = repository.getAppPercent(null, buckets)
@@ -107,8 +107,8 @@ class AppDataUsageRepositoryTest {
            getPackageName = { if (it.key == APP_ID_1) HIDING_PACKAGE_NAME else null },
        )
        val buckets = listOf(
-            Bucket(uid = APP_ID_1, bytes = 1),
-            Bucket(uid = APP_ID_2, bytes = 2),
+            Bucket(uid = APP_ID_1, bytes = 1, startTimeStamp = 0, endTimeStamp = 0),
+            Bucket(uid = APP_ID_2, bytes = 2, startTimeStamp = 0, endTimeStamp = 0),
        )

        val appPercentList = repository.getAppPercent(HIDING_CARRIER_ID, buckets)
--- a/tests/spa_unit/src/com/android/settings/datausage/lib/NetworkCycleBucketRepositoryTest.kt
+++ b/tests/spa_unit/src/com/android/settings/datausage/lib/NetworkCycleBucketRepositoryTest.kt
@@ -0,0 +1,178 @@
+/*
+ * 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.NetworkPolicy
+import android.net.NetworkTemplate
+import android.text.format.DateUtils
+import android.util.Range
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.settings.datausage.lib.NetworkStatsRepository.Companion.Bucket
+import com.google.common.truth.Truth.assertThat
+import java.time.Instant
+import java.time.ZoneId
+import java.time.ZonedDateTime
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.stub
+
+@RunWith(AndroidJUnit4::class)
+class NetworkCycleBucketRepositoryTest {
+    private val context: Context = ApplicationProvider.getApplicationContext()
+
+    private val template = mock<NetworkTemplate>()
+
+    private val mockNetworkCycleDataRepository = mock<NetworkCycleDataRepository>()
+
+    @Test
+    fun loadCycles_byPolicy() {
+        val policy = mock<NetworkPolicy> {
+            on { cycleIterator() } doReturn listOf(
+                Range(zonedDateTime(CYCLE1_START_TIME), zonedDateTime(CYCLE1_END_TIME)),
+            ).iterator()
+        }
+        mockNetworkCycleDataRepository.stub {
+            on { getPolicy() } doReturn policy
+        }
+        val repository = NetworkCycleBucketRepository(
+            context = context,
+            networkTemplate = template,
+            buckets = listOf(
+                Bucket(
+                    uid = 0,
+                    bytes = CYCLE1_BYTES,
+                    startTimeStamp = CYCLE1_START_TIME,
+                    endTimeStamp = CYCLE1_END_TIME,
+                )
+            ),
+            networkCycleDataRepository = mockNetworkCycleDataRepository,
+        )
+
+        val cycles = repository.loadCycles()
+
+        assertThat(cycles).containsExactly(
+            NetworkUsageData(
+                startTime = CYCLE1_START_TIME,
+                endTime = CYCLE1_END_TIME,
+                usage = CYCLE1_BYTES,
+            ),
+        )
+    }
+
+    @Test
+    fun loadCycles_asFourWeeks() {
+        mockNetworkCycleDataRepository.stub {
+            on { getPolicy() } doReturn null
+        }
+        val repository = NetworkCycleBucketRepository(
+            context = context,
+            networkTemplate = template,
+            buckets = listOf(
+                Bucket(
+                    uid = 0,
+                    bytes = CYCLE2_BYTES,
+                    startTimeStamp = CYCLE2_START_TIME,
+                    endTimeStamp = CYCLE2_END_TIME,
+                )
+            ),
+            networkCycleDataRepository = mockNetworkCycleDataRepository,
+        )
+
+        val cycles = repository.loadCycles()
+
+        assertThat(cycles).containsExactly(
+            NetworkUsageData(
+                startTime = CYCLE2_END_TIME - DateUtils.WEEK_IN_MILLIS * 4,
+                endTime = CYCLE2_END_TIME,
+                usage = CYCLE2_BYTES,
+            ),
+        )
+    }
+
+    @Test
+    fun queryChartData() {
+        val cycle = NetworkUsageData(
+            startTime = CYCLE3_START_TIME,
+            endTime = CYCLE4_END_TIME,
+            usage = CYCLE3_BYTES + CYCLE4_BYTES,
+        )
+        val repository = NetworkCycleBucketRepository(
+            context = context,
+            networkTemplate = template,
+            buckets = listOf(
+                Bucket(
+                    uid = 0,
+                    bytes = CYCLE3_BYTES,
+                    startTimeStamp = CYCLE3_START_TIME,
+                    endTimeStamp = CYCLE3_END_TIME,
+                ),
+                Bucket(
+                    uid = 0,
+                    bytes = CYCLE4_BYTES,
+                    startTimeStamp = CYCLE4_START_TIME,
+                    endTimeStamp = CYCLE4_END_TIME,
+                ),
+            ),
+            networkCycleDataRepository = mockNetworkCycleDataRepository,
+        )
+
+        val summary = repository.queryChartData(cycle)
+
+        assertThat(summary).isEqualTo(
+            NetworkCycleChartData(
+                total = cycle,
+                dailyUsage = listOf(
+                    NetworkUsageData(
+                        startTime = CYCLE3_START_TIME,
+                        endTime = CYCLE3_END_TIME,
+                        usage = CYCLE3_BYTES,
+                    ),
+                    NetworkUsageData(
+                        startTime = CYCLE4_START_TIME,
+                        endTime = CYCLE4_END_TIME,
+                        usage = CYCLE4_BYTES,
+                    ),
+                ),
+            )
+        )
+    }
+
+    private fun zonedDateTime(epochMilli: Long): ZonedDateTime? =
+        ZonedDateTime.ofInstant(Instant.ofEpochMilli(epochMilli), ZoneId.systemDefault())
+
+    private companion object {
+        const val CYCLE1_START_TIME = 1L
+        const val CYCLE1_END_TIME = 2L
+        const val CYCLE1_BYTES = 11L
+
+        const val CYCLE2_START_TIME = 1695555555000L
+        const val CYCLE2_END_TIME = 1695566666000L
+        const val CYCLE2_BYTES = 22L
+
+        const val CYCLE3_START_TIME = 1695555555000L
+        const val CYCLE3_END_TIME = CYCLE3_START_TIME + DateUtils.DAY_IN_MILLIS
+        const val CYCLE3_BYTES = 33L
+
+        const val CYCLE4_START_TIME = CYCLE3_END_TIME
+        const val CYCLE4_END_TIME = CYCLE4_START_TIME + DateUtils.DAY_IN_MILLIS
+        const val CYCLE4_BYTES = 44L
+    }
+}
--- a/tests/spa_unit/src/com/android/settings/datausage/lib/NetworkCycleDataRepositoryTest.kt
+++ b/tests/spa_unit/src/com/android/settings/datausage/lib/NetworkCycleDataRepositoryTest.kt
@@ -63,7 +63,7 @@ class NetworkCycleDataRepositoryTest {
        spy(NetworkCycleDataRepository(context, template, mockNetworkStatsRepository))

    @Test
-    fun loadCycles_byPolicy() = runTest {
+    fun loadFirstCycle_byPolicy() = runTest {
        val policy = mock<NetworkPolicy> {
            on { cycleIterator() } doReturn listOf(
                Range(zonedDateTime(CYCLE1_START_TIME), zonedDateTime(CYCLE1_END_TIME))
@@ -71,23 +71,23 @@ class NetworkCycleDataRepositoryTest {
        }
        doReturn(policy).whenever(repository).getPolicy()

-        val cycles = repository.loadCycles()
+        val firstCycle = repository.loadFirstCycle()

-        assertThat(cycles).containsExactly(
+        assertThat(firstCycle).isEqualTo(
            NetworkUsageData(startTime = 1, endTime = 2, usage = CYCLE1_BYTES),
        )
    }

    @Test
-    fun loadCycles_asFourWeeks() = runTest {
+    fun loadFirstCycle_asFourWeeks() = runTest {
        doReturn(null).whenever(repository).getPolicy()
        mockNetworkStatsRepository.stub {
            on { getTimeRange() } doReturn Range(CYCLE2_START_TIME, CYCLE2_END_TIME)
        }

-        val cycles = repository.loadCycles()
+        val firstCycle = repository.loadFirstCycle()

-        assertThat(cycles).containsExactly(
+        assertThat(firstCycle).isEqualTo(
            NetworkUsageData(
                startTime = CYCLE2_END_TIME - DateUtils.WEEK_IN_MILLIS * 4,
                endTime = CYCLE2_END_TIME,
@@ -96,33 +96,6 @@ class NetworkCycleDataRepositoryTest {
        )
    }

-    @Test
-    fun querySummary() = runTest {
-        val summary = repository.queryChartData(CYCLE3_START_TIME, CYCLE4_END_TIME)
-
-        assertThat(summary).isEqualTo(
-            NetworkCycleChartData(
-                total = NetworkUsageData(
-                    startTime = CYCLE3_START_TIME,
-                    endTime = CYCLE4_END_TIME,
-                    usage = CYCLE3_BYTES + CYCLE4_BYTES,
-                ),
-                dailyUsage = listOf(
-                    NetworkUsageData(
-                        startTime = CYCLE3_START_TIME,
-                        endTime = CYCLE3_END_TIME,
-                        usage = CYCLE3_BYTES,
-                    ),
-                    NetworkUsageData(
-                        startTime = CYCLE4_START_TIME,
-                        endTime = CYCLE4_END_TIME,
-                        usage = CYCLE4_BYTES,
-                    ),
-                ),
-            )
-        )
-    }
-
    private fun zonedDateTime(epochMilli: Long): ZonedDateTime? =
        ZonedDateTime.ofInstant(Instant.ofEpochMilli(epochMilli), ZoneId.systemDefault())