diff --git a/src/com/android/settings/datausage/ChartDataUsagePreferenceController.kt b/src/com/android/settings/datausage/ChartDataUsagePreferenceController.kt
index 0479be4..49235b5 100644
--- a/src/com/android/settings/datausage/ChartDataUsagePreferenceController.kt
+++ b/src/com/android/settings/datausage/ChartDataUsagePreferenceController.kt
@@ -74,7 +74,7 @@
         preference.setTime(startTime, endTime)
         lifecycleScope.launch {
             val chartData = withContext(Dispatchers.Default) {
-                repository.querySummary(startTime, endTime)
+                repository.queryChartData(startTime, endTime)
             }
             preference.setNetworkCycleData(chartData)
         }
diff --git a/src/com/android/settings/datausage/DataUsageList.kt b/src/com/android/settings/datausage/DataUsageList.kt
index 9ac7161..1eb883c 100644
--- a/src/com/android/settings/datausage/DataUsageList.kt
+++ b/src/com/android/settings/datausage/DataUsageList.kt
@@ -195,8 +195,7 @@
         lastDisplayedUsageData = usageData
         Log.d(TAG, "showing cycle $usageData")
 
-        val totalPhrase = DataUsageUtils.formatDataUsage(requireContext(), usageData.usage)
-        usageAmount.title = getString(R.string.data_used_template, totalPhrase)
+        usageAmount.title = usageData.getDataUsedString(requireContext())
 
         updateChart(usageData)
         updateApps(usageData)
diff --git a/src/com/android/settings/datausage/lib/NetworkCycleDataRepository.kt b/src/com/android/settings/datausage/lib/NetworkCycleDataRepository.kt
index 1ff7a81..4e97c6d 100644
--- a/src/com/android/settings/datausage/lib/NetworkCycleDataRepository.kt
+++ b/src/com/android/settings/datausage/lib/NetworkCycleDataRepository.kt
@@ -23,15 +23,13 @@
 import android.text.format.DateUtils
 import android.util.Range
 import com.android.settingslib.NetworkPolicyEditor
-import kotlinx.coroutines.async
-import kotlinx.coroutines.awaitAll
-import kotlinx.coroutines.coroutineScope
+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 querySummary(startTime: Long, endTime: Long): NetworkCycleChartData?
+    suspend fun queryChartData(startTime: Long, endTime: Long): NetworkCycleChartData?
 }
 
 class NetworkCycleDataRepository(
@@ -46,6 +44,8 @@
     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>> {
         val policy = getPolicy() ?: return queryCyclesAsFourWeeks()
         return policy.cycleIterator().asSequence().map {
@@ -68,7 +68,7 @@
             getPolicy(networkTemplate)
         }
 
-    override suspend fun querySummary(startTime: Long, endTime: Long): NetworkCycleChartData? {
+    override suspend fun queryChartData(startTime: Long, endTime: Long): NetworkCycleChartData? {
         val usage = networkStatsRepository.querySummaryForDevice(startTime, endTime)
         if (usage > 0L) {
             return NetworkCycleChartData(
@@ -83,17 +83,14 @@
         return null
     }
 
-    private suspend fun List<Range<Long>>.queryUsage(): List<NetworkUsageData> = coroutineScope {
-        map { range ->
-            async {
-                NetworkUsageData(
-                    startTime = range.lower,
-                    endTime = range.upper,
-                    usage = networkStatsRepository.querySummaryForDevice(range.lower, range.upper),
-                )
-            }
-        }.awaitAll()
-    }
+    private suspend fun List<Range<Long>>.queryUsage(): List<NetworkUsageData> =
+        asyncMap { queryUsage(it) }
+
+    fun queryUsage(range: Range<Long>) = NetworkUsageData(
+        startTime = range.lower,
+        endTime = range.upper,
+        usage = networkStatsRepository.querySummaryForDevice(range.lower, range.upper),
+    )
 
     private fun bucketRange(startTime: Long, endTime: Long, bucketSize: Long): List<Range<Long>> {
         val buckets = mutableListOf<Range<Long>>()
diff --git a/src/com/android/settings/datausage/lib/NetworkUsageData.kt b/src/com/android/settings/datausage/lib/NetworkUsageData.kt
index 5bdd7ed..93cde5f 100644
--- a/src/com/android/settings/datausage/lib/NetworkUsageData.kt
+++ b/src/com/android/settings/datausage/lib/NetworkUsageData.kt
@@ -16,7 +16,11 @@
 
 package com.android.settings.datausage.lib
 
+import android.content.Context
+import android.text.format.DateUtils
 import android.util.Range
+import com.android.settings.R
+import com.android.settings.datausage.DataUsageUtils
 
 /**
  * Base data structure representing usage data in a period.
@@ -27,6 +31,21 @@
     val usage: Long,
 ) {
     val timeRange = Range(startTime, endTime)
+
+    fun formatStartDate(context: Context): String =
+        DateUtils.formatDateTime(context, startTime, DATE_FORMAT)
+
+    fun formatDateRange(context: Context): String =
+        DateUtils.formatDateRange(context, startTime, endTime, DATE_FORMAT)
+
+    fun formatUsage(context: Context): CharSequence = DataUsageUtils.formatDataUsage(context, usage)
+
+    fun getDataUsedString(context: Context): String =
+        context.getString(R.string.data_used_template, formatUsage(context))
+
+    private companion object {
+        const val DATE_FORMAT = DateUtils.FORMAT_SHOW_DATE or DateUtils.FORMAT_ABBREV_MONTH
+    }
 }
 
 fun List<NetworkUsageData>.aggregate(): NetworkUsageData? = when {
diff --git a/src/com/android/settings/network/telephony/DataUsagePreferenceController.kt b/src/com/android/settings/network/telephony/DataUsagePreferenceController.kt
index 34433c4..d133955 100644
--- a/src/com/android/settings/network/telephony/DataUsagePreferenceController.kt
+++ b/src/com/android/settings/network/telephony/DataUsagePreferenceController.kt
@@ -31,7 +31,8 @@
 import com.android.settings.R
 import com.android.settings.datausage.DataUsageUtils
 import com.android.settings.datausage.lib.DataUsageLib
-import com.android.settingslib.net.DataUsageController
+import com.android.settings.datausage.lib.NetworkCycleDataRepository
+import com.android.settings.datausage.lib.NetworkStatsRepository.Companion.AllTimeRange
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
@@ -45,9 +46,6 @@
     private lateinit var preference: Preference
     private var networkTemplate: NetworkTemplate? = null
 
-    @VisibleForTesting
-    var dataUsageControllerFactory: (Context) -> DataUsageController = { DataUsageController(it) }
-
     fun init(subId: Int) {
         mSubId = subId
     }
@@ -103,25 +101,21 @@
         else -> null
     }
 
+    @VisibleForTesting
+    fun createNetworkCycleDataRepository(): NetworkCycleDataRepository? =
+        networkTemplate?.let { NetworkCycleDataRepository(mContext, it) }
+
     private fun getDataUsageSummary(): String? {
-        val networkTemplate = networkTemplate ?: return null
-        val controller = dataUsageControllerFactory(mContext).apply {
-            setSubscriptionId(mSubId)
-        }
-        val usageInfo = controller.getDataUsageInfo(networkTemplate)
-        if (usageInfo != null && usageInfo.usageLevel > 0) {
+        val repository = createNetworkCycleDataRepository() ?: return null
+        repository.loadFirstCycle()?.takeIf { it.usage > 0 }?.let { usageData ->
             return mContext.getString(
                 R.string.data_usage_template,
-                DataUsageUtils.formatDataUsage(mContext, usageInfo.usageLevel),
-                usageInfo.period,
+                usageData.formatUsage(mContext),
+                usageData.formatDateRange(mContext),
             )
         }
 
-        return controller.getHistoricalUsageLevel(networkTemplate).takeIf { it > 0 }?.let {
-            mContext.getString(
-                R.string.data_used_template,
-                DataUsageUtils.formatDataUsage(mContext, it),
-            )
-        }
+        return repository.queryUsage(AllTimeRange).takeIf { it.usage > 0 }
+            ?.getDataUsedString(mContext)
     }
 }
diff --git a/src/com/android/settings/spa/app/appinfo/AppDataUsagePreference.kt b/src/com/android/settings/spa/app/appinfo/AppDataUsagePreference.kt
index 6b14da0..ceb3986 100644
--- a/src/com/android/settings/spa/app/appinfo/AppDataUsagePreference.kt
+++ b/src/com/android/settings/spa/app/appinfo/AppDataUsagePreference.kt
@@ -19,8 +19,6 @@
 import android.content.Context
 import android.content.pm.ApplicationInfo
 import android.net.NetworkTemplate
-import android.text.format.DateUtils
-import android.text.format.Formatter
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberCoroutineScope
@@ -114,8 +112,8 @@
         } else {
             context.getString(
                 R.string.data_summary_format,
-                Formatter.formatFileSize(context, appUsageData.usage, Formatter.FLAG_IEC_UNITS),
-                DateUtils.formatDateTime(context, appUsageData.startTime, DATE_FORMAT),
+                appUsageData.formatUsage(context),
+                appUsageData.formatStartDate(context),
             )
         }
     }
@@ -128,8 +126,4 @@
             AppInfoSettingsProvider.METRICS_CATEGORY,
         )
     }
-
-    private companion object {
-        const val DATE_FORMAT = DateUtils.FORMAT_SHOW_DATE or DateUtils.FORMAT_ABBREV_MONTH
-    }
 }
diff --git a/tests/spa_unit/src/com/android/settings/datausage/ChartDataUsagePreferenceControllerTest.kt b/tests/spa_unit/src/com/android/settings/datausage/ChartDataUsagePreferenceControllerTest.kt
index e0eb789..ae09ef9 100644
--- a/tests/spa_unit/src/com/android/settings/datausage/ChartDataUsagePreferenceControllerTest.kt
+++ b/tests/spa_unit/src/com/android/settings/datausage/ChartDataUsagePreferenceControllerTest.kt
@@ -43,7 +43,7 @@
         override fun getCycles() = emptyList<Range<Long>>()
         override fun getPolicy() = null
 
-        override suspend fun querySummary(startTime: Long, endTime: Long) = when {
+        override suspend fun queryChartData(startTime: Long, endTime: Long) = when {
             startTime == START_TIME && endTime == END_TIME -> CycleChartDate
             else -> null
         }
diff --git a/tests/spa_unit/src/com/android/settings/datausage/DataUsageListHeaderControllerTest.kt b/tests/spa_unit/src/com/android/settings/datausage/DataUsageListHeaderControllerTest.kt
index 581f7ba..3580e68 100644
--- a/tests/spa_unit/src/com/android/settings/datausage/DataUsageListHeaderControllerTest.kt
+++ b/tests/spa_unit/src/com/android/settings/datausage/DataUsageListHeaderControllerTest.kt
@@ -51,7 +51,7 @@
         override suspend fun loadCycles() = emptyList<NetworkUsageData>()
         override fun getCycles() = emptyList<Range<Long>>()
         override fun getPolicy() = null
-        override suspend fun querySummary(startTime: Long, endTime: Long) = null
+        override suspend fun queryChartData(startTime: Long, endTime: Long) = null
     }
 
     private val header =
diff --git a/tests/spa_unit/src/com/android/settings/datausage/lib/NetworkCycleDataRepositoryTest.kt b/tests/spa_unit/src/com/android/settings/datausage/lib/NetworkCycleDataRepositoryTest.kt
index feb195d..5678503 100644
--- a/tests/spa_unit/src/com/android/settings/datausage/lib/NetworkCycleDataRepositoryTest.kt
+++ b/tests/spa_unit/src/com/android/settings/datausage/lib/NetworkCycleDataRepositoryTest.kt
@@ -98,7 +98,7 @@
 
     @Test
     fun querySummary() = runTest {
-        val summary = repository.querySummary(CYCLE3_START_TIME, CYCLE4_END_TIME)
+        val summary = repository.queryChartData(CYCLE3_START_TIME, CYCLE4_END_TIME)
 
         assertThat(summary).isEqualTo(
             NetworkCycleChartData(
diff --git a/tests/spa_unit/src/com/android/settings/network/telephony/DataUsagePreferenceControllerTest.kt b/tests/spa_unit/src/com/android/settings/network/telephony/DataUsagePreferenceControllerTest.kt
index a6d1531..069145d 100644
--- a/tests/spa_unit/src/com/android/settings/network/telephony/DataUsagePreferenceControllerTest.kt
+++ b/tests/spa_unit/src/com/android/settings/network/telephony/DataUsagePreferenceControllerTest.kt
@@ -25,7 +25,7 @@
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.testing.TestLifecycleOwner
 import androidx.preference.Preference
-import androidx.preference.PreferenceScreen
+import androidx.preference.PreferenceManager
 import androidx.test.core.app.ApplicationProvider
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.dx.mockito.inline.extended.ExtendedMockito
@@ -33,45 +33,46 @@
 import com.android.settings.core.BasePreferenceController.AVAILABLE_UNSEARCHABLE
 import com.android.settings.datausage.DataUsageUtils
 import com.android.settings.datausage.lib.DataUsageLib
-import com.android.settingslib.net.DataUsageController
-import com.android.settingslib.net.DataUsageController.DataUsageInfo
+import com.android.settings.datausage.lib.NetworkCycleDataRepository
+import com.android.settings.datausage.lib.NetworkUsageData
 import com.android.settingslib.spa.testutils.waitUntil
 import com.google.common.truth.Truth.assertThat
-import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.runBlocking
 import org.junit.After
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.ArgumentCaptor
-import org.mockito.Mock
-import org.mockito.Mockito.any
-import org.mockito.Mockito.doNothing
-import org.mockito.Mockito.verify
 import org.mockito.MockitoSession
-import org.mockito.Spy
+import org.mockito.kotlin.any
+import org.mockito.kotlin.argumentCaptor
+import org.mockito.kotlin.doNothing
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.spy
+import org.mockito.kotlin.stub
+import org.mockito.kotlin.verify
+import org.mockito.kotlin.whenever
 import org.mockito.quality.Strictness
-import org.mockito.Mockito.`when` as whenever
 
 @RunWith(AndroidJUnit4::class)
 class DataUsagePreferenceControllerTest {
 
     private lateinit var mockSession: MockitoSession
 
-    @Spy
-    private val context: Context = ApplicationProvider.getApplicationContext()
+    private val context: Context = spy(ApplicationProvider.getApplicationContext()) {
+        doNothing().whenever(mock).startActivity(any())
+    }
 
-    private lateinit var controller: DataUsagePreferenceController
+    private val preference = Preference(context).apply { key = TEST_KEY }
+    private val preferenceScreen = PreferenceManager(context).createPreferenceScreen(context)
+    private val networkTemplate = mock<NetworkTemplate>()
+    private val repository = mock<NetworkCycleDataRepository> {
+        on { queryUsage(any()) } doReturn NetworkUsageData(START_TIME, END_TIME, 0L)
+    }
 
-    private val preference = Preference(context)
-
-    @Mock
-    private lateinit var networkTemplate: NetworkTemplate
-
-    @Mock
-    private lateinit var dataUsageController: DataUsageController
-
-    @Mock
-    private lateinit var preferenceScreen: PreferenceScreen
+    private val controller = spy(DataUsagePreferenceController(context, TEST_KEY)) {
+        doReturn(repository).whenever(mock).createNetworkCycleDataRepository()
+    }
 
     @Before
     fun setUp() {
@@ -85,17 +86,15 @@
 
         whenever(SubscriptionManager.isValidSubscriptionId(SUB_ID)).thenReturn(true)
         ExtendedMockito.doReturn(true).`when` { DataUsageUtils.hasMobileData(context) }
-        ExtendedMockito.doReturn(networkTemplate)
-            .`when` { DataUsageLib.getMobileTemplate(context, SUB_ID) }
-        preference.key = TEST_KEY
-        whenever(preferenceScreen.findPreference<Preference>(TEST_KEY)).thenReturn(preference)
+        ExtendedMockito.doReturn(networkTemplate).`when` {
+            DataUsageLib.getMobileTemplate(context, SUB_ID)
+        }
 
-        controller =
-            DataUsagePreferenceController(context, TEST_KEY).apply {
-                init(SUB_ID)
-                displayPreference(preferenceScreen)
-                dataUsageControllerFactory = { dataUsageController }
-            }
+        preferenceScreen.addPreference(preference)
+        controller.apply {
+            init(SUB_ID)
+            displayPreference(preferenceScreen)
+        }
     }
 
     @After
@@ -116,26 +115,25 @@
     }
 
     @Test
-    fun handlePreferenceTreeClick_startActivity() = runTest {
-        val usageInfo = DataUsageInfo().apply {
-            usageLevel = DataUnit.MEBIBYTES.toBytes(1)
+    fun handlePreferenceTreeClick_startActivity() = runBlocking {
+        val usageData = NetworkUsageData(START_TIME, END_TIME, 1L)
+        repository.stub {
+            on { loadFirstCycle() } doReturn usageData
         }
-        whenever(dataUsageController.getDataUsageInfo(networkTemplate)).thenReturn(usageInfo)
-        doNothing().`when`(context).startActivity(any())
         controller.onViewCreated(TestLifecycleOwner(initialState = Lifecycle.State.STARTED))
         waitUntil { preference.summary != null }
 
         controller.handlePreferenceTreeClick(preference)
 
-        val captor = ArgumentCaptor.forClass(Intent::class.java)
-        verify(context).startActivity(captor.capture())
-        val intent = captor.value
+        val intent = argumentCaptor<Intent> {
+            verify(context).startActivity(capture())
+        }.firstValue
         assertThat(intent.action).isEqualTo(Settings.ACTION_MOBILE_DATA_USAGE)
         assertThat(intent.getIntExtra(Settings.EXTRA_SUB_ID, 0)).isEqualTo(SUB_ID)
     }
 
     @Test
-    fun updateState_invalidSubId_disabled() = runTest {
+    fun updateState_invalidSubId_disabled() = runBlocking {
         controller.init(SubscriptionManager.INVALID_SUBSCRIPTION_ID)
 
         controller.onViewCreated(TestLifecycleOwner(initialState = Lifecycle.State.STARTED))
@@ -144,9 +142,11 @@
     }
 
     @Test
-    fun updateState_noUsageData_shouldDisablePreference() = runTest {
-        val usageInfo = DataUsageInfo()
-        whenever(dataUsageController.getDataUsageInfo(networkTemplate)).thenReturn(usageInfo)
+    fun updateState_noUsageData_shouldDisablePreference() = runBlocking {
+        val usageData = NetworkUsageData(START_TIME, END_TIME, 0L)
+        repository.stub {
+            on { loadFirstCycle() } doReturn usageData
+        }
 
         controller.onViewCreated(TestLifecycleOwner(initialState = Lifecycle.State.STARTED))
 
@@ -154,11 +154,11 @@
     }
 
     @Test
-    fun updateState_shouldUseIecUnit() = runTest {
-        val usageInfo = DataUsageInfo().apply {
-            usageLevel = DataUnit.MEBIBYTES.toBytes(1)
+    fun updateState_shouldUseIecUnit() = runBlocking {
+        val usageData = NetworkUsageData(START_TIME, END_TIME, DataUnit.MEBIBYTES.toBytes(1))
+        repository.stub {
+            on { loadFirstCycle() } doReturn usageData
         }
-        whenever(dataUsageController.getDataUsageInfo(networkTemplate)).thenReturn(usageInfo)
 
         controller.onViewCreated(TestLifecycleOwner(initialState = Lifecycle.State.STARTED))
 
@@ -168,5 +168,7 @@
     private companion object {
         const val TEST_KEY = "test_key"
         const val SUB_ID = 2
+        const val START_TIME = 10L
+        const val END_TIME = 30L
     }
 }