diff --git a/services/core/java/com/android/server/ConnectivityService.java b/services/core/java/com/android/server/ConnectivityService.java
index 5673f96..b6232a0 100644
--- a/services/core/java/com/android/server/ConnectivityService.java
+++ b/services/core/java/com/android/server/ConnectivityService.java
@@ -6108,7 +6108,7 @@
         nai.networkAgentPortalData = lp.getCaptivePortalData();
     }
 
-    private void updateLinkProperties(NetworkAgentInfo networkAgent, @NonNull LinkProperties newLp,
+    private void updateLinkProperties(NetworkAgentInfo networkAgent, LinkProperties newLp,
             @NonNull LinkProperties oldLp) {
         int netId = networkAgent.network.getNetId();
 
@@ -6117,7 +6117,8 @@
         // the LinkProperties for the network are accurate.
         networkAgent.clatd.fixupLinkProperties(oldLp, newLp);
 
-        updateInterfaces(newLp, oldLp, netId, networkAgent.networkCapabilities);
+        updateInterfaces(newLp, oldLp, netId, networkAgent.networkCapabilities,
+                networkAgent.networkInfo.getType());
 
         // update filtering rules, need to happen after the interface update so netd knows about the
         // new interface (the interface name -> index map becomes initialized)
@@ -6256,7 +6257,7 @@
 
     private void updateInterfaces(final @Nullable LinkProperties newLp,
             final @Nullable LinkProperties oldLp, final int netId,
-            final @NonNull NetworkCapabilities caps) {
+            final @Nullable NetworkCapabilities caps, final int legacyType) {
         final CompareResult<String> interfaceDiff = new CompareResult<>(
                 oldLp != null ? oldLp.getAllInterfaceNames() : null,
                 newLp != null ? newLp.getAllInterfaceNames() : null);
@@ -6267,7 +6268,7 @@
                     if (DBG) log("Adding iface " + iface + " to network " + netId);
                     mNetd.networkAddInterface(netId, iface);
                     wakeupModifyInterface(iface, caps, true);
-                    bs.noteNetworkInterfaceForTransports(iface, caps.getTransportTypes());
+                    bs.noteNetworkInterfaceType(iface, legacyType);
                 } catch (Exception e) {
                     loge("Exception adding interface: " + e);
                 }
@@ -6539,7 +6540,6 @@
      * maintained here that the NetworkAgent is not aware of (e.g., validated, captive portal,
      * and foreground status).
      */
-    @NonNull
     private NetworkCapabilities mixInCapabilities(NetworkAgentInfo nai, NetworkCapabilities nc) {
         // Once a NetworkAgent is connected, complain if some immutable capabilities are removed.
          // Don't complain for VPNs since they're not driven by requests and there is no risk of
diff --git a/services/core/java/com/android/server/connectivity/NetworkAgentInfo.java b/services/core/java/com/android/server/connectivity/NetworkAgentInfo.java
index b282484..ab0360b 100644
--- a/services/core/java/com/android/server/connectivity/NetworkAgentInfo.java
+++ b/services/core/java/com/android/server/connectivity/NetworkAgentInfo.java
@@ -329,7 +329,7 @@
     private final QosCallbackTracker mQosCallbackTracker;
 
     public NetworkAgentInfo(INetworkAgent na, Network net, NetworkInfo info,
-            @NonNull LinkProperties lp, @NonNull NetworkCapabilities nc, int score, Context context,
+            LinkProperties lp, NetworkCapabilities nc, int score, Context context,
             Handler handler, NetworkAgentConfig config, ConnectivityService connService, INetd netd,
             IDnsResolver dnsResolver, INetworkManagementService nms, int factorySerialNumber,
             int creatorUid, QosCallbackTracker qosCallbackTracker) {
diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index 47f2b32..c5e6c35 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -7367,6 +7367,7 @@
         mCm.unregisterNetworkCallback(trustedCallback);
     }
 
+    @Ignore // 40%+ flakiness : figure out why and re-enable.
     @Test
     public final void testBatteryStatsNetworkType() throws Exception {
         final LinkProperties cellLp = new LinkProperties();
@@ -7374,8 +7375,8 @@
         mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR, cellLp);
         mCellNetworkAgent.connect(true);
         waitForIdle();
-        verify(mBatteryStatsService).noteNetworkInterfaceForTransports(cellLp.getInterfaceName(),
-                new int[] { TRANSPORT_CELLULAR });
+        verify(mBatteryStatsService).noteNetworkInterfaceType(cellLp.getInterfaceName(),
+                TYPE_MOBILE);
         reset(mBatteryStatsService);
 
         final LinkProperties wifiLp = new LinkProperties();
@@ -7383,20 +7384,18 @@
         mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI, wifiLp);
         mWiFiNetworkAgent.connect(true);
         waitForIdle();
-        verify(mBatteryStatsService).noteNetworkInterfaceForTransports(wifiLp.getInterfaceName(),
-                new int[] { TRANSPORT_WIFI });
+        verify(mBatteryStatsService).noteNetworkInterfaceType(wifiLp.getInterfaceName(),
+                TYPE_WIFI);
         reset(mBatteryStatsService);
 
         mCellNetworkAgent.disconnect();
-        mWiFiNetworkAgent.disconnect();
 
         cellLp.setInterfaceName("wifi0");
         mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR, cellLp);
         mCellNetworkAgent.connect(true);
         waitForIdle();
-        verify(mBatteryStatsService).noteNetworkInterfaceForTransports(cellLp.getInterfaceName(),
-                new int[] { TRANSPORT_CELLULAR });
-        mCellNetworkAgent.disconnect();
+        verify(mBatteryStatsService).noteNetworkInterfaceType(cellLp.getInterfaceName(),
+                TYPE_MOBILE);
     }
 
     /**
@@ -7469,8 +7468,8 @@
         assertRoutesAdded(cellNetId, ipv6Subnet, defaultRoute);
         verify(mMockDnsResolver, times(1)).createNetworkCache(eq(cellNetId));
         verify(mMockNetd, times(1)).networkAddInterface(cellNetId, MOBILE_IFNAME);
-        verify(mBatteryStatsService).noteNetworkInterfaceForTransports(cellLp.getInterfaceName(),
-                new int[] { TRANSPORT_CELLULAR });
+        verify(mBatteryStatsService).noteNetworkInterfaceType(cellLp.getInterfaceName(),
+                TYPE_MOBILE);
 
         networkCallback.expectAvailableThenValidatedCallbacks(mCellNetworkAgent);
         verify(mMockDnsResolver, times(1)).startPrefix64Discovery(cellNetId);
@@ -7490,8 +7489,7 @@
         // Make sure BatteryStats was not told about any v4- interfaces, as none should have
         // come online yet.
         waitForIdle();
-        verify(mBatteryStatsService, never()).noteNetworkInterfaceForTransports(startsWith("v4-"),
-                any());
+        verify(mBatteryStatsService, never()).noteNetworkInterfaceType(startsWith("v4-"), anyInt());
 
         verifyNoMoreInteractions(mMockNetd);
         verifyNoMoreInteractions(mMockDnsResolver);
@@ -7544,8 +7542,8 @@
         assertTrue(ArrayUtils.contains(resolvrParams.servers, "8.8.8.8"));
 
         for (final LinkProperties stackedLp : stackedLpsAfterChange) {
-            verify(mBatteryStatsService).noteNetworkInterfaceForTransports(
-                    stackedLp.getInterfaceName(), new int[] { TRANSPORT_CELLULAR });
+            verify(mBatteryStatsService).noteNetworkInterfaceType(stackedLp.getInterfaceName(),
+                    TYPE_MOBILE);
         }
         reset(mMockNetd);
         when(mMockNetd.interfaceGetCfg(CLAT_PREFIX + MOBILE_IFNAME))