diff --git a/tests/cts/net/util/java/android/net/cts/util/CtsNetUtils.java b/tests/cts/net/util/java/android/net/cts/util/CtsNetUtils.java
index b32218b..4abbecc 100644
--- a/tests/cts/net/util/java/android/net/cts/util/CtsNetUtils.java
+++ b/tests/cts/net/util/java/android/net/cts/util/CtsNetUtils.java
@@ -529,12 +529,14 @@
         }
         // restore private DNS setting
         if (PRIVATE_DNS_MODE_STRICT.equals(mOldPrivateDnsMode)) {
-            setPrivateDnsStrictMode(mOldPrivateDnsSpecifier);
-
-            // In case of invalid setting, still restore it but fail the test
+            // In case of invalid setting, set to opportunistic to avoid a bad state and fail
             if (mOldPrivateDnsSpecifier == null) {
+                Settings.Global.putString(mCR, Settings.Global.PRIVATE_DNS_MODE,
+                        PRIVATE_DNS_MODE_OPPORTUNISTIC);
                 fail("Invalid private DNS setting: no hostname specified in strict mode");
             }
+            setPrivateDnsStrictMode(mOldPrivateDnsSpecifier);
+
             awaitPrivateDnsSetting("restorePrivateDnsSetting timeout",
                     mCm.getActiveNetwork(),
                     mOldPrivateDnsSpecifier, true);
diff --git a/tests/unit/Android.bp b/tests/unit/Android.bp
index 6c4bb90..6b7921b 100644
--- a/tests/unit/Android.bp
+++ b/tests/unit/Android.bp
@@ -56,21 +56,21 @@
         "java/**/*.kt",
     ],
     test_suites: ["device-tests"],
-    certificate: "platform",
     jarjar_rules: "jarjar-rules.txt",
     static_libs: [
         "androidx.test.rules",
         "bouncycastle-repackaged-unbundled",
-        "FrameworksNetCommonTests",
         "frameworks-base-testutils",
         "frameworks-net-integration-testutils",
         "framework-protos",
-        "mockito-target-minus-junit4",
+        "mockito-target-extended-minus-junit4",
         "net-tests-utils",
         "platform-test-annotations",
         "service-connectivity-pre-jarjar",
         "services.core",
         "services.net",
+        // At the bottom to prefer test libraries above instead of those in the common tests
+        "FrameworksNetCommonTests",
     ],
     libs: [
         "android.net.ipsec.ike.stubs.module_lib",
@@ -80,6 +80,9 @@
         "ServiceConnectivityResources",
     ],
     jni_libs: [
+        // For mockito extended
+        "libdexmakerjvmtiagent",
+        "libstaticjvmtiagent",
         "libservice-connectivity",
     ],
 }
diff --git a/tests/unit/AndroidManifest.xml b/tests/unit/AndroidManifest.xml
index d08b2f8..18653dc 100644
--- a/tests/unit/AndroidManifest.xml
+++ b/tests/unit/AndroidManifest.xml
@@ -50,7 +50,8 @@
     <uses-permission android:name="android.permission.NETWORK_STATS_PROVIDER" />
     <uses-permission android:name="android.permission.CONTROL_OEM_PAID_NETWORK_PREFERENCE" />
 
-    <application>
+    <!-- Debuggable is necessary for mockito extended -->
+    <application android:debuggable="true">
         <uses-library android:name="android.test.runner" />
         <uses-library android:name="android.net.ipsec.ike" />
     </application>
diff --git a/tests/unit/java/android/net/ConnectivityManagerTest.java b/tests/unit/java/android/net/ConnectivityManagerTest.java
index 591e0cc..b1b4218 100644
--- a/tests/unit/java/android/net/ConnectivityManagerTest.java
+++ b/tests/unit/java/android/net/ConnectivityManagerTest.java
@@ -46,6 +46,7 @@
 import static org.mockito.ArgumentMatchers.nullable;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.anyInt;
+import static org.mockito.Mockito.doCallRealMethod;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.reset;
@@ -216,6 +217,7 @@
         ConnectivityManager manager = new ConnectivityManager(mCtx, mService);
         NetworkRequest request = makeRequest(1);
         NetworkCallback callback = mock(ConnectivityManager.NetworkCallback.class);
+        doCallRealMethod().when(callback).onAvailable(any(), any(), any(), anyInt());
         Handler handler = new Handler(Looper.getMainLooper());
         ArgumentCaptor<Messenger> captor = ArgumentCaptor.forClass(Messenger.class);
 
@@ -244,6 +246,7 @@
         NetworkRequest req1 = makeRequest(1);
         NetworkRequest req2 = makeRequest(2);
         NetworkCallback callback = mock(ConnectivityManager.NetworkCallback.class);
+        doCallRealMethod().when(callback).onAvailable(any(), any(), any(), anyInt());
         Handler handler = new Handler(Looper.getMainLooper());
         ArgumentCaptor<Messenger> captor = ArgumentCaptor.forClass(Messenger.class);
 
diff --git a/tests/unit/java/android/net/nsd/NsdManagerTest.java b/tests/unit/java/android/net/nsd/NsdManagerTest.java
index b0a9b8a..370179c 100644
--- a/tests/unit/java/android/net/nsd/NsdManagerTest.java
+++ b/tests/unit/java/android/net/nsd/NsdManagerTest.java
@@ -20,12 +20,12 @@
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.fail;
 import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.timeout;
 import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
 
 import android.content.Context;
 import android.os.Handler;
@@ -66,7 +66,7 @@
         MockitoAnnotations.initMocks(this);
 
         mServiceHandler = spy(MockServiceHandler.create(mContext));
-        when(mService.getMessenger()).thenReturn(new Messenger(mServiceHandler));
+        doReturn(new Messenger(mServiceHandler)).when(mService).getMessenger();
 
         mManager = makeManager();
     }
diff --git a/tests/unit/java/android/net/util/MultinetworkPolicyTrackerTest.kt b/tests/unit/java/android/net/util/MultinetworkPolicyTrackerTest.kt
index 25aa626..9291532 100644
--- a/tests/unit/java/android/net/util/MultinetworkPolicyTrackerTest.kt
+++ b/tests/unit/java/android/net/util/MultinetworkPolicyTrackerTest.kt
@@ -74,6 +74,7 @@
         doReturn(Context.TELEPHONY_SERVICE).`when`(it)
                 .getSystemServiceName(TelephonyManager::class.java)
         doReturn(telephonyManager).`when`(it).getSystemService(Context.TELEPHONY_SERVICE)
+        doReturn(telephonyManager).`when`(it).getSystemService(TelephonyManager::class.java)
         doReturn(subscriptionManager).`when`(it)
                 .getSystemService(Context.TELEPHONY_SUBSCRIPTION_SERVICE)
         doReturn(resolver).`when`(it).contentResolver
diff --git a/tests/unit/java/com/android/server/ConnectivityServiceTest.java b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
index a9612de..2c19cb8 100644
--- a/tests/unit/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
@@ -18,10 +18,14 @@
 
 import static android.Manifest.permission.CHANGE_NETWORK_STATE;
 import static android.Manifest.permission.CONNECTIVITY_USE_RESTRICTED_NETWORKS;
+import static android.Manifest.permission.CONTROL_OEM_PAID_NETWORK_PREFERENCE;
 import static android.Manifest.permission.DUMP;
+import static android.Manifest.permission.GET_INTENT_SENDER_INTENT;
 import static android.Manifest.permission.LOCAL_MAC_ADDRESS;
 import static android.Manifest.permission.NETWORK_FACTORY;
 import static android.Manifest.permission.NETWORK_SETTINGS;
+import static android.Manifest.permission.NETWORK_STACK;
+import static android.Manifest.permission.PACKET_KEEPALIVE_OFFLOAD;
 import static android.app.PendingIntent.FLAG_IMMUTABLE;
 import static android.content.Intent.ACTION_PACKAGE_ADDED;
 import static android.content.Intent.ACTION_PACKAGE_REMOVED;
@@ -130,6 +134,7 @@
 import static com.android.testutils.MiscAsserts.assertRunsInAtMost;
 import static com.android.testutils.MiscAsserts.assertSameElements;
 import static com.android.testutils.MiscAsserts.assertThrows;
+import static com.android.testutils.TestPermissionUtil.runAsShell;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -255,6 +260,7 @@
 import android.net.shared.PrivateDnsConfig;
 import android.net.util.MultinetworkPolicyTracker;
 import android.os.BadParcelableException;
+import android.os.BatteryStatsManager;
 import android.os.Binder;
 import android.os.Build;
 import android.os.Bundle;
@@ -463,6 +469,7 @@
     @Mock DeviceIdleInternal mDeviceIdleInternal;
     @Mock INetworkManagementService mNetworkManagementService;
     @Mock NetworkStatsManager mStatsManager;
+    @Mock BatteryStatsManager mBatteryStatsManager;
     @Mock IDnsResolver mMockDnsResolver;
     @Mock INetd mMockNetd;
     @Mock NetworkStackClientBase mNetworkStack;
@@ -573,6 +580,7 @@
             if (Context.NETWORK_POLICY_SERVICE.equals(name)) return mNetworkPolicyManager;
             if (Context.SYSTEM_CONFIG_SERVICE.equals(name)) return mSystemConfigManager;
             if (Context.NETWORK_STATS_SERVICE.equals(name)) return mStatsManager;
+            if (Context.BATTERY_STATS_SERVICE.equals(name)) return mBatteryStatsManager;
             return super.getSystemService(name);
         }
 
@@ -653,6 +661,13 @@
         public void setPermission(String permission, Integer granted) {
             mMockedPermissions.put(permission, granted);
         }
+
+        @Override
+        public Intent registerReceiverForAllUsers(@Nullable BroadcastReceiver receiver,
+                @NonNull IntentFilter filter, @Nullable String broadcastPermission,
+                @Nullable Handler scheduler) {
+            return null;
+        }
     }
 
     private void waitForIdle() {
@@ -1202,7 +1217,22 @@
                             return mDeviceIdleInternal;
                         }
                     },
-                    mNetworkManagementService, mMockNetd, userId, mVpnProfileStore);
+                    mNetworkManagementService, mMockNetd, userId, mVpnProfileStore,
+                    new SystemServices(mServiceContext) {
+                        @Override
+                        public String settingsSecureGetStringForUser(String key, int userId) {
+                            switch (key) {
+                                // Settings keys not marked as @Readable are not readable from
+                                // non-privileged apps, unless marked as testOnly=true
+                                // (atest refuses to install testOnly=true apps), even if mocked
+                                // in the content provider.
+                                case Settings.Secure.ALWAYS_ON_VPN_APP:
+                                    return null;
+                                default:
+                                    return super.settingsSecureGetStringForUser(key, userId);
+                            }
+                        }
+                    }, new Ikev2SessionCreator());
         }
 
         public void setUids(Set<UidRange> uids) {
@@ -1580,6 +1610,11 @@
         mServiceContext = new MockContext(InstrumentationRegistry.getContext(),
                 new FakeSettingsProvider());
         mServiceContext.setUseRegisteredHandlers(true);
+        mServiceContext.setPermission(NETWORK_FACTORY, PERMISSION_GRANTED);
+        mServiceContext.setPermission(NETWORK_STACK, PERMISSION_GRANTED);
+        mServiceContext.setPermission(CONTROL_OEM_PAID_NETWORK_PREFERENCE, PERMISSION_GRANTED);
+        mServiceContext.setPermission(PACKET_KEEPALIVE_OFFLOAD, PERMISSION_GRANTED);
+        mServiceContext.setPermission(CONNECTIVITY_USE_RESTRICTED_NETWORKS, PERMISSION_GRANTED);
 
         mAlarmManagerThread = new HandlerThread("TestAlarmManager");
         mAlarmManagerThread.start();
@@ -1810,6 +1845,16 @@
         assertEmpty(mCm.getAllNetworkStateSnapshots());
     }
 
+    private static PendingIntent wrapPendingIntent(final PendingIntent intent) {
+        final PendingIntent ret = spy(intent);
+        // intentFilterEquals requires GET_INTENT_SENDER_INTENT permission
+        doAnswer(inv -> {
+            final PendingIntent other = inv.getArgument(0);
+            return runAsShell(GET_INTENT_SENDER_INTENT, () -> intent.intentFilterEquals(other));
+        }).when(ret).intentFilterEquals(any());
+        return ret;
+    }
+
     /**
      * Class to simplify expecting broadcasts using BroadcastInterceptingContext.
      * Ensures that the receiver is unregistered after the expected broadcast is received. This
@@ -3258,8 +3303,8 @@
 
     @Test
     public void testNoMutableNetworkRequests() throws Exception {
-        final PendingIntent pendingIntent = PendingIntent.getBroadcast(
-                mContext, 0 /* requestCode */, new Intent("a"), FLAG_IMMUTABLE);
+        final PendingIntent pendingIntent = wrapPendingIntent(PendingIntent.getBroadcast(
+                mContext, 0 /* requestCode */, new Intent("a"), FLAG_IMMUTABLE));
         NetworkRequest request1 = new NetworkRequest.Builder()
                 .addCapability(NET_CAPABILITY_VALIDATED)
                 .build();
@@ -4103,16 +4148,16 @@
                 mCm.registerNetworkCallback(r, new NetworkCallback()));
 
         assertThrows(SecurityException.class, () ->
-                mCm.registerNetworkCallback(r, PendingIntent.getService(
-                        mServiceContext, 0 /* requestCode */, new Intent(), FLAG_IMMUTABLE)));
+                mCm.registerNetworkCallback(r, wrapPendingIntent(PendingIntent.getService(
+                        mServiceContext, 0 /* requestCode */, new Intent(), FLAG_IMMUTABLE))));
 
         // Requesting a Network with signal strength should get IllegalArgumentException.
         assertThrows(IllegalArgumentException.class, () ->
                 mCm.requestNetwork(r, new NetworkCallback()));
 
         assertThrows(IllegalArgumentException.class, () ->
-                mCm.requestNetwork(r, PendingIntent.getService(
-                        mServiceContext, 0 /* requestCode */, new Intent(), FLAG_IMMUTABLE)));
+                mCm.requestNetwork(r, wrapPendingIntent(PendingIntent.getService(
+                        mServiceContext, 0 /* requestCode */, new Intent(), FLAG_IMMUTABLE))));
     }
 
     @Test
@@ -5668,14 +5713,14 @@
         }
         j = 0;
         while (j++ < INTENTS / 2) {
-            final PendingIntent pi = PendingIntent.getBroadcast(mContext, 0 /* requestCode */,
-                    new Intent("a" + j), FLAG_IMMUTABLE);
+            final PendingIntent pi = wrapPendingIntent(PendingIntent.getBroadcast(mContext,
+                    0 /* requestCode */, new Intent("a" + j), FLAG_IMMUTABLE));
             mCm.requestNetwork(networkRequest, pi);
             registered.add(pi);
         }
         while (j++ < INTENTS) {
-            final PendingIntent pi = PendingIntent.getBroadcast(mContext, 0 /* requestCode */,
-                    new Intent("b" + j), FLAG_IMMUTABLE);
+            final PendingIntent pi = wrapPendingIntent(PendingIntent.getBroadcast(mContext,
+                    0 /* requestCode */, new Intent("b" + j), FLAG_IMMUTABLE));
             mCm.registerNetworkCallback(networkRequest, pi);
             registered.add(pi);
         }
@@ -5689,13 +5734,13 @@
         );
         assertThrows(TooManyRequestsException.class, () ->
                 mCm.requestNetwork(networkRequest,
-                        PendingIntent.getBroadcast(mContext, 0 /* requestCode */,
-                                new Intent("c"), FLAG_IMMUTABLE))
+                        wrapPendingIntent(PendingIntent.getBroadcast(mContext, 0 /* requestCode */,
+                                new Intent("c"), FLAG_IMMUTABLE)))
         );
         assertThrows(TooManyRequestsException.class, () ->
                 mCm.registerNetworkCallback(networkRequest,
-                        PendingIntent.getBroadcast(mContext, 0 /* requestCode */,
-                                new Intent("d"), FLAG_IMMUTABLE))
+                        wrapPendingIntent(PendingIntent.getBroadcast(mContext, 0 /* requestCode */,
+                                new Intent("d"), FLAG_IMMUTABLE)))
         );
 
         // The system gets another SYSTEM_ONLY_MAX_REQUESTS slots.
@@ -5775,16 +5820,16 @@
         waitForIdle();
 
         for (int i = 0; i < MAX_REQUESTS; i++) {
-            final PendingIntent pendingIntent = PendingIntent.getBroadcast(
-                    mContext, 0 /* requestCode */, new Intent("e" + i), FLAG_IMMUTABLE);
+            final PendingIntent pendingIntent = wrapPendingIntent(PendingIntent.getBroadcast(
+                    mContext, 0 /* requestCode */, new Intent("e" + i), FLAG_IMMUTABLE));
             mCm.requestNetwork(networkRequest, pendingIntent);
             mCm.unregisterNetworkCallback(pendingIntent);
         }
         waitForIdle();
 
         for (int i = 0; i < MAX_REQUESTS; i++) {
-            final PendingIntent pendingIntent = PendingIntent.getBroadcast(
-                    mContext, 0 /* requestCode */, new Intent("f" + i), FLAG_IMMUTABLE);
+            final PendingIntent pendingIntent = wrapPendingIntent(PendingIntent.getBroadcast(
+                    mContext, 0 /* requestCode */, new Intent("f" + i), FLAG_IMMUTABLE));
             mCm.registerNetworkCallback(networkRequest, pendingIntent);
             mCm.unregisterNetworkCallback(pendingIntent);
         }
@@ -9227,8 +9272,7 @@
         mServiceContext.setPermission(NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK,
                 PERMISSION_DENIED);
         mServiceContext.setPermission(NETWORK_SETTINGS, PERMISSION_DENIED);
-        mServiceContext.setPermission(Manifest.permission.NETWORK_STACK,
-                PERMISSION_DENIED);
+        mServiceContext.setPermission(NETWORK_STACK, PERMISSION_DENIED);
         mServiceContext.setPermission(Manifest.permission.NETWORK_SETUP_WIZARD,
                 PERMISSION_DENIED);
     }
@@ -9669,7 +9713,7 @@
         setupConnectionOwnerUid(vpnOwnerUid, vpnType);
 
         // Test as VPN app
-        mServiceContext.setPermission(android.Manifest.permission.NETWORK_STACK, PERMISSION_DENIED);
+        mServiceContext.setPermission(NETWORK_STACK, PERMISSION_DENIED);
         mServiceContext.setPermission(
                 NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK, PERMISSION_DENIED);
     }
@@ -9709,8 +9753,7 @@
     public void testGetConnectionOwnerUidVpnServiceNetworkStackDoesNotThrow() throws Exception {
         final int myUid = Process.myUid();
         setupConnectionOwnerUid(myUid, VpnManager.TYPE_VPN_SERVICE);
-        mServiceContext.setPermission(
-                android.Manifest.permission.NETWORK_STACK, PERMISSION_GRANTED);
+        mServiceContext.setPermission(NETWORK_STACK, PERMISSION_GRANTED);
 
         assertEquals(42, mService.getConnectionOwnerUid(getTestConnectionInfo()));
     }
@@ -9878,8 +9921,7 @@
     public void testCheckConnectivityDiagnosticsPermissionsNetworkStack() throws Exception {
         final NetworkAgentInfo naiWithoutUid = fakeMobileNai(new NetworkCapabilities());
 
-        mServiceContext.setPermission(
-                android.Manifest.permission.NETWORK_STACK, PERMISSION_GRANTED);
+        mServiceContext.setPermission(NETWORK_STACK, PERMISSION_GRANTED);
         assertTrue(
                 "NetworkStack permission not applied",
                 mService.checkConnectivityDiagnosticsPermissions(
@@ -9895,7 +9937,7 @@
         nc.setAdministratorUids(new int[] {wrongUid});
         final NetworkAgentInfo naiWithUid = fakeWifiNai(nc);
 
-        mServiceContext.setPermission(android.Manifest.permission.NETWORK_STACK, PERMISSION_DENIED);
+        mServiceContext.setPermission(NETWORK_STACK, PERMISSION_DENIED);
 
         assertFalse(
                 "Mismatched uid/package name should not pass the location permission check",
@@ -9905,7 +9947,7 @@
 
     private void verifyConnectivityDiagnosticsPermissionsWithNetworkAgentInfo(
             NetworkAgentInfo info, boolean expectPermission) {
-        mServiceContext.setPermission(android.Manifest.permission.NETWORK_STACK, PERMISSION_DENIED);
+        mServiceContext.setPermission(NETWORK_STACK, PERMISSION_DENIED);
 
         assertEquals(
                 "Unexpected ConnDiags permission",
@@ -9973,7 +10015,7 @@
 
         setupLocationPermissions(Build.VERSION_CODES.Q, true, AppOpsManager.OPSTR_FINE_LOCATION,
                 Manifest.permission.ACCESS_FINE_LOCATION);
-        mServiceContext.setPermission(android.Manifest.permission.NETWORK_STACK, PERMISSION_DENIED);
+        mServiceContext.setPermission(NETWORK_STACK, PERMISSION_DENIED);
 
         assertTrue(
                 "NetworkCapabilities administrator uid permission not applied",
@@ -9990,7 +10032,7 @@
 
         setupLocationPermissions(Build.VERSION_CODES.Q, true, AppOpsManager.OPSTR_FINE_LOCATION,
                 Manifest.permission.ACCESS_FINE_LOCATION);
-        mServiceContext.setPermission(android.Manifest.permission.NETWORK_STACK, PERMISSION_DENIED);
+        mServiceContext.setPermission(NETWORK_STACK, PERMISSION_DENIED);
 
         // Use wrong pid and uid
         assertFalse(
@@ -10016,8 +10058,7 @@
         final NetworkRequest request = new NetworkRequest.Builder().build();
         when(mConnectivityDiagnosticsCallback.asBinder()).thenReturn(mIBinder);
 
-        mServiceContext.setPermission(
-                android.Manifest.permission.NETWORK_STACK, PERMISSION_GRANTED);
+        mServiceContext.setPermission(NETWORK_STACK, PERMISSION_GRANTED);
 
         mService.registerConnectivityDiagnosticsCallback(
                 mConnectivityDiagnosticsCallback, request, mContext.getPackageName());
@@ -10036,8 +10077,7 @@
         final NetworkRequest request = new NetworkRequest.Builder().build();
         when(mConnectivityDiagnosticsCallback.asBinder()).thenReturn(mIBinder);
 
-        mServiceContext.setPermission(
-                android.Manifest.permission.NETWORK_STACK, PERMISSION_GRANTED);
+        mServiceContext.setPermission(NETWORK_STACK, PERMISSION_GRANTED);
 
         mService.registerConnectivityDiagnosticsCallback(
                 mConnectivityDiagnosticsCallback, request, mContext.getPackageName());
@@ -12768,8 +12808,8 @@
     @Test
     public void testNetworkRequestWithSubIdsWithNetworkFactoryPermission() throws Exception {
         mServiceContext.setPermission(NETWORK_FACTORY, PERMISSION_GRANTED);
-        final PendingIntent pendingIntent = PendingIntent.getBroadcast(
-                mContext, 0 /* requestCode */, new Intent("a"), FLAG_IMMUTABLE);
+        final PendingIntent pendingIntent = wrapPendingIntent(PendingIntent.getBroadcast(
+                mContext, 0 /* requestCode */, new Intent("a"), FLAG_IMMUTABLE));
         final NetworkCallback networkCallback1 = new NetworkCallback();
         final NetworkCallback networkCallback2 = new NetworkCallback();
 
@@ -12785,8 +12825,8 @@
     @Test
     public void testNetworkRequestWithSubIdsWithoutNetworkFactoryPermission() throws Exception {
         mServiceContext.setPermission(NETWORK_FACTORY, PERMISSION_DENIED);
-        final PendingIntent pendingIntent = PendingIntent.getBroadcast(
-                mContext, 0 /* requestCode */, new Intent("a"), FLAG_IMMUTABLE);
+        final PendingIntent pendingIntent = wrapPendingIntent(PendingIntent.getBroadcast(
+                mContext, 0 /* requestCode */, new Intent("a"), FLAG_IMMUTABLE));
 
         final Class<SecurityException> expected = SecurityException.class;
         assertThrows(
diff --git a/tests/unit/java/com/android/server/connectivity/MultipathPolicyTrackerTest.java b/tests/unit/java/com/android/server/connectivity/MultipathPolicyTrackerTest.java
index 38f6d7f..d6bd08e 100644
--- a/tests/unit/java/com/android/server/connectivity/MultipathPolicyTrackerTest.java
+++ b/tests/unit/java/com/android/server/connectivity/MultipathPolicyTrackerTest.java
@@ -116,6 +116,7 @@
     private <T> void mockService(String serviceName, Class<T> serviceClass, T service) {
         when(mContext.getSystemServiceName(serviceClass)).thenReturn(serviceName);
         when(mContext.getSystemService(serviceName)).thenReturn(service);
+        when(mContext.getSystemService(serviceClass)).thenReturn(service);
     }
 
     @Before
diff --git a/tests/unit/java/com/android/server/connectivity/PermissionMonitorTest.java b/tests/unit/java/com/android/server/connectivity/PermissionMonitorTest.java
index e98f5db..0bfc129 100644
--- a/tests/unit/java/com/android/server/connectivity/PermissionMonitorTest.java
+++ b/tests/unit/java/com/android/server/connectivity/PermissionMonitorTest.java
@@ -139,6 +139,7 @@
                 .thenReturn(Context.SYSTEM_CONFIG_SERVICE);
         when(mContext.getSystemService(Context.SYSTEM_CONFIG_SERVICE))
                 .thenReturn(mSystemConfigManager);
+        when(mContext.getSystemService(SystemConfigManager.class)).thenReturn(mSystemConfigManager);
         when(mSystemConfigManager.getSystemPermissionUids(anyString())).thenReturn(new int[0]);
         final Context asUserCtx = mock(Context.class, AdditionalAnswers.delegatesTo(mContext));
         doReturn(UserHandle.ALL).when(asUserCtx).getUser();
diff --git a/tests/unit/java/com/android/server/connectivity/VpnTest.java b/tests/unit/java/com/android/server/connectivity/VpnTest.java
index b725b82..a086451 100644
--- a/tests/unit/java/com/android/server/connectivity/VpnTest.java
+++ b/tests/unit/java/com/android/server/connectivity/VpnTest.java
@@ -219,19 +219,11 @@
 
         when(mContext.getPackageName()).thenReturn(TEST_VPN_PKG);
         when(mContext.getOpPackageName()).thenReturn(TEST_VPN_PKG);
-        when(mContext.getSystemServiceName(UserManager.class))
-                .thenReturn(Context.USER_SERVICE);
-        when(mContext.getSystemService(eq(Context.USER_SERVICE))).thenReturn(mUserManager);
-        when(mContext.getSystemService(eq(Context.APP_OPS_SERVICE))).thenReturn(mAppOps);
-        when(mContext.getSystemServiceName(NotificationManager.class))
-                .thenReturn(Context.NOTIFICATION_SERVICE);
-        when(mContext.getSystemService(eq(Context.NOTIFICATION_SERVICE)))
-                .thenReturn(mNotificationManager);
-        when(mContext.getSystemService(eq(Context.CONNECTIVITY_SERVICE)))
-                .thenReturn(mConnectivityManager);
-        when(mContext.getSystemServiceName(eq(ConnectivityManager.class)))
-                .thenReturn(Context.CONNECTIVITY_SERVICE);
-        when(mContext.getSystemService(eq(Context.IPSEC_SERVICE))).thenReturn(mIpSecManager);
+        mockService(UserManager.class, Context.USER_SERVICE, mUserManager);
+        mockService(AppOpsManager.class, Context.APP_OPS_SERVICE, mAppOps);
+        mockService(NotificationManager.class, Context.NOTIFICATION_SERVICE, mNotificationManager);
+        mockService(ConnectivityManager.class, Context.CONNECTIVITY_SERVICE, mConnectivityManager);
+        mockService(IpSecManager.class, Context.IPSEC_SERVICE, mIpSecManager);
         when(mContext.getString(R.string.config_customVpnAlwaysOnDisconnectedDialogComponent))
                 .thenReturn(Resources.getSystem().getString(
                         R.string.config_customVpnAlwaysOnDisconnectedDialogComponent));
@@ -259,6 +251,12 @@
                 .thenReturn(tunnelResp);
     }
 
+    private <T> void mockService(Class<T> clazz, String name, T service) {
+        doReturn(service).when(mContext).getSystemService(clazz);
+        doReturn(service).when(mContext).getSystemService(name);
+        doReturn(name).when(mContext).getSystemServiceName(clazz);
+    }
+
     private Set<Range<Integer>> rangeSet(Range<Integer> ... ranges) {
         final Set<Range<Integer>> range = new ArraySet<>();
         for (Range<Integer> r : ranges) range.add(r);
diff --git a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
index 0ba5f7d..da2fdf4 100644
--- a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
+++ b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
@@ -16,12 +16,16 @@
 
 package com.android.server.net;
 
+import static android.Manifest.permission.READ_NETWORK_USAGE_HISTORY;
+import static android.Manifest.permission.UPDATE_DEVICE_STATS;
 import static android.content.Intent.ACTION_UID_REMOVED;
 import static android.content.Intent.EXTRA_UID;
+import static android.content.pm.PackageManager.PERMISSION_GRANTED;
 import static android.net.ConnectivityManager.TYPE_MOBILE;
 import static android.net.ConnectivityManager.TYPE_WIFI;
 import static android.net.NetworkIdentity.OEM_PAID;
 import static android.net.NetworkIdentity.OEM_PRIVATE;
+import static android.net.NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK;
 import static android.net.NetworkStats.DEFAULT_NETWORK_ALL;
 import static android.net.NetworkStats.DEFAULT_NETWORK_NO;
 import static android.net.NetworkStats.DEFAULT_NETWORK_YES;
@@ -106,6 +110,7 @@
 import android.provider.Settings;
 import android.telephony.TelephonyManager;
 
+import androidx.annotation.Nullable;
 import androidx.test.InstrumentationRegistry;
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
@@ -199,6 +204,26 @@
             if (Context.TELEPHONY_SERVICE.equals(name)) return mTelephonyManager;
             return mBaseContext.getSystemService(name);
         }
+
+        @Override
+        public void enforceCallingOrSelfPermission(String permission, @Nullable String message) {
+            if (checkCallingOrSelfPermission(permission) != PERMISSION_GRANTED) {
+                super.enforceCallingOrSelfPermission(permission, message);
+            }
+        }
+
+        @Override
+        public int checkCallingOrSelfPermission(String permission) {
+            switch (permission) {
+                case PERMISSION_MAINLINE_NETWORK_STACK:
+                case READ_NETWORK_USAGE_HISTORY:
+                case UPDATE_DEVICE_STATS:
+                    return PERMISSION_GRANTED;
+                default:
+                    return super.checkCallingOrSelfPermission(permission);
+            }
+
+        }
     }
 
     private final Clock mClock = new SimpleClock(ZoneOffset.UTC) {