Add loader for net cycle data for specific Uid.

- Add the loader to retrive usage data for a particular uid. - restructure the NetworkCycleData and the loader: - NetworkCycleData now becomes the base class of the cycle data. It only contains the cycle start/end and the total usage. - Add NetworkCycleChartData which has the bucketized usage for plotting the usage chart. - Add NetworkCycleDataForUid which shows detail usage for a specific app. Bug: 111751694 Test: make RunSettingsLibRoboTests Change-Id: Ifd95445730429373dc45aaceae37eb0d79b4377c
2018-09-27 14:57:42 -07:00
parent 9c1c9c2acb
commit e43477a25a
10 changed files with 576 additions and 111 deletions
--- a/packages/SettingsLib/src/com/android/settingslib/net/ChartDataLoaderCompat.java
+++ b/packages/SettingsLib/src/com/android/settingslib/net/ChartDataLoaderCompat.java
@@ -37,7 +37,8 @@ import com.android.settingslib.AppItem;
 /**
 * Loader for historical chart data for both network and UID details.
 *
- * Deprecated in favor of {@link NetworkCycleDataLoader}
+ * Deprecated in favor of {@link NetworkCycleChartDataLoader} and
+ * {@link NetworkCycleDataForUidLoader}
 *
 * @deprecated
 */
--- a/packages/SettingsLib/src/com/android/settingslib/net/NetworkCycleChartData.java
+++ b/packages/SettingsLib/src/com/android/settingslib/net/NetworkCycleChartData.java
@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2018 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.settingslib.net;
+
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Usage data in a billing cycle with bucketized data for plotting the usage chart.
+ */
+public class NetworkCycleChartData extends NetworkCycleData {
+    public static final long BUCKET_DURATION_MS = TimeUnit.DAYS.toMillis(1);
+
+    private List<NetworkCycleData> mUsageBuckets;
+
+    private NetworkCycleChartData() {
+    }
+
+    public List<NetworkCycleData> getUsageBuckets() {
+        return mUsageBuckets;
+    }
+
+    public static class Builder extends NetworkCycleData.Builder {
+        private NetworkCycleChartData mObject = new NetworkCycleChartData();
+
+        public Builder setUsageBuckets(List<NetworkCycleData> buckets) {
+            getObject().mUsageBuckets = buckets;
+            return this;
+        }
+
+        @Override
+        protected NetworkCycleChartData getObject() {
+            return mObject;
+        }
+
+        @Override
+        public NetworkCycleChartData build() {
+            return getObject();
+        }
+    }
+
+}
--- a/packages/SettingsLib/src/com/android/settingslib/net/NetworkCycleChartDataLoader.java
+++ b/packages/SettingsLib/src/com/android/settingslib/net/NetworkCycleChartDataLoader.java
@@ -0,0 +1,106 @@
+/*
+ * Copyright (C) 2018 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.settingslib.net;
+
+import android.app.usage.NetworkStats;
+import android.content.Context;
+import android.os.RemoteException;
+import android.util.Log;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Loader for network data usage history. It returns a list of usage data per billing cycle with
+ * bucketized usages.
+ */
+public class NetworkCycleChartDataLoader
+        extends NetworkCycleDataLoader<List<NetworkCycleChartData>> {
+
+    private static final String TAG = "NetworkCycleChartLoader";
+
+    private final List<NetworkCycleChartData> mData;
+
+    private NetworkCycleChartDataLoader(Builder builder) {
+        super(builder);
+        mData = new ArrayList<NetworkCycleChartData>();
+    }
+
+    @Override
+    void recordUsage(long start, long end) {
+        try {
+            final NetworkStats stats = mNetworkStatsManager.querySummary(
+                mNetworkType, mSubId, start, end);
+            final long total = getTotalUsage(stats);
+            if (total > 0L) {
+                final NetworkCycleChartData.Builder builder = new NetworkCycleChartData.Builder();
+                builder.setUsageBuckets(getUsageBuckets(start, end))
+                    .setStartTime(start)
+                    .setEndTime(end)
+                    .setTotalUsage(total);
+                mData.add(builder.build());
+            }
+        } catch (RemoteException e) {
+            Log.e(TAG, "Exception querying network detail.", e);
+        }
+    }
+
+    @Override
+    List<NetworkCycleChartData> getCycleUsage() {
+        return mData;
+    }
+
+    public static Builder<?> builder(Context context) {
+        return new Builder<NetworkCycleChartDataLoader>(context) {
+            @Override
+            public NetworkCycleChartDataLoader build() {
+                return new NetworkCycleChartDataLoader(this);
+            }
+        };
+    }
+
+    private List<NetworkCycleData> getUsageBuckets(long start, long end) {
+        final List<NetworkCycleData> data = new ArrayList<>();
+        long bucketStart = start;
+        long bucketEnd = start + NetworkCycleChartData.BUCKET_DURATION_MS;
+        while (bucketEnd <= end) {
+            long usage = 0L;
+            try {
+                final NetworkStats stats = mNetworkStatsManager.querySummary(
+                    mNetworkType, mSubId, bucketStart, bucketEnd);
+                usage = getTotalUsage(stats);
+            } catch (RemoteException e) {
+                Log.e(TAG, "Exception querying network detail.", e);
+            }
+            data.add(new NetworkCycleData.Builder()
+                .setStartTime(bucketStart).setEndTime(bucketEnd).setTotalUsage(usage).build());
+            bucketStart = bucketEnd;
+            bucketEnd += NetworkCycleChartData.BUCKET_DURATION_MS;
+        }
+        return data;
+    }
+
+    public static abstract class Builder<T extends NetworkCycleChartDataLoader>
+            extends NetworkCycleDataLoader.Builder<T> {
+
+        public Builder(Context context) {
+            super(context);
+        }
+
+    }
+
+}
--- a/packages/SettingsLib/src/com/android/settingslib/net/NetworkCycleData.java
+++ b/packages/SettingsLib/src/com/android/settingslib/net/NetworkCycleData.java
@@ -16,54 +16,55 @@

 package com.android.settingslib.net;

-import java.util.List;
-import java.util.concurrent.TimeUnit;
-
 /**
- * Data structure representing usage data in a billing cycle.
+ * Base data structure representing usage data in a billing cycle.
 */
 public class NetworkCycleData {
-    public static final long BUCKET_DURATION_MS = TimeUnit.DAYS.toMillis(1);
-    public long startTime;
-    public long endTime;
-    public long totalUsage;
-    public List<NetworkCycleData> usageBuckets;

-    private NetworkCycleData(Builder builder) {
-        startTime = builder.mStart;
-        endTime = builder.mEnd;
-        totalUsage = builder.mTotalUsage;
-        usageBuckets = builder.mUsageBuckets;
+    private long mStartTime;
+    private long mEndTime;
+    private long mTotalUsage;
+
+    protected NetworkCycleData() {
+    }
+
+    public long getStartTime() {
+        return mStartTime;
+    }
+
+    public long getEndTime() {
+        return mEndTime;
+    }
+
+    public long getTotalUsage() {
+        return mTotalUsage;
    }

    public static class Builder {
-        private long mStart;
-        private long mEnd;
-        private long mTotalUsage;
-        private List<NetworkCycleData> mUsageBuckets;
+
+        private NetworkCycleData mObject = new NetworkCycleData();

        public Builder setStartTime(long start) {
-            mStart = start;
+            getObject().mStartTime = start;
            return this;
        }

        public Builder setEndTime(long end) {
-            mEnd = end;
+            getObject().mEndTime = end;
            return this;
        }

        public Builder setTotalUsage(long total) {
-            mTotalUsage = total;
+            getObject().mTotalUsage = total;
            return this;
        }

-        public Builder setUsageBuckets(List<NetworkCycleData> buckets) {
-            mUsageBuckets = buckets;
-            return this;
+        protected NetworkCycleData getObject() {
+            return mObject;
        }

        public NetworkCycleData build() {
-            return new NetworkCycleData(this);
+            return getObject();
        }
    }
 }
--- a/packages/SettingsLib/src/com/android/settingslib/net/NetworkCycleDataForUid.java
+++ b/packages/SettingsLib/src/com/android/settingslib/net/NetworkCycleDataForUid.java
@@ -0,0 +1,65 @@
+/*
+ * Copyright (C) 2018 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.settingslib.net;
+
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Usage data in a billing cycle for a specific Uid.
+ */
+public class NetworkCycleDataForUid extends NetworkCycleData {
+
+    private long mBackgroudUsage;
+    private long mForegroudUsage;
+
+    private NetworkCycleDataForUid() {
+    }
+
+    public long getBackgroudUsage() {
+        return mBackgroudUsage;
+    }
+
+    public long getForegroudUsage() {
+        return mForegroudUsage;
+    }
+
+    public static class Builder extends NetworkCycleData.Builder {
+
+        private NetworkCycleDataForUid mObject = new NetworkCycleDataForUid();
+
+        public Builder setBackgroundUsage(long backgroundUsage) {
+            getObject().mBackgroudUsage = backgroundUsage;
+            return this;
+        }
+
+        public Builder setForegroundUsage(long foregroundUsage) {
+            getObject().mForegroudUsage = foregroundUsage;
+            return this;
+        }
+
+        @Override
+        public NetworkCycleDataForUid getObject() {
+            return mObject;
+        }
+
+        @Override
+        public NetworkCycleDataForUid build() {
+            return getObject();
+        }
+    }
+
+}
--- a/packages/SettingsLib/src/com/android/settingslib/net/NetworkCycleDataForUidLoader.java
+++ b/packages/SettingsLib/src/com/android/settingslib/net/NetworkCycleDataForUidLoader.java
@@ -0,0 +1,102 @@
+/*
+ * Copyright (C) 2018 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.settingslib.net;
+
+import static android.app.usage.NetworkStats.Bucket.STATE_FOREGROUND;
+import static android.net.NetworkStats.TAG_NONE;
+
+import android.app.usage.NetworkStats;
+import android.content.Context;
+import android.util.Log;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Loader for network data usage history. It returns a list of usage data per billing cycle for a
+ * specific Uid.
+ */
+public class NetworkCycleDataForUidLoader extends
+        NetworkCycleDataLoader<List<NetworkCycleDataForUid>> {
+    private static final String TAG = "NetworkDataForUidLoader";
+
+    private final List<NetworkCycleDataForUid> mData;
+    private final int mUid;
+
+    private NetworkCycleDataForUidLoader(Builder builder) {
+        super(builder);
+        mUid = builder.mUid;
+        mData = new ArrayList<NetworkCycleDataForUid>();
+    }
+
+    @Override
+    void recordUsage(long start, long end) {
+        try {
+            final NetworkStats stats = mNetworkStatsManager.queryDetailsForUid(
+                mNetworkType, mSubId, start, end, mUid);
+            final long total = getTotalUsage(stats);
+            if (total > 0L) {
+                final long foreground = getForegroundUsage(start, end);
+                final NetworkCycleDataForUid.Builder builder = new NetworkCycleDataForUid.Builder();
+                builder.setBackgroundUsage(total - foreground)
+                    .setForegroundUsage(foreground)
+                    .setStartTime(start)
+                    .setEndTime(end)
+                    .setTotalUsage(total);
+                mData.add(builder.build());
+            }
+        } catch (Exception e) {
+            Log.e(TAG, "Exception querying network detail.", e);
+        }
+    }
+
+    @Override
+    List<NetworkCycleDataForUid> getCycleUsage() {
+        return mData;
+    }
+
+    public static Builder<?> builder(Context context) {
+        return new Builder<NetworkCycleDataForUidLoader>(context) {
+            @Override
+            public NetworkCycleDataForUidLoader build() {
+                return new NetworkCycleDataForUidLoader(this);
+            }
+        };
+    }
+
+    private long getForegroundUsage(long start, long end) {
+        final NetworkStats stats = mNetworkStatsManager.queryDetailsForUidTagState(
+            mNetworkType, mSubId, start, end, mUid, TAG_NONE, STATE_FOREGROUND);
+        return getTotalUsage(stats);
+    }
+
+    public static abstract class Builder<T extends NetworkCycleDataForUidLoader>
+            extends NetworkCycleDataLoader.Builder<T> {
+
+        private int mUid;
+
+        public Builder(Context context) {
+            super(context);
+        }
+
+        public Builder<T> setUid(int uid) {
+            mUid = uid;
+            return this;
+        }
+    }
+
+}
--- a/packages/SettingsLib/src/com/android/settingslib/net/NetworkCycleDataLoader.java
+++ b/packages/SettingsLib/src/com/android/settingslib/net/NetworkCycleDataLoader.java
@@ -32,32 +32,28 @@ import android.net.TrafficStats;
 import android.os.RemoteException;
 import android.os.ServiceManager;
 import android.text.format.DateUtils;
-import android.util.Log;
 import android.util.Pair;

 import java.time.ZonedDateTime;
-import java.util.ArrayList;
 import java.util.Iterator;
-import java.util.List;

-import androidx.annotation.NonNull;
 import androidx.annotation.VisibleForTesting;
 import androidx.loader.content.AsyncTaskLoader;

 /**
 * Loader for network data usage history. It returns a list of usage data per billing cycle.
 */
-public class NetworkCycleDataLoader extends AsyncTaskLoader<List<NetworkCycleData>> {
-    private static final String TAG = "CycleDataSummaryLoader";
-    private final NetworkStatsManager mNetworkStatsManager;
-    private final String mSubId;
-    private final int mNetworkType;
+public abstract class NetworkCycleDataLoader<D> extends AsyncTaskLoader<D> {
+    private static final String TAG = "NetworkCycleDataLoader";
+    protected final NetworkStatsManager mNetworkStatsManager;
+    protected final String mSubId;
+    protected final int mNetworkType;
    private final NetworkPolicy mPolicy;
    private final NetworkTemplate mNetworkTemplate;
    @VisibleForTesting
    final INetworkStatsService mNetworkStatsService;

-    private NetworkCycleDataLoader(Builder builder) {
+    protected NetworkCycleDataLoader(Builder<?> builder) {
        super(builder.mContext);
        mPolicy = builder.mPolicy;
        mSubId = builder.mSubId;
@@ -75,21 +71,25 @@ public class NetworkCycleDataLoader extends AsyncTaskLoader<List<NetworkCycleDat
        forceLoad();
    }

-    @Override
-    public List<NetworkCycleData> loadInBackground() {
+    public D loadInBackground() {
        if (mPolicy == null) {
-            return loadFourWeeksData();
+            loadFourWeeksData();
+        } else {
+            loadPolicyData();
        }
-        final List<NetworkCycleData> data = new ArrayList<>();
-        final Iterator<Pair<ZonedDateTime, ZonedDateTime>> iterator = NetworkPolicyManager
-            .cycleIterator(mPolicy);
+        return getCycleUsage();
+    }
+
+    @VisibleForTesting
+    void loadPolicyData() {
+        final Iterator<Pair<ZonedDateTime, ZonedDateTime>> iterator =
+            NetworkPolicyManager.cycleIterator(mPolicy);
        while (iterator.hasNext()) {
            final Pair<ZonedDateTime, ZonedDateTime> cycle = iterator.next();
            final long cycleStart = cycle.first.toInstant().toEpochMilli();
            final long cycleEnd = cycle.second.toInstant().toEpochMilli();
-            getUsage(cycleStart, cycleEnd, data);
+            recordUsage(cycleStart, cycleEnd);
        }
-        return data;
    }

    @Override
@@ -105,8 +105,7 @@ public class NetworkCycleDataLoader extends AsyncTaskLoader<List<NetworkCycleDat
    }

    @VisibleForTesting
-    List<NetworkCycleData> loadFourWeeksData() {
-        final List<NetworkCycleData> data = new ArrayList<>();
+    void loadFourWeeksData() {
        try {
            final INetworkStatsSession networkSession = mNetworkStatsService.openSession();
            final NetworkStatsHistory networkHistory = networkSession.getHistoryForNetwork(
@@ -117,7 +116,7 @@ public class NetworkCycleDataLoader extends AsyncTaskLoader<List<NetworkCycleDat
            long cycleEnd = historyEnd;
            while (cycleEnd > historyStart) {
                final long cycleStart = cycleEnd - (DateUtils.WEEK_IN_MILLIS * 4);
-                getUsage(cycleStart, cycleEnd, data);
+                recordUsage(cycleStart, cycleEnd);
                cycleEnd = cycleStart;
            }

@@ -125,29 +124,23 @@ public class NetworkCycleDataLoader extends AsyncTaskLoader<List<NetworkCycleDat
        } catch (RemoteException e) {
            throw new RuntimeException(e);
        }
-        return data;
    }

    @VisibleForTesting
-    void getUsage(long start, long end, @NonNull List<NetworkCycleData> data) {
-        try {
-            final NetworkStats stats = mNetworkStatsManager.querySummary(
-                mNetworkType, mSubId, start, end);
-            final long total = getTotalUsage(stats);
-            if (total > 0L) {
-                data.add(new NetworkCycleData.Builder()
-                    .setStartTime(start)
-                    .setEndTime(end)
-                    .setTotalUsage(total)
-                    .setUsageBuckets(getUsageBuckets(start, end))
-                    .build());
+    abstract void recordUsage(long start, long end);
+
+    abstract D getCycleUsage();
+
+    public static Builder<?> builder(Context context) {
+        return new Builder<NetworkCycleDataLoader>(context) {
+            @Override
+            public NetworkCycleDataLoader build() {
+                return null;
            }
-        } catch (RemoteException e) {
-            Log.e(TAG, "Exception querying network detail.", e);
-        }
+        };
    }

-    private long getTotalUsage(NetworkStats stats) {
+    protected long getTotalUsage(NetworkStats stats) {
        long bytes = 0L;
        if (stats != null) {
            final NetworkStats.Bucket bucket = new NetworkStats.Bucket();
@@ -159,61 +152,38 @@ public class NetworkCycleDataLoader extends AsyncTaskLoader<List<NetworkCycleDat
        return bytes;
    }

-    private List<NetworkCycleData> getUsageBuckets(long start, long end) {
-        final List<NetworkCycleData> data = new ArrayList<>();
-        long bucketStart = start;
-        long bucketEnd = start + NetworkCycleData.BUCKET_DURATION_MS;
-        while (bucketEnd <= end) {
-            long usage = 0L;
-            try {
-                final NetworkStats stats = mNetworkStatsManager.querySummary(
-                    mNetworkType, mSubId, bucketStart, bucketEnd);
-                usage = getTotalUsage(stats);
-            } catch (RemoteException e) {
-                Log.e(TAG, "Exception querying network detail.", e);
-            }
-            data.add(new NetworkCycleData.Builder()
-                .setStartTime(bucketStart).setEndTime(bucketEnd).setTotalUsage(usage).build());
-            bucketStart = bucketEnd;
-            bucketEnd += NetworkCycleData.BUCKET_DURATION_MS;
-        }
-        return data;
-    }
-
-    public static class Builder {
+    public static abstract class Builder<T extends NetworkCycleDataLoader> {
        private final Context mContext;
        private NetworkPolicy mPolicy;
        private String mSubId;
        private int mNetworkType;
        private NetworkTemplate mNetworkTemplate;

-        public Builder(Context context) {
+        public Builder (Context context) {
            mContext = context;
        }

-        public Builder setNetworkPolicy(NetworkPolicy policy) {
+        public Builder<T> setNetworkPolicy(NetworkPolicy policy) {
            mPolicy = policy;
            return this;
        }

-        public Builder setSubscriberId(String subId) {
+        public Builder<T> setSubscriberId(String subId) {
            mSubId = subId;
            return this;
        }

-        public Builder setNetworkType(int networkType) {
+        public Builder<T> setNetworkType(int networkType) {
            mNetworkType = networkType;
            return this;
        }

-        public Builder setNetworkTemplate(NetworkTemplate template) {
+        public Builder<T> setNetworkTemplate(NetworkTemplate template) {
            mNetworkTemplate = template;
            return this;
        }

-        public NetworkCycleDataLoader build() {
-            return new NetworkCycleDataLoader(this);
-        }
+        public abstract T build();
    }

 }
--- a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/net/NetworkCycleChartDataLoaderTest.java
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/net/NetworkCycleChartDataLoaderTest.java
@@ -0,0 +1,66 @@
+/*
+ * Copyright (C) 2018 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.settingslib.net;
+
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import android.app.usage.NetworkStatsManager;
+import android.content.Context;
+import android.net.ConnectivityManager;
+import android.os.RemoteException;
+import android.text.format.DateUtils;
+
+import com.android.settingslib.SettingsLibRobolectricTestRunner;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+@RunWith(SettingsLibRobolectricTestRunner.class)
+public class NetworkCycleChartDataLoaderTest {
+
+    @Mock
+    private NetworkStatsManager mNetworkStatsManager;
+    @Mock
+    private Context mContext;
+
+    private NetworkCycleChartDataLoader mLoader;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        when(mContext.getSystemService(Context.NETWORK_STATS_SERVICE))
+            .thenReturn(mNetworkStatsManager);
+    }
+
+    @Test
+    public void recordUsage_shouldQueryNetworkSummary() throws RemoteException {
+        final long end = System.currentTimeMillis();
+        final long start = end - (DateUtils.WEEK_IN_MILLIS * 4);
+        final int networkType = ConnectivityManager.TYPE_MOBILE;
+        final String subId = "TestSubscriber";
+        mLoader = NetworkCycleChartDataLoader.builder(mContext)
+            .setNetworkType(networkType).setSubscriberId(subId).build();
+
+        mLoader.recordUsage(start, end);
+
+        verify(mNetworkStatsManager).querySummary(networkType, subId, start, end);
+    }
+}
--- a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/net/NetworkCycleDataForUidLoaderTest.java
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/net/NetworkCycleDataForUidLoaderTest.java
@@ -0,0 +1,66 @@
+/*
+ * Copyright (C) 2018 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.settingslib.net;
+
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import android.app.usage.NetworkStatsManager;
+import android.content.Context;
+import android.net.ConnectivityManager;
+import android.text.format.DateUtils;
+
+import com.android.settingslib.SettingsLibRobolectricTestRunner;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+@RunWith(SettingsLibRobolectricTestRunner.class)
+public class NetworkCycleDataForUidLoaderTest {
+
+    @Mock
+    private NetworkStatsManager mNetworkStatsManager;
+    @Mock
+    private Context mContext;
+
+    private NetworkCycleDataForUidLoader mLoader;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        when(mContext.getSystemService(Context.NETWORK_STATS_SERVICE))
+            .thenReturn(mNetworkStatsManager);
+    }
+
+    @Test
+    public void recordUsage_shouldQueryNetworkDetailsForUid() {
+        final long end = System.currentTimeMillis();
+        final long start = end - (DateUtils.WEEK_IN_MILLIS * 4);
+        final int networkType = ConnectivityManager.TYPE_MOBILE;
+        final String subId = "TestSubscriber";
+        final int uid = 1;
+        mLoader = NetworkCycleDataForUidLoader.builder(mContext)
+            .setUid(uid).setNetworkType(networkType).setSubscriberId(subId).build();
+
+        mLoader.recordUsage(start, end);
+
+        verify(mNetworkStatsManager).queryDetailsForUid(networkType, subId, start, end, uid);
+    }
+}
--- 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,13 +16,10 @@

 package com.android.settingslib.net;

-import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyInt;
-import static org.mockito.Mockito.anyLong;
-import static org.mockito.Mockito.doReturn;
-import static org.mockito.Mockito.eq;
-import static org.mockito.Mockito.mock;
 import static org.mockito.Matchers.nullable;
+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;
@@ -50,6 +47,7 @@ import org.robolectric.util.ReflectionHelpers;

 import java.time.ZonedDateTime;
 import java.util.Iterator;
+import java.util.List;

@RunWith(SettingsLibRobolectricTestRunner.class)
 public class NetworkCycleDataLoaderTest {
@@ -64,8 +62,11 @@ public class NetworkCycleDataLoaderTest {
    private Iterator<Range<ZonedDateTime>> mIterator;
    @Mock
    private INetworkStatsService mNetworkStatsService;
+    @Mock
+    private NetworkCycleDataLoader.Builder mBuilder;

-    private NetworkCycleDataLoader mLoader;
+
+    private NetworkCycleDataTestLoader mLoader;

    @Before
    public void setUp() {
@@ -77,8 +78,8 @@ public class NetworkCycleDataLoaderTest {

    @Test
    public void loadInBackground_noNetworkPolicy_shouldLoad4WeeksData() {
-        mLoader = spy(new NetworkCycleDataLoader.Builder(mContext).build());
-        doReturn(null).when(mLoader).loadFourWeeksData();
+        mLoader = spy(new NetworkCycleDataTestLoader(mContext));
+        doNothing().when(mLoader).loadFourWeeksData();

        mLoader.loadInBackground();

@@ -86,31 +87,45 @@ public class NetworkCycleDataLoaderTest {
    }

    @Test
-    public void loadInBackground_shouldQueryNetworkSummary() throws RemoteException {
+    public void loadInBackground_hasNetworkPolicy_shouldLoadPolicyData() {
+        mLoader = spy(new NetworkCycleDataTestLoader(mContext));
+        ReflectionHelpers.setField(mLoader, "mPolicy", mPolicy);
+
+        mLoader.loadInBackground();
+
+        verify(mLoader).loadPolicyData();
+    }
+
+    @Test
+    public void loadPolicyData_shouldRecordUsageFromPolicyCycle() {
        final int networkType = ConnectivityManager.TYPE_MOBILE;
        final String subId = "TestSubscriber";
        final ZonedDateTime now = ZonedDateTime.now();
        final Range<ZonedDateTime> cycle = new Range<>(now, now);
+        final long nowInMs = now.toInstant().toEpochMilli();
        // mock 1 cycle data.
        // hasNext() will be called internally in next(), hence setting it to return true twice.
        when(mIterator.hasNext()).thenReturn(true).thenReturn(true).thenReturn(false);
        when(mIterator.next()).thenReturn(cycle);
-        mLoader = new NetworkCycleDataLoader.Builder(mContext)
-            .setNetworkPolicy(mPolicy).setNetworkType(networkType).setSubscriberId(subId).build();
+        mLoader = spy(new NetworkCycleDataTestLoader(mContext));
+        ReflectionHelpers.setField(mLoader, "mPolicy", mPolicy);
+        ReflectionHelpers.setField(mLoader, "mNetworkType", networkType);
+        ReflectionHelpers.setField(mLoader, "mSubId", subId);

-        mLoader.loadInBackground();
+        mLoader.loadPolicyData();

-        verify(mNetworkStatsManager).querySummary(eq(networkType), eq(subId), anyLong(), anyLong());
+        verify(mLoader).recordUsage(nowInMs, nowInMs);
    }

    @Test
-    public void loadFourWeeksData_shouldGetUsageForLast4Weeks() throws RemoteException {
-        mLoader = spy(new NetworkCycleDataLoader.Builder(mContext).build());
+    public void loadFourWeeksData_shouldRecordUsageForLast4Weeks() throws RemoteException {
+        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);
+        when(networkSession.getHistoryForNetwork(nullable(NetworkTemplate.class), anyInt()))
+            .thenReturn(networkHistory);
        final long now = System.currentTimeMillis();
        final long fourWeeksAgo = now - (DateUtils.WEEK_IN_MILLIS * 4);
        when(networkHistory.getStart()).thenReturn(fourWeeksAgo);
@@ -118,6 +133,23 @@ public class NetworkCycleDataLoaderTest {

        mLoader.loadFourWeeksData();

-        verify(mLoader).getUsage(eq(fourWeeksAgo), eq(now), any());
+        verify(mLoader).recordUsage(fourWeeksAgo, now);
+    }
+
+    public class NetworkCycleDataTestLoader extends NetworkCycleDataLoader<List<NetworkCycleData>> {
+
+        private NetworkCycleDataTestLoader(Context context) {
+            super(NetworkCycleDataLoader.builder(mContext));
+            mContext = context;
+        }
+
+        @Override
+        void recordUsage(long start, long end) {
+        }
+
+        @Override
+        List<NetworkCycleData> getCycleUsage() {
+            return null;
+        }
    }
 }