diff --git a/tests/robotests/src/com/android/settings/applications/appinfo/AppDataUsagePreferenceControllerTest.java b/tests/robotests/src/com/android/settings/applications/appinfo/AppDataUsagePreferenceControllerTest.java
index e77106d2a01..1adc6781aec 100644
--- a/tests/robotests/src/com/android/settings/applications/appinfo/AppDataUsagePreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/applications/appinfo/AppDataUsagePreferenceControllerTest.java
@@ -16,6 +16,8 @@
 
 package com.android.settings.applications.appinfo;
 
+import static com.android.settings.core.FeatureFlags.DATA_USAGE_V2;
+
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
@@ -33,6 +35,7 @@ import android.content.pm.ApplicationInfo;
 import android.net.ConnectivityManager;
 import android.net.INetworkStatsSession;
 import android.os.Bundle;
+import android.util.FeatureFlagUtils;
 
 import androidx.loader.app.LoaderManager;
 import androidx.preference.Preference;
@@ -65,6 +68,7 @@ public class AppDataUsagePreferenceControllerTest {
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         mContext = spy(RuntimeEnvironment.application.getApplicationContext());
+        FeatureFlagUtils.setEnabled(mContext, DATA_USAGE_V2, false);
         mController = spy(new AppDataUsagePreferenceController(mContext, "test_key"));
         mController.setParentFragment(mFragment);
     }
diff --git a/tests/robotests/src/com/android/settings/datausage/DataUsageSummaryTest.java b/tests/robotests/src/com/android/settings/datausage/DataUsageSummaryTest.java
index 49e3e864c9f..9fee9f0ffa2 100644
--- a/tests/robotests/src/com/android/settings/datausage/DataUsageSummaryTest.java
+++ b/tests/robotests/src/com/android/settings/datausage/DataUsageSummaryTest.java
@@ -28,6 +28,7 @@ import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 
+import android.app.usage.NetworkStatsManager;
 import android.content.Context;
 import android.net.NetworkPolicyManager;
 
@@ -68,6 +69,8 @@ public class DataUsageSummaryTest {
     private SummaryLoader mSummaryLoader;
     @Mock
     private NetworkPolicyManager mNetworkPolicyManager;
+    @Mock
+    private NetworkStatsManager mNetworkStatsManager;
     private Context mContext;
     private FragmentActivity mActivity;
     private SummaryLoader.SummaryProvider mSummaryProvider;
@@ -86,6 +89,7 @@ public class DataUsageSummaryTest {
 
         mContext = RuntimeEnvironment.application;
         mActivity = spy(Robolectric.buildActivity(FragmentActivity.class).get());
+        doReturn(mNetworkStatsManager).when(mActivity).getSystemService(NetworkStatsManager.class);
 
         mSummaryProvider = DataUsageSummary.SUMMARY_PROVIDER_FACTORY
                 .createSummaryProvider(mActivity, mSummaryLoader);
diff --git a/tests/robotests/src/com/android/settings/network/telephony/DataUsagePreferenceControllerTest.java b/tests/robotests/src/com/android/settings/network/telephony/DataUsagePreferenceControllerTest.java
index ea6f903cecc..1d91bb875ba 100644
--- a/tests/robotests/src/com/android/settings/network/telephony/DataUsagePreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/network/telephony/DataUsagePreferenceControllerTest.java
@@ -23,6 +23,7 @@ import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 
 import android.app.Activity;
+import android.app.usage.NetworkStatsManager;
 import android.content.Context;
 import android.content.Intent;
 import android.provider.Settings;
@@ -48,6 +49,8 @@ public class DataUsagePreferenceControllerTest {
 
     @Mock
     private TelephonyManager mTelephonyManager;
+    @Mock
+    private NetworkStatsManager mNetworkStatsManager;
     private DataUsagePreferenceController mController;
     private SwitchPreference mPreference;
     private Context mContext;
@@ -59,6 +62,7 @@ public class DataUsagePreferenceControllerTest {
         mContext = spy(Robolectric.setupActivity(Activity.class));
         doReturn(mTelephonyManager).when(mContext).getSystemService(Context.TELEPHONY_SERVICE);
         doReturn(mTelephonyManager).when(mTelephonyManager).createForSubscriptionId(SUB_ID);
+        doReturn(mNetworkStatsManager).when(mContext).getSystemService(NetworkStatsManager.class);
 
         mPreference = new SwitchPreference(mContext);
         mController = new DataUsagePreferenceController(mContext, "data_usage");