[MS46] Remove INetworkStatsSession usage from NetworkCycleDataLoader

Use public API instead. Also, since the android.app.usage.NetworkStats
is not friendly for test injection, add several util methods for
test injection.

Test: make RunSettingsLibRoboTests \ ROBOTEST_FILTER=NetworkCycleDataLoaderTest
Bug: 204830222
Change-Id: I1b8328010803050b5e1f176b47654ae18e8e1a0b
diff --git a/packages/SettingsLib/src/com/android/settingslib/net/NetworkCycleDataLoader.java b/packages/SettingsLib/src/com/android/settingslib/net/NetworkCycleDataLoader.java
index 3e95b01..5e9ac5a 100644
--- a/packages/SettingsLib/src/com/android/settingslib/net/NetworkCycleDataLoader.java
+++ b/packages/SettingsLib/src/com/android/settingslib/net/NetworkCycleDataLoader.java
@@ -16,23 +16,16 @@
 
 package com.android.settingslib.net;
 
-import static android.net.NetworkStatsHistory.FIELD_RX_BYTES;
-import static android.net.NetworkStatsHistory.FIELD_TX_BYTES;
-
+import android.annotation.NonNull;
 import android.app.usage.NetworkStats;
 import android.app.usage.NetworkStatsManager;
 import android.content.Context;
-import android.net.INetworkStatsService;
-import android.net.INetworkStatsSession;
 import android.net.NetworkPolicy;
 import android.net.NetworkPolicyManager;
-import android.net.NetworkStatsHistory;
 import android.net.NetworkTemplate;
-import android.net.TrafficStats;
-import android.os.RemoteException;
-import android.os.ServiceManager;
 import android.text.format.DateUtils;
 import android.util.Pair;
+import android.util.Range;
 
 import androidx.annotation.VisibleForTesting;
 import androidx.loader.content.AsyncTaskLoader;
@@ -52,8 +45,6 @@
     protected final NetworkTemplate mNetworkTemplate;
     private final NetworkPolicy mPolicy;
     private final ArrayList<Long> mCycles;
-    @VisibleForTesting
-    final INetworkStatsService mNetworkStatsService;
 
     protected NetworkCycleDataLoader(Builder<?> builder) {
         super(builder.mContext);
@@ -61,8 +52,6 @@
         mCycles = builder.mCycles;
         mNetworkStatsManager = (NetworkStatsManager)
             builder.mContext.getSystemService(Context.NETWORK_STATS_SERVICE);
-        mNetworkStatsService = INetworkStatsService.Stub.asInterface(
-            ServiceManager.getService(Context.NETWORK_STATS_SERVICE));
         final NetworkPolicyEditor policyEditor =
             new NetworkPolicyEditor(NetworkPolicyManager.from(builder.mContext));
         policyEditor.read();
@@ -112,23 +101,20 @@
 
     @VisibleForTesting
     void loadFourWeeksData() {
+        if (mNetworkTemplate == null) return;
+        final NetworkStats stats = mNetworkStatsManager.queryDetailsForDevice(
+                mNetworkTemplate, Long.MIN_VALUE, Long.MAX_VALUE);
         try {
-            final INetworkStatsSession networkSession = mNetworkStatsService.openSession();
-            final NetworkStatsHistory networkHistory = networkSession.getHistoryForNetwork(
-                mNetworkTemplate, FIELD_RX_BYTES | FIELD_TX_BYTES);
-            final long historyStart = networkHistory.getStart();
-            final long historyEnd = networkHistory.getEnd();
+            final Range<Long> historyTimeRange = getTimeRangeOf(stats);
 
-            long cycleEnd = historyEnd;
-            while (cycleEnd > historyStart) {
+            long cycleEnd = historyTimeRange.getUpper();
+            while (cycleEnd > historyTimeRange.getLower()) {
                 final long cycleStart = cycleEnd - (DateUtils.WEEK_IN_MILLIS * 4);
                 recordUsage(cycleStart, cycleEnd);
                 cycleEnd = cycleStart;
             }
-
-            TrafficStats.closeQuietly(networkSession);
-        } catch (RemoteException e) {
-            throw new RuntimeException(e);
+        } catch (IllegalArgumentException e) {
+            // Empty history, ignore.
         }
     }
 
@@ -169,6 +155,32 @@
         return bytes;
     }
 
+    @NonNull
+    @VisibleForTesting
+    Range getTimeRangeOf(@NonNull NetworkStats stats) {
+        long start = Long.MAX_VALUE;
+        long end = Long.MIN_VALUE;
+        while (hasNextBucket(stats)) {
+            final NetworkStats.Bucket bucket = getNextBucket(stats);
+            start = Math.min(start, bucket.getStartTimeStamp());
+            end = Math.max(end, bucket.getEndTimeStamp());
+        }
+        return new Range(start, end);
+    }
+
+    @VisibleForTesting
+    boolean hasNextBucket(@NonNull NetworkStats stats) {
+        return stats.hasNextBucket();
+    }
+
+    @NonNull
+    @VisibleForTesting
+    NetworkStats.Bucket getNextBucket(@NonNull NetworkStats stats) {
+        NetworkStats.Bucket bucket = new NetworkStats.Bucket();
+        stats.getNextBucket(bucket);
+        return bucket;
+    }
+
     @VisibleForTesting(otherwise = VisibleForTesting.NONE)
     public ArrayList<Long> getCycles() {
         return mCycles;
diff --git a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/net/NetworkCycleDataLoaderTest.java b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/net/NetworkCycleDataLoaderTest.java
index 74b9151..c79440e 100644
--- a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/net/NetworkCycleDataLoaderTest.java
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/net/NetworkCycleDataLoaderTest.java
@@ -16,24 +16,24 @@
 
 package com.android.settingslib.net;
 
-import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.ArgumentMatchers.nullable;
+import static android.app.usage.NetworkStats.Bucket.UID_ALL;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.annotation.NonNull;
+import android.app.usage.NetworkStats;
 import android.app.usage.NetworkStatsManager;
 import android.content.Context;
 import android.net.ConnectivityManager;
-import android.net.INetworkStatsService;
-import android.net.INetworkStatsSession;
 import android.net.NetworkPolicy;
 import android.net.NetworkPolicyManager;
-import android.net.NetworkStatsHistory;
 import android.net.NetworkTemplate;
-import android.os.RemoteException;
 import android.text.format.DateUtils;
 import android.util.Range;
 
@@ -49,6 +49,8 @@
 import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Queue;
+import java.util.concurrent.LinkedBlockingQueue;
 
 @RunWith(RobolectricTestRunner.class)
 public class NetworkCycleDataLoaderTest {
@@ -63,8 +65,6 @@
     private NetworkPolicy mPolicy;
     @Mock
     private Iterator<Range<ZonedDateTime>> mIterator;
-    @Mock
-    private INetworkStatsService mNetworkStatsService;
 
     private NetworkCycleDataTestLoader mLoader;
 
@@ -132,20 +132,24 @@
         verify(mLoader).recordUsage(nowInMs, nowInMs);
     }
 
+    private NetworkStats.Bucket makeMockBucket(int uid, long rxBytes, long txBytes,
+            long start, long end) {
+        NetworkStats.Bucket ret = mock(NetworkStats.Bucket.class);
+        when(ret.getUid()).thenReturn(uid);
+        when(ret.getRxBytes()).thenReturn(rxBytes);
+        when(ret.getTxBytes()).thenReturn(txBytes);
+        when(ret.getStartTimeStamp()).thenReturn(start);
+        when(ret.getEndTimeStamp()).thenReturn(end);
+        return ret;
+    }
+
     @Test
-    public void loadFourWeeksData_shouldRecordUsageForLast4Weeks() throws RemoteException {
+    public void loadFourWeeksData_shouldRecordUsageForLast4Weeks() {
         mLoader = spy(new NetworkCycleDataTestLoader(mContext));
-        ReflectionHelpers.setField(mLoader, "mNetworkStatsService", mNetworkStatsService);
-        final INetworkStatsSession networkSession = mock(INetworkStatsSession.class);
-        when(mNetworkStatsService.openSession()).thenReturn(networkSession);
-        final NetworkStatsHistory networkHistory = mock(NetworkStatsHistory.class);
-        when(networkSession.getHistoryForNetwork(nullable(NetworkTemplate.class), anyInt()))
-            .thenReturn(networkHistory);
         final long now = System.currentTimeMillis();
         final long fourWeeksAgo = now - (DateUtils.WEEK_IN_MILLIS * 4);
         final long twoDaysAgo = now - (DateUtils.DAY_IN_MILLIS * 2);
-        when(networkHistory.getStart()).thenReturn(twoDaysAgo);
-        when(networkHistory.getEnd()).thenReturn(now);
+        mLoader.addBucket(makeMockBucket(UID_ALL, 123, 456, twoDaysAgo, now));
 
         mLoader.loadFourWeeksData();
 
@@ -173,10 +177,31 @@
         verify(mLoader).recordUsage(thirtyDaysAgo, twentyDaysAgo);
     }
 
+    @Test
+    public void getTimeRangeOf() {
+        mLoader = spy(new NetworkCycleDataTestLoader(mContext));
+        // If empty, new Range(MAX_VALUE, MIN_VALUE) will be constructed. Hence, the function
+        // should throw.
+        assertThrows(IllegalArgumentException.class,
+                () -> mLoader.getTimeRangeOf(mock(NetworkStats.class)));
+
+        mLoader.addBucket(makeMockBucket(UID_ALL, 123, 456, 0, 10));
+        // Feed the function with unused NetworkStats. The actual data injection is
+        // done by addBucket.
+        assertEquals(new Range(0L, 10L), mLoader.getTimeRangeOf(mock(NetworkStats.class)));
+
+        mLoader.addBucket(makeMockBucket(UID_ALL, 123, 456, 0, 10));
+        mLoader.addBucket(makeMockBucket(UID_ALL, 123, 456, 30, 40));
+        mLoader.addBucket(makeMockBucket(UID_ALL, 123, 456, 10, 25));
+        assertEquals(new Range(0L, 40L), mLoader.getTimeRangeOf(mock(NetworkStats.class)));
+    }
+
     public class NetworkCycleDataTestLoader extends NetworkCycleDataLoader<List<NetworkCycleData>> {
+        private final Queue<NetworkStats.Bucket> mMockedBuckets = new LinkedBlockingQueue<>();
 
         private NetworkCycleDataTestLoader(Context context) {
-            super(NetworkCycleDataLoader.builder(mContext));
+            super(NetworkCycleDataLoader.builder(mContext)
+                    .setNetworkTemplate(mock(NetworkTemplate.class)));
             mContext = context;
         }
 
@@ -188,5 +213,19 @@
         List<NetworkCycleData> getCycleUsage() {
             return null;
         }
+
+        public void addBucket(NetworkStats.Bucket bucket) {
+            mMockedBuckets.add(bucket);
+        }
+
+        @Override
+        public boolean hasNextBucket(@NonNull NetworkStats unused) {
+            return !mMockedBuckets.isEmpty();
+        }
+
+        @Override
+        public NetworkStats.Bucket getNextBucket(@NonNull NetworkStats unused) {
+            return mMockedBuckets.remove();
+        }
     }
 }