diff --git a/tests/componenttests/Android.bp b/tests/componenttests/Android.bp
index 77ee164..5c03aa9 100644
--- a/tests/componenttests/Android.bp
+++ b/tests/componenttests/Android.bp
@@ -26,6 +26,7 @@
         "androidx.test.runner",
         "androidx.test.rules",
         "androidx.test.ext.junit",
+        "mockito-target",
     ],
 
     test_suites: ["device-tests"],
diff --git a/tests/componenttests/src/com/android/settings/network/telephony/MobileDataPreferenceControllerComponentTest.java b/tests/componenttests/src/com/android/settings/network/telephony/MobileDataPreferenceControllerComponentTest.java
index 78a2c92..df3dd2e 100644
--- a/tests/componenttests/src/com/android/settings/network/telephony/MobileDataPreferenceControllerComponentTest.java
+++ b/tests/componenttests/src/com/android/settings/network/telephony/MobileDataPreferenceControllerComponentTest.java
@@ -20,6 +20,8 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.Mockito.mock;
+
 import android.app.Instrumentation;
 import android.content.Context;
 import android.content.Intent;
@@ -31,6 +33,7 @@
 
 import androidx.fragment.app.FragmentActivity;
 import androidx.fragment.app.FragmentManager;
+import androidx.lifecycle.LifecycleOwner;
 import androidx.test.core.app.ActivityScenario;
 import androidx.test.ext.junit.rules.ActivityScenarioRule;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
@@ -40,6 +43,9 @@
 import com.android.settings.Settings;
 import com.android.settings.testutils.CommonUtils;
 import com.android.settings.testutils.UiUtils;
+import com.android.settingslib.core.lifecycle.Lifecycle;
+import com.android.settingslib.mobile.dataservice.MobileNetworkInfoEntity;
+import com.android.settingslib.mobile.dataservice.SubscriptionInfoEntity;
 
 import org.junit.After;
 import org.junit.Assume;
@@ -47,6 +53,8 @@
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.Mockito;
 
 import java.io.IOException;
 import java.net.URL;
@@ -54,6 +62,12 @@
 @RunWith(AndroidJUnit4.class)
 @SmallTest
 public class MobileDataPreferenceControllerComponentTest {
+
+    @Mock
+    private Lifecycle mLifecycle;
+    @Mock
+    private LifecycleOwner mLifecycleOwner;
+
     public static final int TIMEOUT = 2000;
     private static int sSubscriptionId = 2;
     public final String TAG = this.getClass().getName();
@@ -116,9 +130,11 @@
             try {
                 URL url = new URL("https://www.google.net");
                 MobileDataPreferenceController controller = new MobileDataPreferenceController(
-                        mInstrumentation.getTargetContext(), "mobile_data");
+                        mInstrumentation.getTargetContext(), "mobile_data", mLifecycle,
+                        mLifecycleOwner, sSubscriptionId);
                 FragmentManager manager = ((FragmentActivity) activity).getSupportFragmentManager();
-                controller.init(manager, sSubscriptionId);
+                controller.init(manager, sSubscriptionId, mock(SubscriptionInfoEntity.class), mock(
+                        MobileNetworkInfoEntity.class));
 
                 // Make sure mobile network can connect at first.
                 assertThat(UiUtils.waitUntilCondition(1000,
diff --git a/tests/unit/src/com/android/settings/network/telephony/MobileDataPreferenceControllerTest.java b/tests/unit/src/com/android/settings/network/telephony/MobileDataPreferenceControllerTest.java
index 45f3693..16e3963 100644
--- a/tests/unit/src/com/android/settings/network/telephony/MobileDataPreferenceControllerTest.java
+++ b/tests/unit/src/com/android/settings/network/telephony/MobileDataPreferenceControllerTest.java
@@ -28,18 +28,24 @@
 
 import android.app.Instrumentation;
 import android.content.Context;
+import android.os.Looper;
 import android.telephony.SubscriptionInfo;
 import android.telephony.SubscriptionManager;
 import android.telephony.TelephonyManager;
 
 import androidx.fragment.app.FragmentManager;
 import androidx.fragment.app.FragmentTransaction;
+import androidx.lifecycle.LifecycleOwner;
+import androidx.lifecycle.LifecycleRegistry;
 import androidx.preference.SwitchPreference;
 import androidx.test.core.app.ApplicationProvider;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.platform.app.InstrumentationRegistry;
 
 import com.android.settings.testutils.ResourcesUtils;
+import com.android.settingslib.core.lifecycle.Lifecycle;
+import com.android.settingslib.mobile.dataservice.MobileNetworkInfoEntity;
+import com.android.settingslib.mobile.dataservice.SubscriptionInfoEntity;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -49,6 +55,10 @@
 
 @RunWith(AndroidJUnit4.class)
 public class MobileDataPreferenceControllerTest {
+    private static final String SUB_ID_1 = "1";
+    private static final String SUB_ID_2 = "2";
+    private static final String DISPLAY_NAME_1 = "Sub 1";
+    private static final String DISPLAY_NAME_2 = "Sub 2";
     private static final int SUB_ID = 2;
     private static final int SUB_ID_OTHER = 3;
 
@@ -64,7 +74,15 @@
     private SubscriptionInfo mSubscriptionInfo;
     @Mock
     private FragmentTransaction mFragmentTransaction;
-
+    @Mock
+    private Lifecycle mLifecycle;
+    @Mock
+    private LifecycleOwner mLifecycleOwner;
+    private SubscriptionInfoEntity mSubInfo1;
+    private SubscriptionInfoEntity mSubInfo2;
+    private MobileNetworkInfoEntity mNetworkInfo1;
+    private MobileNetworkInfoEntity mNetworkInfo2;
+    private LifecycleRegistry mLifecycleRegistry;
     private MobileDataPreferenceController mController;
     private SwitchPreference mPreference;
     private Context mContext;
@@ -73,6 +91,10 @@
     public void setUp() {
         MockitoAnnotations.initMocks(this);
 
+        if (Looper.myLooper() == null) {
+            Looper.prepare();
+        }
+
         mContext = spy(ApplicationProvider.getApplicationContext());
         doReturn(mTelephonyManager).when(mContext).getSystemService(Context.TELEPHONY_SERVICE);
 
@@ -83,22 +105,45 @@
         doReturn(mFragmentTransaction).when(mFragmentManager).beginTransaction();
 
         mPreference = new SwitchPreference(mContext);
-        mController = new MobileDataPreferenceController(mContext, "mobile_data");
-        mController.init(mFragmentManager, SUB_ID);
+        mController = new MobileDataPreferenceController(mContext, "mobile_data", mLifecycle,
+                mLifecycleOwner, SUB_ID);
+        mController.init(mFragmentManager, SUB_ID, mSubInfo1, mNetworkInfo1);
         mPreference.setKey(mController.getPreferenceKey());
+        mLifecycleRegistry = new LifecycleRegistry(mLifecycleOwner);
+        when(mLifecycleOwner.getLifecycle()).thenReturn(mLifecycleRegistry);
+    }
+
+    private SubscriptionInfoEntity setupSubscriptionInfoEntity(String subId, String displayName,
+            boolean isOpportunistic, boolean isValid, boolean isActive, boolean isAvailable,
+            boolean isDefaultData) {
+        int id = Integer.parseInt(subId);
+        return new SubscriptionInfoEntity(subId, id, id,
+                displayName, displayName, 0, "mcc", "mnc", "countryIso", false, id,
+                TelephonyManager.DEFAULT_PORT_INDEX, isOpportunistic, null,
+                SubscriptionManager.SUBSCRIPTION_TYPE_LOCAL_SIM, displayName, false,
+                "1234567890", true, "default", false, isValid, true, isActive, isAvailable, false,
+                false, isDefaultData, false, false);
+    }
+
+    private MobileNetworkInfoEntity setupMobileNetworkInfoEntity(String subId,
+            boolean isDatEnabled) {
+        return new MobileNetworkInfoEntity(subId, false, false, isDatEnabled, false, false, false,
+                false, false, false, false, false);
     }
 
     @Test
     public void getAvailabilityStatus_invalidSubscription_returnAvailableUnsearchable() {
-        mController.init(mFragmentManager, SubscriptionManager.INVALID_SUBSCRIPTION_ID);
+        mController.init(mFragmentManager, SubscriptionManager.INVALID_SUBSCRIPTION_ID, mSubInfo1,
+                mNetworkInfo1);
 
         assertThat(mController.getAvailabilityStatus()).isEqualTo(AVAILABLE_UNSEARCHABLE);
     }
 
     @Test
     public void isDialogNeeded_disableSingleSim_returnFalse() {
-        doReturn(true).when(mTelephonyManager).isDataEnabled();
-        doReturn(mSubscriptionInfo).when(mSubscriptionManager).getActiveSubscriptionInfo(SUB_ID);
+        mSubInfo1 = setupSubscriptionInfoEntity(SUB_ID_1, DISPLAY_NAME_1, false, true, true, true,
+                true);
+        mNetworkInfo1 = setupMobileNetworkInfoEntity(String.valueOf(SUB_ID), true);
         doReturn(1).when(mTelephonyManager).getActiveModemCount();
 
         assertThat(mController.isDialogNeeded()).isFalse();
@@ -106,12 +151,15 @@
 
     @Test
     public void isDialogNeeded_enableNonDefaultSimInMultiSimMode_returnTrue() {
-        doReturn(false).when(mTelephonyManager).isDataEnabled();
-        doReturn(mSubscriptionInfo).when(mSubscriptionManager)
-                .getActiveSubscriptionInfo(SUB_ID);
+        mSubInfo1 = setupSubscriptionInfoEntity(SUB_ID_1, DISPLAY_NAME_1, false, true, true, true,
+                false);
+        mNetworkInfo1 = setupMobileNetworkInfoEntity(String.valueOf(SUB_ID), false);
+        doReturn(1).when(mTelephonyManager).getActiveModemCount();
         // Ideally, it would be better if we could set the default data subscription to
         // SUB_ID_OTHER, and set that as an active subscription id.
-        when(mSubscriptionManager.isActiveSubscriptionId(anyInt())).thenReturn(true);
+        mSubInfo2 = setupSubscriptionInfoEntity(SUB_ID_2, DISPLAY_NAME_2, false, true, true, true,
+                true);
+        mNetworkInfo1 = setupMobileNetworkInfoEntity(String.valueOf(SUB_ID), true);
         doReturn(2).when(mTelephonyManager).getActiveModemCount();
 
         assertThat(mController.isDialogNeeded()).isTrue();
@@ -132,8 +180,11 @@
 
     @Test
     public void onPreferenceChange_singleSim_On_shouldEnableData() {
-        doReturn(true).when(mTelephonyManager).isDataEnabled();
-        doReturn(mSubscriptionInfo).when(mSubscriptionManager).getActiveSubscriptionInfo(SUB_ID);
+        mSubInfo1 = setupSubscriptionInfoEntity(SUB_ID_1, DISPLAY_NAME_1, true, true, true, true,
+                true);
+        mNetworkInfo1 = setupMobileNetworkInfoEntity(String.valueOf(SUB_ID), true);
+        mController.setSubscriptionInfoEntity(mSubInfo1);
+        mController.setMobileNetworkInfoEntity(mNetworkInfo1);
         doReturn(1).when(mTelephonyManager).getActiveModemCount();
 
         mController.onPreferenceChange(mPreference, true);
@@ -143,8 +194,11 @@
 
     @Test
     public void onPreferenceChange_multiSim_On_shouldEnableData() {
-        doReturn(true).when(mTelephonyManager).isDataEnabled();
-        doReturn(mSubscriptionInfo).when(mSubscriptionManager).getActiveSubscriptionInfo(SUB_ID);
+        mSubInfo1 = setupSubscriptionInfoEntity(SUB_ID_1, DISPLAY_NAME_1, true, true, true, true,
+                true);
+        mNetworkInfo1 = setupMobileNetworkInfoEntity(String.valueOf(SUB_ID), true);
+        mController.setSubscriptionInfoEntity(mSubInfo1);
+        mController.setMobileNetworkInfoEntity(mNetworkInfo1);
         doReturn(2).when(mTelephonyManager).getActiveModemCount();
 
         mController.onPreferenceChange(mPreference, true);
@@ -154,18 +208,20 @@
 
     @Test
     public void isChecked_returnUserDataEnabled() {
-        mController.init(mFragmentManager, SUB_ID);
+        mNetworkInfo1 = setupMobileNetworkInfoEntity(String.valueOf(SUB_ID), false);
+        mController.init(mFragmentManager, SUB_ID, mSubInfo1, mNetworkInfo1);
         assertThat(mController.isChecked()).isFalse();
 
-        doReturn(true).when(mTelephonyManager).isDataEnabled();
+        mNetworkInfo1 = setupMobileNetworkInfoEntity(String.valueOf(SUB_ID), true);
+        mController.setMobileNetworkInfoEntity(mNetworkInfo1);
         assertThat(mController.isChecked()).isTrue();
     }
 
     @Test
     public void updateState_opportunistic_disabled() {
-        doReturn(mSubscriptionInfo).when(mSubscriptionManager).getActiveSubscriptionInfo(SUB_ID);
-        mController.init(mFragmentManager, SUB_ID);
-        doReturn(true).when(mSubscriptionInfo).isOpportunistic();
+        mSubInfo1 = setupSubscriptionInfoEntity(SUB_ID_1, DISPLAY_NAME_1, true, true, true, true,
+                true);
+        mController.init(mFragmentManager, SUB_ID, mSubInfo1, mNetworkInfo1);
         mController.updateState(mPreference);
 
         assertThat(mPreference.isEnabled()).isFalse();
@@ -176,9 +232,9 @@
 
     @Test
     public void updateState_notOpportunistic_enabled() {
-        doReturn(mSubscriptionInfo).when(mSubscriptionManager).getActiveSubscriptionInfo(SUB_ID);
-        mController.init(mFragmentManager, SUB_ID);
-        doReturn(false).when(mSubscriptionInfo).isOpportunistic();
+        mSubInfo1 = setupSubscriptionInfoEntity(SUB_ID_1, DISPLAY_NAME_1, false, true, true, true,
+                true);
+        mController.init(mFragmentManager, SUB_ID, mSubInfo1, mNetworkInfo1);
         mController.updateState(mPreference);
 
         assertThat(mPreference.isEnabled()).isTrue();