Add explicit subscription id support to DataUsageController

The existing code in DataUsageController only returns usage stats for
the default mobile data subscription. But for dual-sim, dual-standby, we
want to be able to get usage stats for other subscriptions as well, even
when they aren't currently set to be the default. So this CL adds a way
to request stats for an explicit subscription id.

Bug: 122670283
Test: make RunSettingsLibRoboTests
Change-Id: Ieb7eef512bf91292b3fb9f37b9bf589a98faa3a8

2 files changed, 60 insertions, 7 deletions

diff --git a/packages/SettingsLib/src/com/android/settingslib/net/DataUsageController.java b/packages/SettingsLib/src/com/android/settingslib/net/DataUsageController.java
index 180b77e7503b..8cb252ed6283 100644
--- a/packages/SettingsLib/src/com/android/settingslib/net/DataUsageController.java
+++ b/packages/SettingsLib/src/com/android/settingslib/net/DataUsageController.java
@@ -66,6 +66,7 @@ public class DataUsageController {
     private INetworkStatsSession mSession;
     private Callback mCallback;
     private NetworkNameProvider mNetworkController;
+    private int mSubscriptionId;
 
     public DataUsageController(Context context) {
         mContext = context;
@@ -75,6 +76,7 @@ public class DataUsageController {
                 ServiceManager.getService(Context.NETWORK_STATS_SERVICE));
         mPolicyManager = NetworkPolicyManager.from(mContext);
         mNetworkStatsManager = context.getSystemService(NetworkStatsManager.class);
+        mSubscriptionId = SubscriptionManager.INVALID_SUBSCRIPTION_ID;
     }
 
     public void setNetworkController(NetworkNameProvider networkController) {
@@ -82,6 +84,15 @@ public class DataUsageController {
     }
 
     /**
+     * By default this class will just get data usage information for the default data subscription,
+     * but this method can be called to require it to use an explicit subscription id which may be
+     * different from the default one (this is useful for the case of multi-SIM devices).
+     */
+    public void setSubscriptionId(int subscriptionId) {
+        mSubscriptionId = subscriptionId;
+    }
+
+    /**
      * Returns the default warning level in bytes.
      */
     public long getDefaultWarningLevel() {
@@ -99,7 +110,7 @@ public class DataUsageController {
     }
 
     public DataUsageInfo getDataUsageInfo() {
-        final String subscriberId = getActiveSubscriberId(mContext);
+        final String subscriberId = getActiveSubscriberId();
         if (subscriberId == null) {
             return warn("no subscriber id");
         }
@@ -164,7 +175,8 @@ public class DataUsageController {
     private long getUsageLevel(NetworkTemplate template, long start, long end) {
         try {
             final Bucket bucket = mNetworkStatsManager.querySummaryForDevice(
-                getNetworkType(template), getActiveSubscriberId(mContext), start, end);
+                    getNetworkType(template), getActiveSubscriberId(),
+                    start, end);
             if (bucket != null) {
                 return bucket.getRxBytes() + bucket.getTxBytes();
             }
@@ -237,10 +249,13 @@ public class DataUsageController {
         }
     }
 
-    private static String getActiveSubscriberId(Context context) {
-        final TelephonyManager tele = TelephonyManager.from(context);
-        final String actualSubscriberId = tele.getSubscriberId(
-                SubscriptionManager.getDefaultDataSubscriptionId());
+    private String getActiveSubscriberId() {
+        final TelephonyManager tele = TelephonyManager.from(mContext);
+        int subscriptionId = mSubscriptionId;
+        if (subscriptionId == SubscriptionManager.INVALID_SUBSCRIPTION_ID) {
+            subscriptionId = SubscriptionManager.getDefaultDataSubscriptionId();
+        }
+        final String actualSubscriberId = tele.getSubscriberId(subscriptionId);
         return actualSubscriberId;
     }
 
diff --git a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/net/DataUsageControllerTest.java b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/net/DataUsageControllerTest.java
index acf99a24e96c..220463b1f06c 100644
--- a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/net/DataUsageControllerTest.java
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/net/DataUsageControllerTest.java
@@ -45,6 +45,7 @@ import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.robolectric.RobolectricTestRunner;
+import org.robolectric.shadows.ShadowSubscriptionManager;
 
 @RunWith(RobolectricTestRunner.class)
 public class DataUsageControllerTest {
@@ -73,7 +74,9 @@ public class DataUsageControllerTest {
                 new NetworkStatsHistory(DateUtils.DAY_IN_MILLIS /* bucketDuration */));
         doReturn(mNetworkStatsHistory)
                 .when(mSession).getHistoryForNetwork(any(NetworkTemplate.class), anyInt());
-        doReturn(SUB_ID).when(mTelephonyManager).getSubscriberId(anyInt());
+        final int defaultSubscriptionId = 1234;
+        ShadowSubscriptionManager.setDefaultDataSubscriptionId(defaultSubscriptionId);
+        doReturn(SUB_ID).when(mTelephonyManager).getSubscriberId(eq(defaultSubscriptionId));
     }
 
     @Test
@@ -107,4 +110,39 @@ public class DataUsageControllerTest {
         assertThat(mController.getHistoricalUsageLevel(NetworkTemplate.buildTemplateWifiWildcard()))
                 .isEqualTo(receivedBytes + transmittedBytes);
     }
+
+    @Test
+    public void getDataUsageInfo_hasUsageData_shouldReturnCorrectUsageForExplicitSubId()
+            throws Exception {
+        // First setup a stats bucket for the default subscription / subscriber ID.
+        final long defaultSubRx = 1234567L;
+        final long defaultSubTx = 123456L;
+        final NetworkStats.Bucket defaultSubscriberBucket = mock(NetworkStats.Bucket.class);
+        when(defaultSubscriberBucket.getRxBytes()).thenReturn(defaultSubRx);
+        when(defaultSubscriberBucket.getTxBytes()).thenReturn(defaultSubTx);
+        when(mNetworkStatsManager.querySummaryForDevice(eq(ConnectivityManager.TYPE_MOBILE),
+                eq(SUB_ID), eq(0L)/* startTime */, anyLong() /* endTime */)).thenReturn(
+                defaultSubscriberBucket);
+
+        // Now setup a stats bucket for a different, non-default subscription / subscriber ID.
+        final long nonDefaultSubRx = 7654321L;
+        final long nonDefaultSubTx = 654321L;
+        final NetworkStats.Bucket nonDefaultSubscriberBucket = mock(NetworkStats.Bucket.class);
+        when(nonDefaultSubscriberBucket.getRxBytes()).thenReturn(nonDefaultSubRx);
+        when(nonDefaultSubscriberBucket.getTxBytes()).thenReturn(nonDefaultSubTx);
+        final int explictSubscriptionId = 55;
+        final String subscriberId2 = "Test Subscriber 2";
+        when(mNetworkStatsManager.querySummaryForDevice(eq(ConnectivityManager.TYPE_MOBILE),
+                eq(subscriberId2), eq(0L)/* startTime */, anyLong() /* endTime */)).thenReturn(
+                nonDefaultSubscriberBucket);
+        doReturn(subscriberId2).when(mTelephonyManager).getSubscriberId(explictSubscriptionId);
+
+        // Now verify that when we're asking for stats on the non-default subscription, we get
+        // the data back for that subscription and *not* the default one.
+        mController.setSubscriptionId(explictSubscriptionId);
+
+        assertThat(mController.getHistoricalUsageLevel(
+                NetworkTemplate.buildTemplateMobileAll(subscriberId2))).isEqualTo(
+                nonDefaultSubRx + nonDefaultSubTx);
+    }
 }