wifi data usage: Query wifi usage per configuration

- Updated to use the AllNetworkKeys set instead of NetworkId.

- Fix crash issue when mDataUsageController is not created.

- See b/126299427#comment37 for a screenshot.

Bug: 126299427
Test: manual test
make RunSettingsRoboTests
ROBOTEST_FILTER=DataUsageSummaryPreferenceControllerTest
make RunSettingsRoboTests
ROBOTEST_FILTER=WifiDetailPreferenceController2Test
make RunSettingsRoboTests
ROBOTEST_FILTER=WifiDataUsageSummaryPreferenceControllerTest

Change-Id: I26289fde1f18d97631448edc71e034d33b865cc4

src/com/android/settings/datausage/DataUsageSummaryPreferenceController.java

@@ -123,13 +123,13 @@ public class DataUsageSummaryPreferenceController extends TelephonyBasePreferenc
         mDataUsageController = null;
     }
 
-    private void updateConfiguration(Context context,
+    protected void updateConfiguration(Context context,
             int subscriptionId, SubscriptionInfo subInfo) {
         final NetworkPolicyManager policyManager =
                 context.getSystemService(NetworkPolicyManager.class);
         mPolicyEditor = new NetworkPolicyEditor(policyManager);
 
-        mDataUsageController = new DataUsageController(context);
+        mDataUsageController = createDataUsageController(context);
         mDataUsageController.setSubscriptionId(subscriptionId);
         mDataInfoController = new DataUsageInfoController();
 
@@ -145,6 +145,11 @@ public class DataUsageSummaryPreferenceController extends TelephonyBasePreferenc
         }
     }
 
+    @VisibleForTesting
+    DataUsageController createDataUsageController(Context context) {
+        return new DataUsageController(context);
+    }
+
     @VisibleForTesting
     DataUsageSummaryPreferenceController(
             DataUsageController dataUsageController,
@@ -186,8 +191,7 @@ public class DataUsageSummaryPreferenceController extends TelephonyBasePreferenc
                 .getSubscriptionPlans(subscriptionId);
     }
 
-    @VisibleForTesting
-    SubscriptionInfo getSubscriptionInfo(int subscriptionId) {
+    protected SubscriptionInfo getSubscriptionInfo(int subscriptionId) {
         if (!mHasMobileData) {
             return null;
         }

src/com/android/settings/datausage/WifiDataUsageSummaryPreferenceController.java

@@ -26,23 +26,19 @@ import androidx.preference.PreferenceFragmentCompat;
 import com.android.settingslib.core.lifecycle.Lifecycle;
 import com.android.settingslib.net.DataUsageController;
 
+import java.util.HashSet;
 import java.util.Set;
 
 /**
  * The controller displays a data usage chart for the specified Wi-Fi network.
  */
 public class WifiDataUsageSummaryPreferenceController extends DataUsageSummaryPreferenceController {
-    final String mNetworkId;
+    final Set<String> mAllNetworkKeys;
 
-    public WifiDataUsageSummaryPreferenceController(Activity activity,
-            Lifecycle lifecycle, PreferenceFragmentCompat fragment, CharSequence networkId) {
+    public WifiDataUsageSummaryPreferenceController(Activity activity, Lifecycle lifecycle,
+            PreferenceFragmentCompat fragment, Set<String> allNetworkKeys) {
         super(activity, lifecycle, fragment, SubscriptionManager.INVALID_SUBSCRIPTION_ID);
-
-        if (networkId == null) {
-            mNetworkId = null;
-        } else {
-            mNetworkId = String.valueOf(networkId);
-        }
+        mAllNetworkKeys = new HashSet<>(allNetworkKeys);
     }
 
     @Override
@@ -52,10 +48,11 @@ public class WifiDataUsageSummaryPreferenceController extends DataUsageSummaryPr
         }
 
         final DataUsageSummaryPreference mPreference = (DataUsageSummaryPreference) preference;
-        // TODO(b/126299427): Currently gets data usage of whole Wi-Fi networks, but should get
-        //  specified one.
         final NetworkTemplate template = new NetworkTemplate.Builder(NetworkTemplate.MATCH_WIFI)
-                .setWifiNetworkKeys(Set.of(mNetworkId)).build();
+                .setWifiNetworkKeys(mAllNetworkKeys).build();
+        if (mDataUsageController == null) {
+            updateConfiguration(mContext, mSubId, getSubscriptionInfo(mSubId));
+        }
         final DataUsageController.DataUsageInfo info = mDataUsageController.getDataUsageInfo(
                 template);
         mDataInfoController.updateDataLimit(info, mPolicyEditor.getPolicy(template));

src/com/android/settings/wifi/details2/WifiDetailPreferenceController2.java

@@ -435,9 +435,9 @@ public class WifiDetailPreferenceController2 extends AbstractPreferenceControlle
             mDataUsageSummaryPref = screen.findPreference(KEY_DATA_USAGE_HEADER);
             mDataUsageSummaryPref.setVisible(true);
             mSummaryHeaderController =
-                new WifiDataUsageSummaryPreferenceController(mFragment.getActivity(),
-                        mLifecycle, (PreferenceFragmentCompat) mFragment,
-                        mWifiEntry.getTitle());
+                    new WifiDataUsageSummaryPreferenceController(mFragment.getActivity(),
+                            mLifecycle, (PreferenceFragmentCompat) mFragment,
+                            mWifiEntry.getWifiConfiguration().getAllNetworkKeys());
             return;
         }
 

tests/robotests/src/com/android/settings/datausage/WifiDataUsageSummaryPreferenceControllerTest.java

@@ -0,0 +1,94 @@
+/*
+ * Copyright (C) 2022 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.settings.datausage;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
+
+import android.content.Context;
+import android.net.NetworkPolicyManager;
+import android.telephony.TelephonyManager;
+
+import androidx.fragment.app.FragmentActivity;
+import androidx.test.core.app.ApplicationProvider;
+
+import com.android.settingslib.core.lifecycle.Lifecycle;
+import com.android.settingslib.net.DataUsageController;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.Spy;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
+import org.robolectric.RobolectricTestRunner;
+
+import java.util.HashSet;
+import java.util.Set;
+
+@RunWith(RobolectricTestRunner.class)
+public class WifiDataUsageSummaryPreferenceControllerTest {
+
+    @Rule
+    public final MockitoRule mMockitoRule = MockitoJUnit.rule();
+    @Spy
+    Context mContext = ApplicationProvider.getApplicationContext();
+    @Mock
+    FragmentActivity mActivity;
+    @Mock
+    Lifecycle mLifecycle;
+    @Mock
+    TelephonyManager mTelephonyManager;
+    @Mock
+    NetworkPolicyManager mNetworkPolicyManager;
+    @Mock
+    DataUsageSummaryPreference mSummaryPreference;
+    @Mock
+    DataUsageController mDataUsageController;
+    @Mock
+    DataUsageController.DataUsageInfo mDataUsageInfo;
+
+    WifiDataUsageSummaryPreferenceController mController;
+    Set<String> mAllNetworkKeys = new HashSet<>();
+
+    @Before
+    public void setUp() {
+        doReturn(mContext.getResources()).when(mActivity).getResources();
+        doReturn(mTelephonyManager).when(mActivity).getSystemService(TelephonyManager.class);
+        doReturn(mNetworkPolicyManager).when(mActivity)
+                .getSystemService(NetworkPolicyManager.class);
+        doNothing().when(mSummaryPreference).setWifiMode(anyBoolean(), anyString(), anyBoolean());
+        doReturn(mDataUsageInfo).when(mDataUsageController).getDataUsageInfo(any());
+
+        mController = spy(new WifiDataUsageSummaryPreferenceController(mActivity, mLifecycle, null,
+                mAllNetworkKeys));
+        doReturn(mDataUsageController).when(mController).createDataUsageController(any());
+    }
+
+    @Test
+    public void updateState_nullOfDataUsageController_shouldNotCrash() {
+        mController.mDataUsageController = null;
+
+        mController.updateState(mSummaryPreference);
+    }
+}