diff --git a/staticlibs/Android.bp b/staticlibs/Android.bp
index 03296e7..74aafd3 100644
--- a/staticlibs/Android.bp
+++ b/staticlibs/Android.bp
@@ -59,6 +59,7 @@
         "//packages/modules/NetworkStack:__subpackages__",
         "//packages/modules/CaptivePortalLogin",
         "//frameworks/libs/net/common/tests:__subpackages__",
+        "//frameworks/libs/net/common/testutils:__subpackages__",
   ],
   static_libs: [
       "net-utils-framework-common",
@@ -68,25 +69,6 @@
   ],
 }
 
-java_library {
-  // Consider using net-tests-utils instead if writing device code.
-  // That library has a lot more useful tools into it for users that
-  // work on Android and includes this lib.
-  name: "net-tests-utils-host-device-common",
-  srcs: [
-      "hostdevice/**/*.java",
-      "hostdevice/**/*.kt",
-  ],
-  host_supported: true,
-  visibility: [
-      "//frameworks/libs/net/common/tests:__subpackages__",
-      "//frameworks/libs/net/client-libs/tests:__subpackages__",
-  ],
-  static_libs: [
-      "kotlin-test"
-  ]
-}
-
 java_defaults {
     name: "lib_mockito_extended",
     static_libs: [
@@ -98,25 +80,6 @@
     ],
 }
 
-java_library {
-    name: "net-tests-utils",
-    srcs: [
-        "devicetests/**/*.java",
-        "devicetests/**/*.kt",
-    ],
-    defaults: ["lib_mockito_extended"],
-    libs: [
-        "androidx.annotation_annotation",
-    ],
-    static_libs: [
-        "androidx.test.ext.junit",
-        "kotlin-reflect",
-        "libnanohttpd",
-        "net-tests-utils-host-device-common",
-        "net-utils-device-common",
-    ],
-}
-
 filegroup {
     name: "net-utils-framework-common-srcs",
     srcs: ["framework/**/*.java"],
diff --git a/staticlibs/framework/com/android/net/module/util/NetworkCapabilitiesUtils.java b/staticlibs/framework/com/android/net/module/util/NetworkCapabilitiesUtils.java
index 3de78c6..4c7d675 100644
--- a/staticlibs/framework/com/android/net/module/util/NetworkCapabilitiesUtils.java
+++ b/staticlibs/framework/com/android/net/module/util/NetworkCapabilitiesUtils.java
@@ -16,6 +16,20 @@
 
 package com.android.net.module.util;
 
+import static android.net.NetworkCapabilities.NET_CAPABILITY_CBS;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_DUN;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_EIMS;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_FOTA;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_IA;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_IMS;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_INTERNET;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_MCX;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_MMS;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_OEM_PAID;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_RCS;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_SUPL;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_WIFI_P2P;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_XCAP;
 import static android.net.NetworkCapabilities.TRANSPORT_BLUETOOTH;
 import static android.net.NetworkCapabilities.TRANSPORT_CELLULAR;
 import static android.net.NetworkCapabilities.TRANSPORT_ETHERNET;
@@ -24,7 +38,9 @@
 import static android.net.NetworkCapabilities.TRANSPORT_WIFI_AWARE;
 
 import android.annotation.NonNull;
+import android.net.NetworkCapabilities;
 
+import com.android.internal.annotations.VisibleForTesting;
 
 /**
  * Utilities to examine {@link android.net.NetworkCapabilities}.
@@ -55,6 +71,83 @@
     };
 
     /**
+     * See android.net.NetworkCapabilities.NET_CAPABILITY_OEM_PRIVATE
+     * TODO: Use API constant when all downstream branches are S-based
+     */
+    public static final int NET_CAPABILITY_OEM_PRIVATE = 26;
+
+    /**
+     * See android.net.NetworkCapabilities.NET_CAPABILITY_VEHICLE_INTERNAL
+     * TODO: Use API constant when all downstream branches are S-based
+     */
+    public static final int NET_CAPABILITY_VEHICLE_INTERNAL = 27;
+
+    /**
+     * See android.net.NetworkCapabilities.NET_CAPABILITY_NOT_VCN_MANAGED
+     * TODO: Use API constant when all downstream branches are S-based
+     */
+    public static final int NET_CAPABILITY_NOT_VCN_MANAGED = 28;
+
+    /**
+     * See android.net.NetworkCapabilities.NET_CAPABILITY_ENTERPRISE
+     * TODO: Use API constant when all downstream branches are S-based
+     */
+    public static final int NET_CAPABILITY_ENTERPRISE = 29;
+
+    /**
+     * See android.net.NetworkCapabilities.NET_CAPABILITY_VSIM
+     * TODO: Use API constant when all downstream branches are S-based
+     */
+    public static final int NET_CAPABILITY_VSIM = 30;
+
+    /**
+     * See android.net.NetworkCapabilities.NET_CAPABILITY_BIP
+     * TODO: Use API constant when all downstream branches are S-based
+     */
+    public static final int NET_CAPABILITY_BIP = 31;
+
+
+    /**
+     * Capabilities that suggest that a network is restricted.
+     * See {@code NetworkCapabilities#maybeMarkCapabilitiesRestricted},
+      * and {@code FORCE_RESTRICTED_CAPABILITIES}.
+     */
+    @VisibleForTesting
+    static final long RESTRICTED_CAPABILITIES =
+            (1 << NET_CAPABILITY_BIP)
+            | (1 << NET_CAPABILITY_CBS)
+            | (1 << NET_CAPABILITY_DUN)
+            | (1 << NET_CAPABILITY_EIMS)
+            | (1 << NET_CAPABILITY_ENTERPRISE)
+            | (1 << NET_CAPABILITY_FOTA)
+            | (1 << NET_CAPABILITY_IA)
+            | (1 << NET_CAPABILITY_IMS)
+            | (1 << NET_CAPABILITY_MCX)
+            | (1 << NET_CAPABILITY_RCS)
+            | (1 << NET_CAPABILITY_VEHICLE_INTERNAL)
+            | (1 << NET_CAPABILITY_VSIM)
+            | (1 << NET_CAPABILITY_XCAP);
+
+    /**
+     * Capabilities that force network to be restricted.
+     * See {@code NetworkCapabilities#maybeMarkCapabilitiesRestricted}.
+     */
+    private static final long FORCE_RESTRICTED_CAPABILITIES =
+            (1 << NET_CAPABILITY_OEM_PAID)
+            | (1 << NET_CAPABILITY_OEM_PRIVATE);
+
+    /**
+     * Capabilities that suggest that a network is unrestricted.
+     * See {@code NetworkCapabilities#maybeMarkCapabilitiesRestricted}.
+     */
+    @VisibleForTesting
+    static final long UNRESTRICTED_CAPABILITIES =
+            (1 << NET_CAPABILITY_INTERNET)
+            | (1 << NET_CAPABILITY_MMS)
+            | (1 << NET_CAPABILITY_SUPL)
+            | (1 << NET_CAPABILITY_WIFI_P2P);
+
+    /**
      * Get a transport that can be used to classify a network when displaying its info to users.
      *
      * While networks can have multiple transports, users generally think of them as "wifi",
@@ -79,6 +172,41 @@
         return transports[0];
     }
 
+
+    /**
+     * Infers that all the capabilities it provides are typically provided by restricted networks
+     * or not.
+     *
+     * @param nc the {@link NetworkCapabilities} to infer the restricted capabilities.
+     *
+     * @return {@code true} if the network should be restricted.
+     */
+    // TODO: Use packBits(nc.getCapabilities()) to check more easily using bit masks.
+    public static boolean inferRestrictedCapability(NetworkCapabilities nc) {
+        // Check if we have any capability that forces the network to be restricted.
+        for (int capability : unpackBits(FORCE_RESTRICTED_CAPABILITIES)) {
+            if (nc.hasCapability(capability)) {
+                return true;
+            }
+        }
+
+        // Verify there aren't any unrestricted capabilities.  If there are we say
+        // the whole thing is unrestricted unless it is forced to be restricted.
+        for (int capability : unpackBits(UNRESTRICTED_CAPABILITIES)) {
+            if (nc.hasCapability(capability)) {
+                return false;
+            }
+        }
+
+        // Must have at least some restricted capabilities.
+        for (int capability : unpackBits(RESTRICTED_CAPABILITIES)) {
+            if (nc.hasCapability(capability)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     /**
      * Unpacks long value into an array of bits.
      */
diff --git a/staticlibs/framework/com/android/net/module/util/NetworkStackConstants.java b/staticlibs/framework/com/android/net/module/util/NetworkStackConstants.java
index 499297c..b7062e7 100644
--- a/staticlibs/framework/com/android/net/module/util/NetworkStackConstants.java
+++ b/staticlibs/framework/com/android/net/module/util/NetworkStackConstants.java
@@ -186,6 +186,27 @@
      */
     public static final int VENDOR_SPECIFIC_IE_ID = 0xdd;
 
+
+    /**
+     * TrafficStats constants.
+     */
+    // These tags are used by the network stack to do traffic for its own purposes. Traffic
+    // tagged with these will be counted toward the network stack and must stay inside the
+    // range defined by
+    // {@link android.net.TrafficStats#TAG_NETWORK_STACK_RANGE_START} and
+    // {@link android.net.TrafficStats#TAG_NETWORK_STACK_RANGE_END}.
+    public static final int TAG_SYSTEM_DHCP = 0xFFFFFE01;
+    public static final int TAG_SYSTEM_NEIGHBOR = 0xFFFFFE02;
+    public static final int TAG_SYSTEM_DHCP_SERVER = 0xFFFFFE03;
+
+    // These tags are used by the network stack to do traffic on behalf of apps. Traffic
+    // tagged with these will be counted toward the app on behalf of which the network
+    // stack is doing this traffic. These values must stay inside the range defined by
+    // {@link android.net.TrafficStats#TAG_NETWORK_STACK_IMPERSONATION_RANGE_START} and
+    // {@link android.net.TrafficStats#TAG_NETWORK_STACK_IMPERSONATION_RANGE_END}.
+    public static final int TAG_SYSTEM_PROBE = 0xFFFFFF81;
+    public static final int TAG_SYSTEM_DNS = 0xFFFFFF82;
+
     // TODO: Move to Inet4AddressUtils
     // See aosp/1455936: NetworkStackConstants can't depend on it as it causes jarjar-related issues
     // for users of both the net-utils-device-common and net-utils-framework-common libraries.
diff --git a/staticlibs/tests/unit/src/com/android/net/module/util/NetworkCapabilitiesUtilsTest.kt b/staticlibs/tests/unit/src/com/android/net/module/util/NetworkCapabilitiesUtilsTest.kt
index df2f459..5f15c6a 100644
--- a/staticlibs/tests/unit/src/com/android/net/module/util/NetworkCapabilitiesUtilsTest.kt
+++ b/staticlibs/tests/unit/src/com/android/net/module/util/NetworkCapabilitiesUtilsTest.kt
@@ -16,6 +16,12 @@
 
 package com.android.net.module.util
 
+import android.net.NetworkCapabilities
+import android.net.NetworkCapabilities.NET_CAPABILITY_CBS
+import android.net.NetworkCapabilities.NET_CAPABILITY_EIMS
+import android.net.NetworkCapabilities.NET_CAPABILITY_INTERNET
+import android.net.NetworkCapabilities.NET_CAPABILITY_INTERNET
+import android.net.NetworkCapabilities.NET_CAPABILITY_OEM_PAID
 import android.net.NetworkCapabilities.TRANSPORT_BLUETOOTH
 import android.net.NetworkCapabilities.TRANSPORT_CELLULAR
 import android.net.NetworkCapabilities.TRANSPORT_ETHERNET
@@ -25,6 +31,8 @@
 import android.net.NetworkCapabilities.TRANSPORT_WIFI_AWARE
 import androidx.test.filters.SmallTest
 import androidx.test.runner.AndroidJUnit4
+import com.android.net.module.util.NetworkCapabilitiesUtils.RESTRICTED_CAPABILITIES
+import com.android.net.module.util.NetworkCapabilitiesUtils.UNRESTRICTED_CAPABILITIES
 import com.android.net.module.util.NetworkCapabilitiesUtils.getDisplayTransport
 import com.android.net.module.util.NetworkCapabilitiesUtils.packBits
 import com.android.net.module.util.NetworkCapabilitiesUtils.unpackBits
@@ -33,6 +41,7 @@
 import java.lang.IllegalArgumentException
 import kotlin.test.assertEquals
 import kotlin.test.assertFailsWith
+import kotlin.test.assertFalse
 import kotlin.test.assertTrue
 
 @RunWith(AndroidJUnit4::class)
@@ -78,4 +87,41 @@
         assertEquals(packedBits, packBits(bits))
         assertTrue(bits contentEquals unpackBits(packedBits))
     }
+
+    @Test
+    fun testInferRestrictedCapability() {
+        val nc = NetworkCapabilities()
+        // Default capabilities don't have restricted capability.
+        assertFalse(NetworkCapabilitiesUtils.inferRestrictedCapability(nc))
+        // If there is a force restricted capability, then the network capabilities is restricted.
+        nc.addCapability(NET_CAPABILITY_OEM_PAID)
+        nc.addCapability(NET_CAPABILITY_INTERNET)
+        assertTrue(NetworkCapabilitiesUtils.inferRestrictedCapability(nc))
+        // Except for the force restricted capability, if there is any unrestricted capability in
+        // capabilities, then the network capabilities is not restricted.
+        nc.removeCapability(NET_CAPABILITY_OEM_PAID)
+        nc.addCapability(NET_CAPABILITY_CBS)
+        assertFalse(NetworkCapabilitiesUtils.inferRestrictedCapability(nc))
+        // Except for the force restricted capability, the network capabilities will only be treated
+        // as restricted when there is no any unrestricted capability.
+        nc.removeCapability(NET_CAPABILITY_INTERNET)
+        assertTrue(NetworkCapabilitiesUtils.inferRestrictedCapability(nc))
+    }
+
+    @Test
+    fun testRestrictedUnrestrictedCapabilities() {
+        // verify EIMS is restricted
+        assertEquals((1 shl NET_CAPABILITY_EIMS).toLong() and RESTRICTED_CAPABILITIES,
+                (1 shl NET_CAPABILITY_EIMS).toLong())
+
+        // verify CBS is also restricted
+        assertEquals((1 shl NET_CAPABILITY_CBS).toLong() and RESTRICTED_CAPABILITIES,
+                (1 shl NET_CAPABILITY_CBS).toLong())
+
+        // verify default is not restricted
+        assertEquals((1 shl NET_CAPABILITY_INTERNET).toLong() and RESTRICTED_CAPABILITIES, 0)
+
+        // just to see
+        assertEquals(RESTRICTED_CAPABILITIES and UNRESTRICTED_CAPABILITIES, 0)
+    }
 }
diff --git a/staticlibs/testutils/Android.bp b/staticlibs/testutils/Android.bp
new file mode 100644
index 0000000..d4465dd
--- /dev/null
+++ b/staticlibs/testutils/Android.bp
@@ -0,0 +1,55 @@
+// Copyright (C) 2021 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+java_library {
+    name: "net-tests-utils",
+    srcs: [
+        "devicetests/**/*.java",
+        "devicetests/**/*.kt",
+    ],
+    defaults: ["lib_mockito_extended"],
+    libs: [
+        "androidx.annotation_annotation",
+    ],
+    static_libs: [
+        "androidx.test.ext.junit",
+        "kotlin-reflect",
+        "libnanohttpd",
+        "net-tests-utils-host-device-common",
+        "net-utils-device-common",
+    ],
+}
+
+java_library {
+  // Consider using net-tests-utils instead if writing device code.
+  // That library has a lot more useful tools into it for users that
+  // work on Android and includes this lib.
+  name: "net-tests-utils-host-device-common",
+  srcs: [
+      "hostdevice/**/*.java",
+      "hostdevice/**/*.kt",
+  ],
+  host_supported: true,
+  visibility: [
+      "//frameworks/libs/net/common/tests:__subpackages__",
+      "//frameworks/libs/net/client-libs/tests:__subpackages__",
+  ],
+  static_libs: [
+      "kotlin-test"
+  ]
+}
diff --git a/staticlibs/devicetests/com/android/testutils/ArpResponder.kt b/staticlibs/testutils/devicetests/com/android/testutils/ArpResponder.kt
similarity index 100%
rename from staticlibs/devicetests/com/android/testutils/ArpResponder.kt
rename to staticlibs/testutils/devicetests/com/android/testutils/ArpResponder.kt
diff --git a/staticlibs/devicetests/com/android/testutils/CompatUtil.kt b/staticlibs/testutils/devicetests/com/android/testutils/CompatUtil.kt
similarity index 100%
rename from staticlibs/devicetests/com/android/testutils/CompatUtil.kt
rename to staticlibs/testutils/devicetests/com/android/testutils/CompatUtil.kt
diff --git a/staticlibs/devicetests/com/android/testutils/ConcurrentInterpreter.kt b/staticlibs/testutils/devicetests/com/android/testutils/ConcurrentInterpreter.kt
similarity index 100%
rename from staticlibs/devicetests/com/android/testutils/ConcurrentInterpreter.kt
rename to staticlibs/testutils/devicetests/com/android/testutils/ConcurrentInterpreter.kt
diff --git a/staticlibs/devicetests/com/android/testutils/DevSdkIgnoreRule.kt b/staticlibs/testutils/devicetests/com/android/testutils/DevSdkIgnoreRule.kt
similarity index 100%
rename from staticlibs/devicetests/com/android/testutils/DevSdkIgnoreRule.kt
rename to staticlibs/testutils/devicetests/com/android/testutils/DevSdkIgnoreRule.kt
diff --git a/staticlibs/devicetests/com/android/testutils/DevSdkIgnoreRunner.kt b/staticlibs/testutils/devicetests/com/android/testutils/DevSdkIgnoreRunner.kt
similarity index 100%
rename from staticlibs/devicetests/com/android/testutils/DevSdkIgnoreRunner.kt
rename to staticlibs/testutils/devicetests/com/android/testutils/DevSdkIgnoreRunner.kt
diff --git a/staticlibs/devicetests/com/android/testutils/FakeDns.kt b/staticlibs/testutils/devicetests/com/android/testutils/FakeDns.kt
similarity index 100%
rename from staticlibs/devicetests/com/android/testutils/FakeDns.kt
rename to staticlibs/testutils/devicetests/com/android/testutils/FakeDns.kt
diff --git a/staticlibs/devicetests/com/android/testutils/HandlerUtils.kt b/staticlibs/testutils/devicetests/com/android/testutils/HandlerUtils.kt
similarity index 100%
rename from staticlibs/devicetests/com/android/testutils/HandlerUtils.kt
rename to staticlibs/testutils/devicetests/com/android/testutils/HandlerUtils.kt
diff --git a/staticlibs/testutils/devicetests/com/android/testutils/NetworkStatsProviderCbStubCompat.java b/staticlibs/testutils/devicetests/com/android/testutils/NetworkStatsProviderCbStubCompat.java
new file mode 100644
index 0000000..e84a224
--- /dev/null
+++ b/staticlibs/testutils/devicetests/com/android/testutils/NetworkStatsProviderCbStubCompat.java
@@ -0,0 +1,43 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.testutils;
+
+import android.net.NetworkStats;
+import android.net.netstats.provider.INetworkStatsProviderCallback;
+import android.os.RemoteException;
+
+/**
+ * A shim class that allows {@link TestableNetworkStatsProviderCbBinder} to be built against
+ * different SDK versions.
+ */
+public class NetworkStatsProviderCbStubCompat extends INetworkStatsProviderCallback.Stub {
+    @Override
+    public void notifyStatsUpdated(int token, NetworkStats ifaceStats, NetworkStats uidStats)
+            throws RemoteException {}
+
+    @Override
+    public void notifyAlertReached() throws RemoteException {}
+
+    // Removed in S.
+    public void notifyLimitReached() throws RemoteException {}
+
+    // Added in S.
+    public void notifyWarningOrLimitReached() throws RemoteException {}
+
+    @Override
+    public void unregister() throws RemoteException {}
+}
diff --git a/staticlibs/testutils/devicetests/com/android/testutils/NetworkStatsProviderStubCompat.java b/staticlibs/testutils/devicetests/com/android/testutils/NetworkStatsProviderStubCompat.java
new file mode 100644
index 0000000..a77aa02
--- /dev/null
+++ b/staticlibs/testutils/devicetests/com/android/testutils/NetworkStatsProviderStubCompat.java
@@ -0,0 +1,37 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.testutils;
+
+import android.net.netstats.provider.INetworkStatsProvider;
+
+/**
+ * A shim class that allows {@link TestableNetworkStatsProviderBinder} to be built against
+ * different SDK versions.
+ */
+public class NetworkStatsProviderStubCompat extends INetworkStatsProvider.Stub {
+    @Override
+    public void onRequestStatsUpdate(int token) {}
+
+    // Removed and won't be called in S+.
+    public void onSetLimit(String iface, long quotaBytes) {}
+
+    @Override
+    public void onSetAlert(long bytes) {}
+
+    // Added in S.
+    public void onSetWarningAndLimit(String iface, long warningBytes, long limitBytes) {}
+}
diff --git a/staticlibs/devicetests/com/android/testutils/NetworkStatsUtils.kt b/staticlibs/testutils/devicetests/com/android/testutils/NetworkStatsUtils.kt
similarity index 100%
rename from staticlibs/devicetests/com/android/testutils/NetworkStatsUtils.kt
rename to staticlibs/testutils/devicetests/com/android/testutils/NetworkStatsUtils.kt
diff --git a/staticlibs/devicetests/com/android/testutils/PacketResponder.kt b/staticlibs/testutils/devicetests/com/android/testutils/PacketResponder.kt
similarity index 100%
rename from staticlibs/devicetests/com/android/testutils/PacketResponder.kt
rename to staticlibs/testutils/devicetests/com/android/testutils/PacketResponder.kt
diff --git a/staticlibs/devicetests/com/android/testutils/ParcelUtils.kt b/staticlibs/testutils/devicetests/com/android/testutils/ParcelUtils.kt
similarity index 100%
rename from staticlibs/devicetests/com/android/testutils/ParcelUtils.kt
rename to staticlibs/testutils/devicetests/com/android/testutils/ParcelUtils.kt
diff --git a/staticlibs/devicetests/com/android/testutils/TapPacketReader.java b/staticlibs/testutils/devicetests/com/android/testutils/TapPacketReader.java
similarity index 100%
rename from staticlibs/devicetests/com/android/testutils/TapPacketReader.java
rename to staticlibs/testutils/devicetests/com/android/testutils/TapPacketReader.java
diff --git a/staticlibs/devicetests/com/android/testutils/TapPacketReaderRule.kt b/staticlibs/testutils/devicetests/com/android/testutils/TapPacketReaderRule.kt
similarity index 100%
rename from staticlibs/devicetests/com/android/testutils/TapPacketReaderRule.kt
rename to staticlibs/testutils/devicetests/com/android/testutils/TapPacketReaderRule.kt
diff --git a/staticlibs/devicetests/com/android/testutils/TestHttpServer.kt b/staticlibs/testutils/devicetests/com/android/testutils/TestHttpServer.kt
similarity index 100%
rename from staticlibs/devicetests/com/android/testutils/TestHttpServer.kt
rename to staticlibs/testutils/devicetests/com/android/testutils/TestHttpServer.kt
diff --git a/staticlibs/devicetests/com/android/testutils/TestNetworkTracker.kt b/staticlibs/testutils/devicetests/com/android/testutils/TestNetworkTracker.kt
similarity index 100%
rename from staticlibs/devicetests/com/android/testutils/TestNetworkTracker.kt
rename to staticlibs/testutils/devicetests/com/android/testutils/TestNetworkTracker.kt
diff --git a/staticlibs/devicetests/com/android/testutils/TestPermissionUtil.kt b/staticlibs/testutils/devicetests/com/android/testutils/TestPermissionUtil.kt
similarity index 100%
rename from staticlibs/devicetests/com/android/testutils/TestPermissionUtil.kt
rename to staticlibs/testutils/devicetests/com/android/testutils/TestPermissionUtil.kt
diff --git a/staticlibs/devicetests/com/android/testutils/TestableNetworkCallback.kt b/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkCallback.kt
similarity index 90%
rename from staticlibs/devicetests/com/android/testutils/TestableNetworkCallback.kt
rename to staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkCallback.kt
index 043654f..1b3d0f6 100644
--- a/staticlibs/devicetests/com/android/testutils/TestableNetworkCallback.kt
+++ b/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkCallback.kt
@@ -25,6 +25,7 @@
 import com.android.net.module.util.ArrayTrackRecord
 import com.android.testutils.RecorderCallback.CallbackEntry.Available
 import com.android.testutils.RecorderCallback.CallbackEntry.BlockedStatus
+import com.android.testutils.RecorderCallback.CallbackEntry.BlockedStatusInt
 import com.android.testutils.RecorderCallback.CallbackEntry.CapabilitiesChanged
 import com.android.testutils.RecorderCallback.CallbackEntry.LinkPropertiesChanged
 import com.android.testutils.RecorderCallback.CallbackEntry.Losing
@@ -82,7 +83,10 @@
             override val network: Network,
             val blocked: Boolean
         ) : CallbackEntry()
-
+        data class BlockedStatusInt(
+            override val network: Network,
+            val blocked: Int
+        ) : CallbackEntry()
         // Convenience constants for expecting a type
         companion object {
             @JvmField
@@ -103,6 +107,8 @@
             val UNAVAILABLE = Unavailable::class
             @JvmField
             val BLOCKED_STATUS = BlockedStatus::class
+            @JvmField
+            val BLOCKED_STATUS_INT = BlockedStatusInt::class
         }
     }
 
@@ -131,6 +137,9 @@
         history.add(BlockedStatus(network, blocked))
     }
 
+    // Cannot do:
+    // fun onBlockedStatusChanged(network: Network, blocked: Int) {
+    // because on S, that needs to be "override fun", and on R, that cannot be "override fun".
     override fun onNetworkSuspended(network: Network) {
         Log.d(TAG, "onNetworkSuspended $network $network")
         history.add(Suspended(network))
@@ -272,13 +281,33 @@
         blocked: Boolean = false,
         tmt: Long = defaultTimeoutMs
     ) {
+        expectAvailableCallbacksCommon(net, suspended, validated, tmt)
+        expectBlockedStatusCallback(blocked, net, tmt)
+    }
+
+    fun expectAvailableCallbacks(
+        net: Network,
+        suspended: Boolean,
+        validated: Boolean,
+        blockedStatus: Int,
+        tmt: Long
+    ) {
+        expectAvailableCallbacksCommon(net, suspended, validated, tmt)
+        expectBlockedStatusCallback(blockedStatus, net)
+    }
+
+    private fun expectAvailableCallbacksCommon(
+        net: Network,
+        suspended: Boolean,
+        validated: Boolean,
+        tmt: Long
+    ) {
         expectCallback<Available>(net, tmt)
         if (suspended) {
             expectCallback<Suspended>(net, tmt)
         }
         expectCapabilitiesThat(net, tmt) { validated == it.hasCapability(NET_CAPABILITY_VALIDATED) }
         expectCallback<LinkPropertiesChanged>(net, tmt)
-        expectBlockedStatusCallback(blocked, net)
     }
 
     // Backward compatibility for existing Java code. Use named arguments instead and remove all
@@ -295,6 +324,12 @@
         }
     }
 
+    fun expectBlockedStatusCallback(blocked: Int, net: Network, tmt: Long = defaultTimeoutMs) {
+        expectCallback<BlockedStatusInt>(net, tmt).also {
+            assertEquals(it.blocked, blocked, "Unexpected blocked status ${it.blocked}")
+        }
+    }
+
     // Expects the available callbacks (where the onCapabilitiesChanged must contain the
     // VALIDATED capability), plus another onCapabilitiesChanged which is identical to the
     // one we just sent.
@@ -314,6 +349,16 @@
         expectCapabilitiesThat(net, tmt) { it.hasCapability(NET_CAPABILITY_VALIDATED) }
     }
 
+    fun expectAvailableThenValidatedCallbacks(
+        net: Network,
+        blockedStatus: Int,
+        tmt: Long = defaultTimeoutMs
+    ) {
+        expectAvailableCallbacks(net, validated = false, suspended = false,
+                blockedStatus = blockedStatus, tmt = tmt)
+        expectCapabilitiesThat(net, tmt) { it.hasCapability(NET_CAPABILITY_VALIDATED) }
+    }
+
     // Temporary Java compat measure : have MockNetworkAgent implement this so that all existing
     // calls with networkAgent can be routed through here without moving MockNetworkAgent.
     // TODO: clean this up, remove this method.
diff --git a/staticlibs/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt b/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt
similarity index 100%
rename from staticlibs/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt
rename to staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt
diff --git a/staticlibs/devicetests/com/android/testutils/TestableNetworkStatsProviderBinder.kt b/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProviderBinder.kt
similarity index 73%
rename from staticlibs/devicetests/com/android/testutils/TestableNetworkStatsProviderBinder.kt
rename to staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProviderBinder.kt
index 02922d8..643346b 100644
--- a/staticlibs/devicetests/com/android/testutils/TestableNetworkStatsProviderBinder.kt
+++ b/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProviderBinder.kt
@@ -16,18 +16,21 @@
 
 package com.android.testutils
 
-import android.net.netstats.provider.INetworkStatsProvider
 import com.android.net.module.util.ArrayTrackRecord
 import kotlin.test.assertEquals
 import kotlin.test.fail
 
 private const val DEFAULT_TIMEOUT_MS = 200L
 
-open class TestableNetworkStatsProviderBinder : INetworkStatsProvider.Stub() {
+open class TestableNetworkStatsProviderBinder : NetworkStatsProviderStubCompat() {
     sealed class CallbackType {
         data class OnRequestStatsUpdate(val token: Int) : CallbackType()
-        data class OnSetLimit(val iface: String?, val quotaBytes: Long) : CallbackType()
         data class OnSetAlert(val quotaBytes: Long) : CallbackType()
+        data class OnSetWarningAndLimit(
+            val iface: String,
+            val warningBytes: Long,
+            val limitBytes: Long
+        ) : CallbackType()
     }
 
     private val history = ArrayTrackRecord<CallbackType>().ReadHead()
@@ -36,20 +39,21 @@
         history.add(CallbackType.OnRequestStatsUpdate(token))
     }
 
-    override fun onSetLimit(iface: String?, quotaBytes: Long) {
-        history.add(CallbackType.OnSetLimit(iface, quotaBytes))
-    }
-
     override fun onSetAlert(quotaBytes: Long) {
         history.add(CallbackType.OnSetAlert(quotaBytes))
     }
 
+    override fun onSetWarningAndLimit(iface: String, warningBytes: Long, limitBytes: Long) {
+        history.add(CallbackType.OnSetWarningAndLimit(iface, warningBytes, limitBytes))
+    }
+
     fun expectOnRequestStatsUpdate(token: Int) {
         assertEquals(CallbackType.OnRequestStatsUpdate(token), history.poll(DEFAULT_TIMEOUT_MS))
     }
 
-    fun expectOnSetLimit(iface: String?, quotaBytes: Long) {
-        assertEquals(CallbackType.OnSetLimit(iface, quotaBytes), history.poll(DEFAULT_TIMEOUT_MS))
+    fun expectOnSetWarningAndLimit(iface: String, warningBytes: Long, limitBytes: Long) {
+        assertEquals(CallbackType.OnSetWarningAndLimit(iface, warningBytes, limitBytes),
+                history.poll(DEFAULT_TIMEOUT_MS))
     }
 
     fun expectOnSetAlert(quotaBytes: Long) {
diff --git a/staticlibs/devicetests/com/android/testutils/TestableNetworkStatsProviderCbBinder.kt b/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProviderCbBinder.kt
similarity index 85%
rename from staticlibs/devicetests/com/android/testutils/TestableNetworkStatsProviderCbBinder.kt
rename to staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProviderCbBinder.kt
index 163473a..f15f610 100644
--- a/staticlibs/devicetests/com/android/testutils/TestableNetworkStatsProviderCbBinder.kt
+++ b/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProviderCbBinder.kt
@@ -17,7 +17,6 @@
 package com.android.testutils
 
 import android.net.NetworkStats
-import android.net.netstats.provider.INetworkStatsProviderCallback
 import com.android.net.module.util.ArrayTrackRecord
 import kotlin.test.assertEquals
 import kotlin.test.assertTrue
@@ -25,14 +24,14 @@
 
 private const val DEFAULT_TIMEOUT_MS = 3000L
 
-open class TestableNetworkStatsProviderCbBinder : INetworkStatsProviderCallback.Stub() {
+open class TestableNetworkStatsProviderCbBinder : NetworkStatsProviderCbStubCompat() {
     sealed class CallbackType {
         data class NotifyStatsUpdated(
             val token: Int,
             val ifaceStats: NetworkStats,
             val uidStats: NetworkStats
         ) : CallbackType()
-        object NotifyLimitReached : CallbackType()
+        object NotifyWarningOrLimitReached : CallbackType()
         object NotifyAlertReached : CallbackType()
         object Unregister : CallbackType()
     }
@@ -43,8 +42,8 @@
         history.add(CallbackType.NotifyStatsUpdated(token, ifaceStats, uidStats))
     }
 
-    override fun notifyLimitReached() {
-        history.add(CallbackType.NotifyLimitReached)
+    override fun notifyWarningOrLimitReached() {
+        history.add(CallbackType.NotifyWarningOrLimitReached)
     }
 
     override fun notifyAlertReached() {
@@ -70,8 +69,8 @@
         assertNetworkStatsEquals(uidStats, event.uidStats)
     }
 
-    fun expectNotifyLimitReached() =
-            assertEquals(CallbackType.NotifyLimitReached, history.poll(DEFAULT_TIMEOUT_MS))
+    fun expectNotifyWarningOrLimitReached() =
+            assertEquals(CallbackType.NotifyWarningOrLimitReached, history.poll(DEFAULT_TIMEOUT_MS))
 
     fun expectNotifyAlertReached() =
             assertEquals(CallbackType.NotifyAlertReached, history.poll(DEFAULT_TIMEOUT_MS))
diff --git a/staticlibs/hostdevice/com/android/net/module/util/TrackRecord.kt b/staticlibs/testutils/hostdevice/com/android/net/module/util/TrackRecord.kt
similarity index 100%
rename from staticlibs/hostdevice/com/android/net/module/util/TrackRecord.kt
rename to staticlibs/testutils/hostdevice/com/android/net/module/util/TrackRecord.kt
diff --git a/staticlibs/hostdevice/com/android/testutils/ConcurrentUtils.kt b/staticlibs/testutils/hostdevice/com/android/testutils/ConcurrentUtils.kt
similarity index 100%
rename from staticlibs/hostdevice/com/android/testutils/ConcurrentUtils.kt
rename to staticlibs/testutils/hostdevice/com/android/testutils/ConcurrentUtils.kt
diff --git a/staticlibs/hostdevice/com/android/testutils/ExceptionUtils.java b/staticlibs/testutils/hostdevice/com/android/testutils/ExceptionUtils.java
similarity index 100%
rename from staticlibs/hostdevice/com/android/testutils/ExceptionUtils.java
rename to staticlibs/testutils/hostdevice/com/android/testutils/ExceptionUtils.java
diff --git a/staticlibs/hostdevice/com/android/testutils/FileUtils.kt b/staticlibs/testutils/hostdevice/com/android/testutils/FileUtils.kt
similarity index 100%
rename from staticlibs/hostdevice/com/android/testutils/FileUtils.kt
rename to staticlibs/testutils/hostdevice/com/android/testutils/FileUtils.kt
diff --git a/staticlibs/hostdevice/com/android/testutils/MiscAsserts.kt b/staticlibs/testutils/hostdevice/com/android/testutils/MiscAsserts.kt
similarity index 100%
rename from staticlibs/hostdevice/com/android/testutils/MiscAsserts.kt
rename to staticlibs/testutils/hostdevice/com/android/testutils/MiscAsserts.kt
diff --git a/staticlibs/hostdevice/com/android/testutils/PacketFilter.kt b/staticlibs/testutils/hostdevice/com/android/testutils/PacketFilter.kt
similarity index 100%
rename from staticlibs/hostdevice/com/android/testutils/PacketFilter.kt
rename to staticlibs/testutils/hostdevice/com/android/testutils/PacketFilter.kt
diff --git a/staticlibs/hostdevice/com/android/testutils/SkipPresubmit.kt b/staticlibs/testutils/hostdevice/com/android/testutils/SkipPresubmit.kt
similarity index 100%
rename from staticlibs/hostdevice/com/android/testutils/SkipPresubmit.kt
rename to staticlibs/testutils/hostdevice/com/android/testutils/SkipPresubmit.kt