Get data type from service state as well

Bug: 26276354
Bug: 26413831
Change-Id: Ic080300db7d5e5200caf611cc2ee1ef3af93736e
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/MobileSignalController.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/MobileSignalController.java
index b65bf43..ad8e3bd 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/MobileSignalController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/MobileSignalController.java
@@ -440,6 +440,7 @@
                         + " dataState=" + state.getDataRegState());
             }
             mServiceState = state;
+            mDataNetType = state.getDataNetworkType();
             updateTelephony();
         }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/NetworkControllerBaseTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/NetworkControllerBaseTest.java
index 30c08cd..61d4e4a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/NetworkControllerBaseTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/NetworkControllerBaseTest.java
@@ -59,8 +59,8 @@
     protected NetworkControllerImpl mNetworkController;
     protected MobileSignalController mMobileSignalController;
     protected PhoneStateListener mPhoneStateListener;
-    private SignalStrength mSignalStrength;
-    private ServiceState mServiceState;
+    protected SignalStrength mSignalStrength;
+    protected ServiceState mServiceState;
     protected ConnectivityManager mMockCm;
     protected WifiManager mMockWm;
     protected SubscriptionManager mMockSm;
@@ -234,7 +234,7 @@
         mPhoneStateListener.onSignalStrengthsChanged(mSignalStrength);
     }
 
-    private void updateServiceState() {
+    protected void updateServiceState() {
         Log.d(TAG, "Sending Service State: " + mServiceState);
         mPhoneStateListener.onServiceStateChanged(mServiceState);
     }
@@ -245,6 +245,7 @@
     }
 
     public void updateDataConnectionState(int dataState, int dataNetType) {
+        when(mServiceState.getDataNetworkType()).thenReturn(dataNetType);
         mPhoneStateListener.onDataConnectionStateChanged(dataState, dataNetType);
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/NetworkControllerDataTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/NetworkControllerDataTest.java
index 0ec8802..3e0da3c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/NetworkControllerDataTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/NetworkControllerDataTest.java
@@ -114,6 +114,21 @@
                 TelephonyIcons.QS_DATA_4G);
     }
 
+    public void testDataChangeWithoutConnectionState() {
+        setupDefaultSignal();
+        updateDataConnectionState(TelephonyManager.DATA_CONNECTED,
+                TelephonyManager.NETWORK_TYPE_LTE);
+
+        verifyDataIndicators(TelephonyIcons.DATA_LTE[1][0 /* No direction */],
+                TelephonyIcons.QS_DATA_LTE);
+
+        Mockito.when(mServiceState.getDataNetworkType())
+                .thenReturn(TelephonyManager.NETWORK_TYPE_HSPA);
+        updateServiceState();
+        verifyDataIndicators(TelephonyIcons.DATA_H[1][0 /* No direction */],
+                TelephonyIcons.QS_DATA_H);
+    }
+
     public void testDataActivity() {
         setupDefaultSignal();