diff --git a/Tethering/src/android/net/ip/IpServer.java b/Tethering/src/android/net/ip/IpServer.java
index 859f23a..cf49683 100644
--- a/Tethering/src/android/net/ip/IpServer.java
+++ b/Tethering/src/android/net/ip/IpServer.java
@@ -52,7 +52,6 @@
 import android.net.util.InterfaceSet;
 import android.net.util.PrefixUtils;
 import android.net.util.SharedLog;
-import android.os.Build;
 import android.os.Handler;
 import android.os.Looper;
 import android.os.Message;
@@ -67,6 +66,7 @@
 import com.android.internal.util.MessageUtils;
 import com.android.internal.util.State;
 import com.android.internal.util.StateMachine;
+import com.android.modules.utils.build.SdkLevel;
 import com.android.networkstack.tethering.BpfCoordinator;
 import com.android.networkstack.tethering.BpfCoordinator.ClientInfo;
 import com.android.networkstack.tethering.BpfCoordinator.Ipv6ForwardingRule;
@@ -676,9 +676,7 @@
             return false;
         }
 
-        // TODO: use ShimUtils instead of explicitly checking the version here.
-        if (Build.VERSION.SDK_INT > Build.VERSION_CODES.R || "S".equals(Build.VERSION.CODENAME)
-                    || "T".equals(Build.VERSION.CODENAME)) {
+        if (SdkLevel.isAtLeastS()) {
             // DAD Proxy starts forwarding packets after IPv6 upstream is present.
             mDadProxy = mDeps.getDadProxy(getHandler(), mInterfaceParams);
         }
diff --git a/framework/src/android/net/NetworkCapabilities.java b/framework/src/android/net/NetworkCapabilities.java
index ec71d3d..5bd0af8 100644
--- a/framework/src/android/net/NetworkCapabilities.java
+++ b/framework/src/android/net/NetworkCapabilities.java
@@ -818,7 +818,7 @@
      * restrictions.
      * @hide
      */
-    public void restrictCapabilitesForTestNetwork(int creatorUid) {
+    public void restrictCapabilitiesForTestNetwork(int creatorUid) {
         final long originalCapabilities = mNetworkCapabilities;
         final long originalTransportTypes = mTransportTypes;
         final NetworkSpecifier originalSpecifier = mNetworkSpecifier;
@@ -828,7 +828,7 @@
         final TransportInfo originalTransportInfo = getTransportInfo();
         final Set<Integer> originalSubIds = getSubscriptionIds();
         clearAll();
-        if (0 != (originalCapabilities & NET_CAPABILITY_NOT_RESTRICTED)) {
+        if (0 != (originalCapabilities & (1 << NET_CAPABILITY_NOT_RESTRICTED))) {
             // If the test network is not restricted, then it is only allowed to declare some
             // specific transports. This is to minimize impact on running apps in case an app
             // run from the shell creates a test a network.
diff --git a/service/src/com/android/server/ConnectivityService.java b/service/src/com/android/server/ConnectivityService.java
index 34e15ca..47bde72 100644
--- a/service/src/com/android/server/ConnectivityService.java
+++ b/service/src/com/android/server/ConnectivityService.java
@@ -3255,7 +3255,7 @@
                         // the Messenger, but if this ever changes, not making a defensive copy
                         // here will give attack vectors to clients using this code path.
                         networkCapabilities = new NetworkCapabilities(networkCapabilities);
-                        networkCapabilities.restrictCapabilitesForTestNetwork(nai.creatorUid);
+                        networkCapabilities.restrictCapabilitiesForTestNetwork(nai.creatorUid);
                     }
                     processCapabilitiesFromAgent(nai, networkCapabilities);
                     updateCapabilities(nai.getCurrentScore(), nai, networkCapabilities);
@@ -6764,7 +6764,7 @@
             // the call to mixInCapabilities below anyway, but sanitizing here means the NAI never
             // sees capabilities that may be malicious, which might prevent mistakes in the future.
             networkCapabilities = new NetworkCapabilities(networkCapabilities);
-            networkCapabilities.restrictCapabilitesForTestNetwork(uid);
+            networkCapabilities.restrictCapabilitiesForTestNetwork(uid);
         }
 
         LinkProperties lp = new LinkProperties(linkProperties);
@@ -9525,6 +9525,10 @@
             @NonNull IConnectivityDiagnosticsCallback callback,
             @NonNull NetworkRequest request,
             @NonNull String callingPackageName) {
+        Objects.requireNonNull(callback, "callback must not be null");
+        Objects.requireNonNull(request, "request must not be null");
+        Objects.requireNonNull(callingPackageName, "callingPackageName must not be null");
+
         if (request.legacyType != TYPE_NONE) {
             throw new IllegalArgumentException("ConnectivityManager.TYPE_* are deprecated."
                     + " Please use NetworkCapabilities instead.");
@@ -9573,6 +9577,9 @@
     @Override
     public void simulateDataStall(int detectionMethod, long timestampMillis,
             @NonNull Network network, @NonNull PersistableBundle extras) {
+        Objects.requireNonNull(network, "network must not be null");
+        Objects.requireNonNull(extras, "extras must not be null");
+
         enforceAnyPermissionOf(android.Manifest.permission.MANAGE_TEST_NETWORKS,
                 android.Manifest.permission.NETWORK_STACK);
         final NetworkCapabilities nc = getNetworkCapabilitiesInternal(network);
diff --git a/tests/common/java/android/net/NetworkCapabilitiesTest.java b/tests/common/java/android/net/NetworkCapabilitiesTest.java
index 493a201..3d0cb92 100644
--- a/tests/common/java/android/net/NetworkCapabilitiesTest.java
+++ b/tests/common/java/android/net/NetworkCapabilitiesTest.java
@@ -1167,4 +1167,48 @@
         // Ensure test case fails if new net cap is added into default cap but no update here.
         assertEquals(0, nc.getCapabilities().length);
     }
+
+    @Test @IgnoreUpTo(Build.VERSION_CODES.R)
+    public void testRestrictCapabilitiesForTestNetwork() {
+        final int ownerUid = 1234;
+        final int[] administratorUids = {1001, ownerUid};
+        final NetworkCapabilities nonRestrictedNc = new NetworkCapabilities.Builder()
+                .addTransportType(TRANSPORT_CELLULAR)
+                .addTransportType(TRANSPORT_VPN)
+                .addCapability(NET_CAPABILITY_MMS)
+                .addCapability(NET_CAPABILITY_NOT_METERED)
+                .setAdministratorUids(administratorUids)
+                .setOwnerUid(ownerUid)
+                .setSubscriptionIds(Set.of(TEST_SUBID1)).build();
+
+        nonRestrictedNc.restrictCapabilitiesForTestNetwork(ownerUid);
+        // TRANSPORT_TEST will be appended
+        assertTrue(nonRestrictedNc.hasTransport(TRANSPORT_TEST));
+        assertEquals(Set.of(TEST_SUBID1), nonRestrictedNc.getSubscriptionIds());
+        // Non-UNRESTRICTED_TEST_NETWORKS_ALLOWED_TRANSPORTS will be removed
+        assertFalse(nonRestrictedNc.hasTransport(TRANSPORT_CELLULAR));
+        assertTrue(nonRestrictedNc.hasTransport(TRANSPORT_VPN));
+        // Only TEST_NETWORKS_ALLOWED_CAPABILITIES will be kept
+        assertFalse(nonRestrictedNc.hasCapability(NET_CAPABILITY_MMS));
+        assertTrue(nonRestrictedNc.hasCapability(NET_CAPABILITY_NOT_METERED));
+
+        final NetworkCapabilities restrictedNc = new NetworkCapabilities.Builder(nonRestrictedNc)
+                .removeCapability(NET_CAPABILITY_NOT_RESTRICTED)
+                .addTransportType(TRANSPORT_CELLULAR)
+                .addCapability(NET_CAPABILITY_MMS).build();
+        restrictedNc.restrictCapabilitiesForTestNetwork(ownerUid);
+        // It may declare any transport if the net cap is restricted
+        assertTrue(restrictedNc.hasTransport(TRANSPORT_CELLULAR));
+        // SubIds will be cleared.
+        assertEquals(new ArraySet<>(), restrictedNc.getSubscriptionIds());
+        // Only retain the owner and administrator UIDs if they match the app registering the remote
+        // caller that registered the network.
+        assertEquals(ownerUid, restrictedNc.getOwnerUid());
+        assertArrayEquals(new int[] {ownerUid}, restrictedNc.getAdministratorUids());
+        // The creator UID does not match the owner and administrator UIDs will clear the owner and
+        // administrator UIDs.
+        restrictedNc.restrictCapabilitiesForTestNetwork(5678);
+        assertEquals(INVALID_UID, restrictedNc.getOwnerUid());
+        assertArrayEquals(new int[0], restrictedNc.getAdministratorUids());
+    }
 }
diff --git a/tests/integration/Android.bp b/tests/integration/Android.bp
index 7c42811..2bc985a 100644
--- a/tests/integration/Android.bp
+++ b/tests/integration/Android.bp
@@ -29,6 +29,7 @@
     ],
     libs: [
         "android.test.mock",
+        "ServiceConnectivityResources",
     ],
     static_libs: [
         "NetworkStackApiStableLib",
diff --git a/tests/integration/src/com/android/server/net/integrationtests/ConnectivityServiceIntegrationTest.kt b/tests/integration/src/com/android/server/net/integrationtests/ConnectivityServiceIntegrationTest.kt
index e039ef0..80338aa 100644
--- a/tests/integration/src/com/android/server/net/integrationtests/ConnectivityServiceIntegrationTest.kt
+++ b/tests/integration/src/com/android/server/net/integrationtests/ConnectivityServiceIntegrationTest.kt
@@ -23,7 +23,9 @@
 import android.content.Context.BIND_IMPORTANT
 import android.content.Intent
 import android.content.ServiceConnection
+import android.content.res.Resources
 import android.net.ConnectivityManager
+import android.net.ConnectivityResources
 import android.net.IDnsResolver
 import android.net.INetd
 import android.net.LinkProperties
@@ -35,6 +37,7 @@
 import android.net.TestNetworkStackClient
 import android.net.Uri
 import android.net.metrics.IpConnectivityLog
+import android.net.util.MultinetworkPolicyTracker
 import android.os.ConditionVariable
 import android.os.IBinder
 import android.os.SystemConfigManager
@@ -43,6 +46,7 @@
 import android.util.Log
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.platform.app.InstrumentationRegistry
+import com.android.connectivity.resources.R
 import com.android.server.ConnectivityService
 import com.android.server.NetworkAgentWrapper
 import com.android.server.TestNetIdManager
@@ -59,6 +63,7 @@
 import org.mockito.Mock
 import org.mockito.Mockito.any
 import org.mockito.Mockito.anyInt
+import org.mockito.Mockito.doAnswer
 import org.mockito.Mockito.doNothing
 import org.mockito.Mockito.doReturn
 import org.mockito.Mockito.eq
@@ -93,6 +98,10 @@
     private lateinit var dnsResolver: IDnsResolver
     @Mock
     private lateinit var systemConfigManager: SystemConfigManager
+    @Mock
+    private lateinit var resources: Resources
+    @Mock
+    private lateinit var resourcesContext: Context
     @Spy
     private var context = TestableContext(realContext)
 
@@ -110,9 +119,11 @@
 
         private val realContext get() = InstrumentationRegistry.getInstrumentation().context
         private val httpProbeUrl get() =
-            realContext.getResources().getString(R.string.config_captive_portal_http_url)
+            realContext.getResources().getString(com.android.server.net.integrationtests.R.string
+                    .config_captive_portal_http_url)
         private val httpsProbeUrl get() =
-            realContext.getResources().getString(R.string.config_captive_portal_https_url)
+            realContext.getResources().getString(com.android.server.net.integrationtests.R.string
+                    .config_captive_portal_https_url)
 
         private class InstrumentationServiceConnection : ServiceConnection {
             override fun onServiceConnected(name: ComponentName?, service: IBinder?) {
@@ -156,6 +167,27 @@
                 .getSystemService(Context.SYSTEM_CONFIG_SERVICE)
         doReturn(IntArray(0)).`when`(systemConfigManager).getSystemPermissionUids(anyString())
 
+        doReturn(60000).`when`(resources).getInteger(R.integer.config_networkTransitionTimeout)
+        doReturn("").`when`(resources).getString(R.string.config_networkCaptivePortalServerUrl)
+        doReturn(arrayOf<String>("test_wlan_wol")).`when`(resources)
+                .getStringArray(R.array.config_wakeonlan_supported_interfaces)
+        doReturn(arrayOf("0,1", "1,3")).`when`(resources)
+                .getStringArray(R.array.config_networkSupportedKeepaliveCount)
+        doReturn(emptyArray<String>()).`when`(resources)
+                .getStringArray(R.array.config_networkNotifySwitches)
+        doReturn(intArrayOf(10, 11, 12, 14, 15)).`when`(resources)
+                .getIntArray(R.array.config_protectedNetworks)
+        // We don't test the actual notification value strings, so just return an empty array.
+        // It doesn't matter what the values are as long as it's not null.
+        doReturn(emptyArray<String>()).`when`(resources).getStringArray(
+                R.array.network_switch_type_name)
+        doReturn(1).`when`(resources).getInteger(R.integer.config_networkAvoidBadWifi)
+        doReturn(R.array.config_networkSupportedKeepaliveCount).`when`(resources)
+                .getIdentifier(eq("config_networkSupportedKeepaliveCount"), eq("array"), any())
+
+        doReturn(resources).`when`(resourcesContext).getResources()
+        ConnectivityResources.setResourcesContextForTest(resourcesContext)
+
         networkStackClient = TestNetworkStackClient(realContext)
         networkStackClient.start()
 
@@ -176,12 +208,19 @@
         doReturn(mock(ProxyTracker::class.java)).`when`(deps).makeProxyTracker(any(), any())
         doReturn(mock(MockableSystemProperties::class.java)).`when`(deps).systemProperties
         doReturn(TestNetIdManager()).`when`(deps).makeNetIdManager()
+        doAnswer { inv ->
+            object : MultinetworkPolicyTracker(inv.getArgument(0), inv.getArgument(1),
+                    inv.getArgument(2)) {
+                override fun getResourcesForActiveSubId() = resources
+            }
+        }.`when`(deps).makeMultinetworkPolicyTracker(any(), any(), any())
         return deps
     }
 
     @After
     fun tearDown() {
         nsInstrumentation.clearAllState()
+        ConnectivityResources.setResourcesContextForTest(null)
     }
 
     @Test
diff --git a/tests/unit/java/com/android/server/ConnectivityServiceTest.java b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
index 7034c92..c19ed61 100644
--- a/tests/unit/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
@@ -204,6 +204,7 @@
 import android.location.LocationManager;
 import android.net.CaptivePortalData;
 import android.net.ConnectionInfo;
+import android.net.ConnectivityDiagnosticsManager.DataStallReport;
 import android.net.ConnectivityManager;
 import android.net.ConnectivityManager.NetworkCallback;
 import android.net.ConnectivityManager.PacketKeepalive;
@@ -285,6 +286,7 @@
 import android.os.Parcel;
 import android.os.ParcelFileDescriptor;
 import android.os.Parcelable;
+import android.os.PersistableBundle;
 import android.os.Process;
 import android.os.RemoteException;
 import android.os.ServiceSpecificException;
@@ -10126,6 +10128,35 @@
         assertTrue(mService.mConnectivityDiagnosticsCallbacks.containsKey(mIBinder));
     }
 
+    @Test(expected = NullPointerException.class)
+    public void testRegisterConnectivityDiagnosticsCallbackNullCallback() {
+        mService.registerConnectivityDiagnosticsCallback(
+                null /* callback */,
+                new NetworkRequest.Builder().build(),
+                mContext.getPackageName());
+    }
+
+    @Test(expected = NullPointerException.class)
+    public void testRegisterConnectivityDiagnosticsCallbackNullNetworkRequest() {
+        mService.registerConnectivityDiagnosticsCallback(
+                mConnectivityDiagnosticsCallback,
+                null /* request */,
+                mContext.getPackageName());
+    }
+
+    @Test(expected = NullPointerException.class)
+    public void testRegisterConnectivityDiagnosticsCallbackNullPackageName() {
+        mService.registerConnectivityDiagnosticsCallback(
+                mConnectivityDiagnosticsCallback,
+                new NetworkRequest.Builder().build(),
+                null /* callingPackageName */);
+    }
+
+    @Test(expected = NullPointerException.class)
+    public void testUnregisterConnectivityDiagnosticsCallbackNullPackageName() {
+        mService.unregisterConnectivityDiagnosticsCallback(null /* callback */);
+    }
+
     public NetworkAgentInfo fakeMobileNai(NetworkCapabilities nc) {
         final NetworkCapabilities cellNc = new NetworkCapabilities.Builder(nc)
                 .addTransportType(TRANSPORT_CELLULAR).build();
@@ -10436,6 +10467,24 @@
                                 areConnDiagCapsRedacted(report.getNetworkCapabilities())));
     }
 
+    @Test(expected = NullPointerException.class)
+    public void testSimulateDataStallNullNetwork() {
+        mService.simulateDataStall(
+                DataStallReport.DETECTION_METHOD_DNS_EVENTS,
+                0L /* timestampMillis */,
+                null /* network */,
+                new PersistableBundle());
+    }
+
+    @Test(expected = NullPointerException.class)
+    public void testSimulateDataStallNullPersistableBundle() {
+        mService.simulateDataStall(
+                DataStallReport.DETECTION_METHOD_DNS_EVENTS,
+                0L /* timestampMillis */,
+                mock(Network.class),
+                null /* extras */);
+    }
+
     @Test
     public void testRouteAddDeleteUpdate() throws Exception {
         final NetworkRequest request = new NetworkRequest.Builder().build();