[FUI17] Migrate NetworkStatsService to use NetworkStateSnapshot

This change migrates NetworkStatsService and related code to use
NetworkStateSnapshot which is used for replacing the NetworkState.

This patch also changes some formating which is suggested by
the linter.

Test: FrameworksNetTests NetworkPolicyManagerServiceTest
Bug: 174123988
Change-Id: I547da8f411cb45bdadc376ac3cadf3f3c55bb282

services/core/java/com/android/server/ConnectivityService.java

@@ -120,6 +120,7 @@ import android.net.NetworkSpecifier;
 import android.net.NetworkStack;
 import android.net.NetworkStackClient;
 import android.net.NetworkState;
+import android.net.NetworkStateSnapshot;
 import android.net.NetworkTestResultParcelable;
 import android.net.NetworkUtils;
 import android.net.NetworkWatchlistManager;
@@ -7948,8 +7949,16 @@ public class ConnectivityService extends IConnectivityManager.Stub
 
         final UnderlyingNetworkInfo[] underlyingNetworkInfos = getAllVpnInfo();
         try {
-            mStatsService.forceUpdateIfaces(getDefaultNetworks(), getAllNetworkState(), activeIface,
+            final ArrayList<NetworkStateSnapshot> snapshots = new ArrayList<>();
-                    underlyingNetworkInfos);
+            // TODO: Directly use NetworkStateSnapshot when feasible.
+            for (final NetworkState state : getAllNetworkState()) {
+                final NetworkStateSnapshot snapshot = new NetworkStateSnapshot(state.linkProperties,
+                        state.networkCapabilities, state.network, state.subscriberId,
+                        state.legacyNetworkType);
+                snapshots.add(snapshot);
+            }
+            mStatsService.forceUpdateIfaces(getDefaultNetworks(), snapshots.toArray(
+                    new NetworkStateSnapshot[0]), activeIface, underlyingNetworkInfos);
         } catch (Exception ignored) {
         }
     }

tests/net/java/android/net/NetworkTemplateTest.kt

@@ -20,14 +20,13 @@ import android.content.Context
 import android.net.ConnectivityManager.TYPE_MOBILE
 import android.net.ConnectivityManager.TYPE_WIFI
 import android.net.NetworkIdentity.SUBTYPE_COMBINED
-import android.net.NetworkIdentity.OEM_NONE;
+import android.net.NetworkIdentity.OEM_NONE
-import android.net.NetworkIdentity.OEM_PAID;
+import android.net.NetworkIdentity.OEM_PAID
-import android.net.NetworkIdentity.OEM_PRIVATE;
+import android.net.NetworkIdentity.OEM_PRIVATE
 import android.net.NetworkIdentity.buildNetworkIdentity
 import android.net.NetworkStats.DEFAULT_NETWORK_ALL
 import android.net.NetworkStats.METERED_ALL
 import android.net.NetworkStats.ROAMING_ALL
-import android.net.NetworkTemplate.MATCH_ETHERNET
 import android.net.NetworkTemplate.MATCH_MOBILE
 import android.net.NetworkTemplate.MATCH_MOBILE_WILDCARD
 import android.net.NetworkTemplate.MATCH_WIFI
@@ -50,7 +49,6 @@ import kotlin.test.assertEquals
 import kotlin.test.assertFalse
 import kotlin.test.assertNotEquals
 import kotlin.test.assertTrue
-import kotlin.test.fail
 
 private const val TEST_IMSI1 = "imsi1"
 private const val TEST_IMSI2 = "imsi2"
@@ -60,17 +58,17 @@ private const val TEST_SSID1 = "ssid1"
 class NetworkTemplateTest {
     private val mockContext = mock(Context::class.java)
 
-    private fun buildMobileNetworkState(subscriberId: String): NetworkState =
+    private fun buildMobileNetworkState(subscriberId: String): NetworkStateSnapshot =
             buildNetworkState(TYPE_MOBILE, subscriberId = subscriberId)
-    private fun buildWifiNetworkState(ssid: String): NetworkState =
+    private fun buildWifiNetworkState(ssid: String): NetworkStateSnapshot =
             buildNetworkState(TYPE_WIFI, ssid = ssid)
 
     private fun buildNetworkState(
         type: Int,
         subscriberId: String? = null,
         ssid: String? = null,
-        oemManaged: Int = OEM_NONE,
+        oemManaged: Int = OEM_NONE
-    ): NetworkState {
+    ): NetworkStateSnapshot {
         val lp = LinkProperties()
         val caps = NetworkCapabilities().apply {
             setCapability(NetworkCapabilities.NET_CAPABILITY_NOT_METERED, false)
@@ -81,7 +79,7 @@ class NetworkTemplateTest {
             setCapability(NetworkCapabilities.NET_CAPABILITY_OEM_PRIVATE,
                     (oemManaged and OEM_PRIVATE) == OEM_PRIVATE)
         }
-        return NetworkState(type, lp, caps, mock(Network::class.java), subscriberId)
+        return NetworkStateSnapshot(lp, caps, mock(Network::class.java), subscriberId, type)
     }
 
     private fun NetworkTemplate.assertMatches(ident: NetworkIdentity) =
@@ -179,7 +177,7 @@ class NetworkTemplateTest {
                 OEM_PAID, OEM_PRIVATE, OEM_PAID or OEM_PRIVATE)
 
         // Verify that "not OEM managed network" constants are equal.
-        assertEquals(OEM_MANAGED_NO, OEM_NONE);
+        assertEquals(OEM_MANAGED_NO, OEM_NONE)
 
         // Verify the constants don't conflict.
         assertEquals(constantValues.size, constantValues.distinct().count())
@@ -201,8 +199,13 @@ class NetworkTemplateTest {
      * @param identSsid If networkType is {@code TYPE_WIFI}, this value must *NOT* be null. Provide
      *         one of {@code TEST_SSID*}.
      */
-    private fun matchOemManagedIdent(networkType: Int, matchType:Int, subscriberId: String? = null,
+    private fun matchOemManagedIdent(
-            templateSsid: String? = null, identSsid: String? = null) {
+        networkType: Int,
+        matchType: Int,
+        subscriberId: String? = null,
+        templateSsid: String? = null,
+        identSsid: String? = null
+    ) {
         val oemManagedStates = arrayOf(OEM_NONE, OEM_PAID, OEM_PRIVATE, OEM_PAID or OEM_PRIVATE)
         // A null subscriberId needs a null matchSubscriberIds argument as well.
         val matchSubscriberIds = if (subscriberId == null) null else arrayOf(subscriberId)

tests/net/java/com/android/server/ConnectivityServiceTest.java

@@ -198,7 +198,7 @@ import android.net.NetworkRequest;
 import android.net.NetworkSpecifier;
 import android.net.NetworkStack;
 import android.net.NetworkStackClient;
-import android.net.NetworkState;
+import android.net.NetworkStateSnapshot;
 import android.net.NetworkTestResultParcelable;
 import android.net.OemNetworkPreferences;
 import android.net.ProxyInfo;
@@ -5484,7 +5484,7 @@ public class ConnectivityServiceTest {
                 UnderlyingNetworkInfo[].class);
 
         verify(mStatsService, atLeastOnce()).forceUpdateIfaces(networksCaptor.capture(),
-                any(NetworkState[].class), eq(defaultIface), vpnInfosCaptor.capture());
+                any(NetworkStateSnapshot[].class), eq(defaultIface), vpnInfosCaptor.capture());
 
         assertSameElementsNoDuplicates(networksCaptor.getValue(), networks);
 
@@ -5554,9 +5554,8 @@ public class ConnectivityServiceTest {
         // Temp metered change shouldn't update ifaces
         mCellNetworkAgent.addCapability(NetworkCapabilities.NET_CAPABILITY_TEMPORARILY_NOT_METERED);
         waitForIdle();
-        verify(mStatsService, never())
+        verify(mStatsService, never()).forceUpdateIfaces(eq(onlyCell), any(
-                .forceUpdateIfaces(eq(onlyCell), any(NetworkState[].class), eq(MOBILE_IFNAME),
+                NetworkStateSnapshot[].class), eq(MOBILE_IFNAME), eq(new UnderlyingNetworkInfo[0]));
-                        eq(new UnderlyingNetworkInfo[0]));
         reset(mStatsService);
 
         // Roaming change should update ifaces
@@ -5638,7 +5637,7 @@ public class ConnectivityServiceTest {
         // Confirm that we never tell NetworkStatsService that cell is no longer the underlying
         // network for the VPN...
         verify(mStatsService, never()).forceUpdateIfaces(any(Network[].class),
-                any(NetworkState[].class), any() /* anyString() doesn't match null */,
+                any(NetworkStateSnapshot[].class), any() /* anyString() doesn't match null */,
                 argThat(infos -> infos[0].underlyingIfaces.size() == 1
                         && WIFI_IFNAME.equals(infos[0].underlyingIfaces.get(0))));
         verifyNoMoreInteractions(mStatsService);
@@ -5652,7 +5651,7 @@ public class ConnectivityServiceTest {
         mEthernetNetworkAgent.connect(false);
         waitForIdle();
         verify(mStatsService).forceUpdateIfaces(any(Network[].class),
-                any(NetworkState[].class), any() /* anyString() doesn't match null */,
+                any(NetworkStateSnapshot[].class), any() /* anyString() doesn't match null */,
                 argThat(vpnInfos -> vpnInfos[0].underlyingIfaces.size() == 1
                         && WIFI_IFNAME.equals(vpnInfos[0].underlyingIfaces.get(0))));
         mEthernetNetworkAgent.disconnect();

tests/net/java/com/android/server/net/NetworkStatsServiceTest.java

@@ -19,9 +19,7 @@ package com.android.server.net;
 import static android.content.Intent.ACTION_UID_REMOVED;
 import static android.content.Intent.EXTRA_UID;
 import static android.net.ConnectivityManager.TYPE_MOBILE;
-import static android.net.ConnectivityManager.TYPE_VPN;
 import static android.net.ConnectivityManager.TYPE_WIFI;
-import static android.net.NetworkIdentity.OEM_NONE;
 import static android.net.NetworkIdentity.OEM_PAID;
 import static android.net.NetworkIdentity.OEM_PRIVATE;
 import static android.net.NetworkStats.DEFAULT_NETWORK_ALL;
@@ -86,7 +84,7 @@ import android.net.INetworkStatsSession;
 import android.net.LinkProperties;
 import android.net.Network;
 import android.net.NetworkCapabilities;
-import android.net.NetworkState;
+import android.net.NetworkStateSnapshot;
 import android.net.NetworkStats;
 import android.net.NetworkStatsHistory;
 import android.net.NetworkTemplate;
@@ -286,7 +284,7 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
         // pretend that wifi network comes online; service should ask about full
         // network state, and poll any existing interfaces before updating.
         expectDefaultSettings();
-        NetworkState[] states = new NetworkState[] {buildWifiState()};
+        NetworkStateSnapshot[] states = new NetworkStateSnapshot[] {buildWifiState()};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
 
@@ -329,7 +327,7 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
         // pretend that wifi network comes online; service should ask about full
         // network state, and poll any existing interfaces before updating.
         expectDefaultSettings();
-        NetworkState[] states = new NetworkState[] {buildWifiState()};
+        NetworkStateSnapshot[] states = new NetworkStateSnapshot[] {buildWifiState()};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
 
@@ -403,7 +401,7 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
         // pretend that wifi network comes online; service should ask about full
         // network state, and poll any existing interfaces before updating.
         expectSettings(0L, HOUR_IN_MILLIS, WEEK_IN_MILLIS);
-        NetworkState[] states = new NetworkState[] {buildWifiState()};
+        NetworkStateSnapshot[] states = new NetworkStateSnapshot[] {buildWifiState()};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
 
@@ -444,7 +442,7 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
     public void testUidStatsAcrossNetworks() throws Exception {
         // pretend first mobile network comes online
         expectDefaultSettings();
-        NetworkState[] states = new NetworkState[] {buildMobile3gState(IMSI_1)};
+        NetworkStateSnapshot[] states = new NetworkStateSnapshot[] {buildMobile3gState(IMSI_1)};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
 
@@ -475,7 +473,7 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
         // disappearing, to verify we don't count backwards.
         incrementCurrentTime(HOUR_IN_MILLIS);
         expectDefaultSettings();
-        states = new NetworkState[] {buildMobile3gState(IMSI_2)};
+        states = new NetworkStateSnapshot[] {buildMobile3gState(IMSI_2)};
         expectNetworkStatsSummary(new NetworkStats(getElapsedRealtime(), 1)
                 .insertEntry(TEST_IFACE, 2048L, 16L, 512L, 4L));
         expectNetworkStatsUidDetail(new NetworkStats(getElapsedRealtime(), 3)
@@ -519,7 +517,7 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
     public void testUidRemovedIsMoved() throws Exception {
         // pretend that network comes online
         expectDefaultSettings();
-        NetworkState[] states = new NetworkState[] {buildWifiState()};
+        NetworkStateSnapshot[] states = new NetworkStateSnapshot[] {buildWifiState()};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
 
@@ -583,7 +581,8 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
                 buildTemplateMobileWithRatType(null, TelephonyManager.NETWORK_TYPE_LTE);
         final NetworkTemplate template5g =
                 buildTemplateMobileWithRatType(null, TelephonyManager.NETWORK_TYPE_NR);
-        final NetworkState[] states = new NetworkState[]{buildMobile3gState(IMSI_1)};
+        final NetworkStateSnapshot[] states =
+                new NetworkStateSnapshot[]{buildMobile3gState(IMSI_1)};
 
         // 3G network comes online.
         expectNetworkStatsSummary(buildEmptyStats());
@@ -673,7 +672,8 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
                 METERED_ALL, ROAMING_ALL, DEFAULT_NETWORK_ALL, NETWORK_TYPE_ALL, OEM_MANAGED_NO);
 
         // OEM_PAID network comes online.
-        NetworkState[] states = new NetworkState[]{buildOemManagedMobileState(IMSI_1, false,
+        NetworkStateSnapshot[] states = new NetworkStateSnapshot[]{
+                buildOemManagedMobileState(IMSI_1, false,
                 new int[]{NetworkCapabilities.NET_CAPABILITY_OEM_PAID})};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
@@ -688,7 +688,7 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
         forcePollAndWaitForIdle();
 
         // OEM_PRIVATE network comes online.
-        states = new NetworkState[]{buildOemManagedMobileState(IMSI_1, false,
+        states = new NetworkStateSnapshot[]{buildOemManagedMobileState(IMSI_1, false,
                 new int[]{NetworkCapabilities.NET_CAPABILITY_OEM_PRIVATE})};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
@@ -703,7 +703,7 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
         forcePollAndWaitForIdle();
 
         // OEM_PAID + OEM_PRIVATE network comes online.
-        states = new NetworkState[]{buildOemManagedMobileState(IMSI_1, false,
+        states = new NetworkStateSnapshot[]{buildOemManagedMobileState(IMSI_1, false,
                 new int[]{NetworkCapabilities.NET_CAPABILITY_OEM_PRIVATE,
                           NetworkCapabilities.NET_CAPABILITY_OEM_PAID})};
         expectNetworkStatsSummary(buildEmptyStats());
@@ -719,7 +719,7 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
         forcePollAndWaitForIdle();
 
         // OEM_NONE network comes online.
-        states = new NetworkState[]{buildOemManagedMobileState(IMSI_1, false, new int[]{})};
+        states = new NetworkStateSnapshot[]{buildOemManagedMobileState(IMSI_1, false, new int[]{})};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
         mService.forceUpdateIfaces(NETWORKS_MOBILE, states, getActiveIface(states),
@@ -771,7 +771,7 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
     public void testSummaryForAllUid() throws Exception {
         // pretend that network comes online
         expectDefaultSettings();
-        NetworkState[] states = new NetworkState[] {buildWifiState()};
+        NetworkStateSnapshot[] states = new NetworkStateSnapshot[] {buildWifiState()};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
 
@@ -830,7 +830,7 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
     public void testDetailedUidStats() throws Exception {
         // pretend that network comes online
         expectDefaultSettings();
-        NetworkState[] states = new NetworkState[] {buildWifiState()};
+        NetworkStateSnapshot[] states = new NetworkStateSnapshot[] {buildWifiState()};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
 
@@ -871,9 +871,9 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
         final String stackedIface = "stacked-test0";
         final LinkProperties stackedProp = new LinkProperties();
         stackedProp.setInterfaceName(stackedIface);
-        final NetworkState wifiState = buildWifiState();
+        final NetworkStateSnapshot wifiState = buildWifiState();
         wifiState.linkProperties.addStackedLink(stackedProp);
-        NetworkState[] states = new NetworkState[] {wifiState};
+        NetworkStateSnapshot[] states = new NetworkStateSnapshot[] {wifiState};
 
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
@@ -929,7 +929,7 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
     public void testForegroundBackground() throws Exception {
         // pretend that network comes online
         expectDefaultSettings();
-        NetworkState[] states = new NetworkState[] {buildWifiState()};
+        NetworkStateSnapshot[] states = new NetworkStateSnapshot[] {buildWifiState()};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
 
@@ -986,8 +986,8 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
     public void testMetered() throws Exception {
         // pretend that network comes online
         expectDefaultSettings();
-        NetworkState[] states =
+        NetworkStateSnapshot[] states =
-                new NetworkState[] {buildWifiState(true /* isMetered */, TEST_IFACE)};
+                new NetworkStateSnapshot[] {buildWifiState(true /* isMetered */, TEST_IFACE)};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
 
@@ -1026,8 +1026,8 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
     public void testRoaming() throws Exception {
         // pretend that network comes online
         expectDefaultSettings();
-        NetworkState[] states =
+        NetworkStateSnapshot[] states =
-            new NetworkState[] {buildMobile3gState(IMSI_1, true /* isRoaming */)};
+            new NetworkStateSnapshot[] {buildMobile3gState(IMSI_1, true /* isRoaming */)};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
 
@@ -1065,7 +1065,8 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
     public void testTethering() throws Exception {
         // pretend first mobile network comes online
         expectDefaultSettings();
-        final NetworkState[] states = new NetworkState[]{buildMobile3gState(IMSI_1)};
+        final NetworkStateSnapshot[] states =
+                new NetworkStateSnapshot[]{buildMobile3gState(IMSI_1)};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
 
@@ -1122,7 +1123,7 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
         // pretend that wifi network comes online; service should ask about full
         // network state, and poll any existing interfaces before updating.
         expectDefaultSettings();
-        NetworkState[] states = new NetworkState[] {buildWifiState()};
+        NetworkStateSnapshot[] states = new NetworkStateSnapshot[] {buildWifiState()};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
 
@@ -1220,8 +1221,8 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
     public void testStatsProviderUpdateStats() throws Exception {
         // Pretend that network comes online.
         expectDefaultSettings();
-        final NetworkState[] states =
+        final NetworkStateSnapshot[] states =
-                new NetworkState[]{buildWifiState(true /* isMetered */, TEST_IFACE)};
+                new NetworkStateSnapshot[]{buildWifiState(true /* isMetered */, TEST_IFACE)};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
 
@@ -1282,8 +1283,8 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
     public void testStatsProviderSetAlert() throws Exception {
         // Pretend that network comes online.
         expectDefaultSettings();
-        NetworkState[] states =
+        NetworkStateSnapshot[] states =
-                new NetworkState[]{buildWifiState(true /* isMetered */, TEST_IFACE)};
+                new NetworkStateSnapshot[]{buildWifiState(true /* isMetered */, TEST_IFACE)};
         mService.forceUpdateIfaces(NETWORKS_WIFI, states, getActiveIface(states),
                 new UnderlyingNetworkInfo[0]);
 
@@ -1326,7 +1327,8 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
                 buildTemplateMobileWithRatType(null, TelephonyManager.NETWORK_TYPE_UNKNOWN);
         final NetworkTemplate templateAll =
                 buildTemplateMobileWithRatType(null, NETWORK_TYPE_ALL);
-        final NetworkState[] states = new NetworkState[]{buildMobile3gState(IMSI_1)};
+        final NetworkStateSnapshot[] states =
+                new NetworkStateSnapshot[]{buildMobile3gState(IMSI_1)};
 
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
@@ -1401,7 +1403,7 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
     public void testOperationCount_nonDefault_traffic() throws Exception {
         // Pretend mobile network comes online, but wifi is the default network.
         expectDefaultSettings();
-        NetworkState[] states = new NetworkState[]{
+        NetworkStateSnapshot[] states = new NetworkStateSnapshot[]{
                 buildWifiState(true /*isMetered*/, TEST_IFACE2), buildMobile3gState(IMSI_1)};
         expectNetworkStatsUidDetail(buildEmptyStats());
         mService.forceUpdateIfaces(NETWORKS_WIFI, states, getActiveIface(states),
@@ -1489,7 +1491,7 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
         expectNetworkStatsSummary(buildEmptyStats());
     }
 
-    private String getActiveIface(NetworkState... states) throws Exception {
+    private String getActiveIface(NetworkStateSnapshot... states) throws Exception {
         if (states == null || states.length == 0 || states[0].linkProperties == null) {
             return null;
         }
@@ -1565,11 +1567,11 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
         assertEquals("unexpected operations", operations, entry.operations);
     }
 
-    private static NetworkState buildWifiState() {
+    private static NetworkStateSnapshot buildWifiState() {
         return buildWifiState(false, TEST_IFACE);
     }
 
-    private static NetworkState buildWifiState(boolean isMetered, @NonNull String iface) {
+    private static NetworkStateSnapshot buildWifiState(boolean isMetered, @NonNull String iface) {
         final LinkProperties prop = new LinkProperties();
         prop.setInterfaceName(iface);
         final NetworkCapabilities capabilities = new NetworkCapabilities();
@@ -1577,35 +1579,30 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
         capabilities.setCapability(NetworkCapabilities.NET_CAPABILITY_NOT_ROAMING, true);
         capabilities.addTransportType(NetworkCapabilities.TRANSPORT_WIFI);
         capabilities.setSSID(TEST_SSID);
-        return new NetworkState(TYPE_WIFI, prop, capabilities, WIFI_NETWORK, null);
+        return new NetworkStateSnapshot(prop, capabilities, WIFI_NETWORK, null, TYPE_WIFI);
     }
 
-    private static NetworkState buildMobile3gState(String subscriberId) {
+    private static NetworkStateSnapshot buildMobile3gState(String subscriberId) {
         return buildMobile3gState(subscriberId, false /* isRoaming */);
     }
 
-    private static NetworkState buildMobile3gState(String subscriberId, boolean isRoaming) {
+    private static NetworkStateSnapshot buildMobile3gState(String subscriberId, boolean isRoaming) {
         final LinkProperties prop = new LinkProperties();
         prop.setInterfaceName(TEST_IFACE);
         final NetworkCapabilities capabilities = new NetworkCapabilities();
         capabilities.setCapability(NetworkCapabilities.NET_CAPABILITY_NOT_METERED, false);
         capabilities.setCapability(NetworkCapabilities.NET_CAPABILITY_NOT_ROAMING, !isRoaming);
         capabilities.addTransportType(NetworkCapabilities.TRANSPORT_CELLULAR);
-        return new NetworkState(TYPE_MOBILE, prop, capabilities, MOBILE_NETWORK, subscriberId);
+        return new NetworkStateSnapshot(
+                prop, capabilities, MOBILE_NETWORK, subscriberId, TYPE_MOBILE);
     }
 
     private NetworkStats buildEmptyStats() {
         return new NetworkStats(getElapsedRealtime(), 0);
     }
 
-    private static NetworkState buildVpnState() {
+    private static NetworkStateSnapshot buildOemManagedMobileState(
-        final LinkProperties prop = new LinkProperties();
+            String subscriberId, boolean isRoaming, int[] oemNetCapabilities) {
-        prop.setInterfaceName(TUN_IFACE);
-        return new NetworkState(TYPE_VPN, prop, new NetworkCapabilities(), VPN_NETWORK, null);
-    }
-
-    private static NetworkState buildOemManagedMobileState(String subscriberId, boolean isRoaming,
-                int[] oemNetCapabilities) {
         final LinkProperties prop = new LinkProperties();
         prop.setInterfaceName(TEST_IFACE);
         final NetworkCapabilities capabilities = new NetworkCapabilities();
@@ -1615,7 +1612,8 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
             capabilities.setCapability(nc, true);
         }
         capabilities.addTransportType(NetworkCapabilities.TRANSPORT_CELLULAR);
-        return new NetworkState(TYPE_MOBILE, prop, capabilities, MOBILE_NETWORK, subscriberId);
+        return new NetworkStateSnapshot(prop, capabilities, MOBILE_NETWORK, subscriberId,
+                TYPE_MOBILE);
     }
 
     private long getElapsedRealtime() {