Merge "Fix data usage when policy has no cycle" into 24D1-dev am: bb90a2c5c8

Original change: https://googleplex-android-review.googlesource.com/c/platform/packages/apps/Settings/+/27209264

Change-Id: Ib349fb3eed75dc0d399d9a07355e170e30d1d985
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>
diff --git a/src/com/android/settings/datausage/lib/NetworkCycleBucketRepository.kt b/src/com/android/settings/datausage/lib/NetworkCycleBucketRepository.kt
index 7e3e183..01f1965 100644
--- a/src/com/android/settings/datausage/lib/NetworkCycleBucketRepository.kt
+++ b/src/com/android/settings/datausage/lib/NetworkCycleBucketRepository.kt
@@ -18,11 +18,10 @@
 
 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.asFourWeeks
 import com.android.settings.datausage.lib.NetworkCycleDataRepository.Companion.bucketRange
 import com.android.settings.datausage.lib.NetworkCycleDataRepository.Companion.getCycles
-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
@@ -39,16 +38,11 @@
         getCycles().map { aggregateUsage(it) }.filter { it.usage > 0 }
 
     private fun getCycles(): List<Range<Long>> =
-        networkCycleDataRepository.getPolicy()?.getCycles() ?: queryCyclesAsFourWeeks()
+        networkCycleDataRepository.getPolicy()?.getCycles().orEmpty()
+            .ifEmpty { queryCyclesAsFourWeeks() }
 
-    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,
-        )
-    }
+    private fun queryCyclesAsFourWeeks(): List<Range<Long>> =
+        buckets.aggregate()?.timeRange.asFourWeeks()
 
     fun queryChartData(usageData: NetworkUsageData) = NetworkCycleChartData(
         total = usageData,
diff --git a/src/com/android/settings/datausage/lib/NetworkCycleDataRepository.kt b/src/com/android/settings/datausage/lib/NetworkCycleDataRepository.kt
index 31052ef..d55524c 100644
--- a/src/com/android/settings/datausage/lib/NetworkCycleDataRepository.kt
+++ b/src/com/android/settings/datausage/lib/NetworkCycleDataRepository.kt
@@ -42,16 +42,10 @@
     fun loadFirstCycle(): NetworkUsageData? = getCycles().firstOrNull()?.let { queryUsage(it) }
 
     override fun getCycles(): List<Range<Long>> =
-        getPolicy()?.getCycles() ?: queryCyclesAsFourWeeks()
+        getPolicy()?.getCycles().orEmpty().ifEmpty { queryCyclesAsFourWeeks() }
 
-    private fun queryCyclesAsFourWeeks(): List<Range<Long>> {
-        val timeRange = networkStatsRepository.getTimeRange() ?: return emptyList()
-        return reverseBucketRange(
-            startTime = timeRange.lower,
-            endTime = timeRange.upper,
-            step = DateUtils.WEEK_IN_MILLIS * 4,
-        )
-    }
+    private fun queryCyclesAsFourWeeks(): List<Range<Long>> =
+        networkStatsRepository.getTimeRange().asFourWeeks()
 
     override fun getPolicy(): NetworkPolicy? =
         with(NetworkPolicyEditor(policyManager)) {
@@ -59,7 +53,6 @@
             getPolicy(networkTemplate)
         }
 
-
     override fun queryUsage(range: Range<Long>) = NetworkUsageData(
         startTime = range.lower,
         endTime = range.upper,
@@ -71,6 +64,15 @@
             Range(it.lower.toInstant().toEpochMilli(), it.upper.toInstant().toEpochMilli())
         }.toList()
 
+        fun Range<Long>?.asFourWeeks(): List<Range<Long>> {
+            val timeRange = this ?: return emptyList()
+            return reverseBucketRange(
+                startTime = timeRange.lower,
+                endTime = timeRange.upper,
+                step = DateUtils.WEEK_IN_MILLIS * 4,
+            )
+        }
+
         fun bucketRange(startTime: Long, endTime: Long, step: Long): List<Range<Long>> =
             (startTime..endTime step step).zipWithNext(::Range)
 
diff --git a/tests/spa_unit/src/com/android/settings/datausage/lib/NetworkCycleBucketRepositoryTest.kt b/tests/spa_unit/src/com/android/settings/datausage/lib/NetworkCycleBucketRepositoryTest.kt
index f83b85f..81b57c9 100644
--- a/tests/spa_unit/src/com/android/settings/datausage/lib/NetworkCycleBucketRepositoryTest.kt
+++ b/tests/spa_unit/src/com/android/settings/datausage/lib/NetworkCycleBucketRepositoryTest.kt
@@ -42,6 +42,13 @@
 
     private val mockNetworkCycleDataRepository = mock<NetworkCycleDataRepository>()
 
+    private fun createRepository(buckets: List<Bucket>) = NetworkCycleBucketRepository(
+        context = context,
+        networkTemplate = template,
+        buckets = buckets,
+        networkCycleDataRepository = mockNetworkCycleDataRepository,
+    )
+
     @Test
     fun loadCycles_byPolicy() {
         val policy = mock<NetworkPolicy> {
@@ -52,9 +59,7 @@
         mockNetworkCycleDataRepository.stub {
             on { getPolicy() } doReturn policy
         }
-        val repository = NetworkCycleBucketRepository(
-            context = context,
-            networkTemplate = template,
+        val repository = createRepository(
             buckets = listOf(
                 Bucket(
                     uid = 0,
@@ -62,8 +67,7 @@
                     startTimeStamp = CYCLE1_START_TIME,
                     endTimeStamp = CYCLE1_END_TIME,
                 )
-            ),
-            networkCycleDataRepository = mockNetworkCycleDataRepository,
+            )
         )
 
         val cycles = repository.loadCycles()
@@ -78,13 +82,14 @@
     }
 
     @Test
-    fun loadCycles_asFourWeeks() {
-        mockNetworkCycleDataRepository.stub {
-            on { getPolicy() } doReturn null
+    fun loadCycles_policyHasNoCycle_asFourWeeks() {
+        val policy = mock<NetworkPolicy> {
+            on { cycleIterator() } doReturn emptyList<Range<ZonedDateTime>>().iterator()
         }
-        val repository = NetworkCycleBucketRepository(
-            context = context,
-            networkTemplate = template,
+        mockNetworkCycleDataRepository.stub {
+            on { getPolicy() } doReturn policy
+        }
+        val repository = createRepository(
             buckets = listOf(
                 Bucket(
                     uid = 0,
@@ -92,8 +97,34 @@
                     startTimeStamp = CYCLE2_START_TIME,
                     endTimeStamp = CYCLE2_END_TIME,
                 )
+            )
+        )
+
+        val cycles = repository.loadCycles()
+
+        assertThat(cycles).containsExactly(
+            NetworkUsageData(
+                startTime = CYCLE2_END_TIME - DateUtils.WEEK_IN_MILLIS * 4,
+                endTime = CYCLE2_END_TIME,
+                usage = CYCLE2_BYTES,
             ),
-            networkCycleDataRepository = mockNetworkCycleDataRepository,
+        )
+    }
+
+    @Test
+    fun loadCycles_noPolicy_asFourWeeks() {
+        mockNetworkCycleDataRepository.stub {
+            on { getPolicy() } doReturn null
+        }
+        val repository = createRepository(
+            buckets = listOf(
+                Bucket(
+                    uid = 0,
+                    bytes = CYCLE2_BYTES,
+                    startTimeStamp = CYCLE2_START_TIME,
+                    endTimeStamp = CYCLE2_END_TIME,
+                )
+            )
         )
 
         val cycles = repository.loadCycles()
@@ -114,9 +145,7 @@
             endTime = CYCLE4_END_TIME,
             usage = CYCLE3_BYTES + CYCLE4_BYTES,
         )
-        val repository = NetworkCycleBucketRepository(
-            context = context,
-            networkTemplate = template,
+        val repository = createRepository(
             buckets = listOf(
                 Bucket(
                     uid = 0,
@@ -131,7 +160,6 @@
                     endTimeStamp = CYCLE4_END_TIME,
                 ),
             ),
-            networkCycleDataRepository = mockNetworkCycleDataRepository,
         )
 
         val summary = repository.queryChartData(cycle)
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 77fe843..c011af2 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
@@ -25,6 +25,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.settings.testutils.zonedDateTime
 import com.google.common.truth.Truth.assertThat
+import java.time.ZonedDateTime
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -77,7 +78,28 @@
     }
 
     @Test
-    fun loadFirstCycle_asFourWeeks() = runTest {
+    fun loadFirstCycle_policyHasNoCycle_asFourWeeks() = runTest {
+        val policy = mock<NetworkPolicy> {
+            on { cycleIterator() } doReturn emptyList<Range<ZonedDateTime>>().iterator()
+        }
+        doReturn(policy).whenever(repository).getPolicy()
+        mockNetworkStatsRepository.stub {
+            on { getTimeRange() } doReturn Range(CYCLE2_START_TIME, CYCLE2_END_TIME)
+        }
+
+        val firstCycle = repository.loadFirstCycle()
+
+        assertThat(firstCycle).isEqualTo(
+            NetworkUsageData(
+                startTime = CYCLE2_END_TIME - DateUtils.WEEK_IN_MILLIS * 4,
+                endTime = CYCLE2_END_TIME,
+                usage = CYCLE2_BYTES,
+            ),
+        )
+    }
+
+    @Test
+    fun loadFirstCycle_noPolicy_asFourWeeks() = runTest {
         doReturn(null).whenever(repository).getPolicy()
         mockNetworkStatsRepository.stub {
             on { getTimeRange() } doReturn Range(CYCLE2_START_TIME, CYCLE2_END_TIME)