diff --git a/tests/unit/java/android/net/NetworkTemplateTest.kt b/tests/unit/java/android/net/NetworkTemplateTest.kt
index 0c3bee3..048597f 100644
--- a/tests/unit/java/android/net/NetworkTemplateTest.kt
+++ b/tests/unit/java/android/net/NetworkTemplateTest.kt
@@ -95,7 +95,7 @@
         oemManaged: Int = OEM_NONE,
         metered: Boolean = true
     ): NetworkStateSnapshot {
-        `when`(mockWifiInfo.getCurrentNetworkKey()).thenReturn(wifiKey)
+        `when`(mockWifiInfo.getNetworkKey()).thenReturn(wifiKey)
         val lp = LinkProperties()
         val caps = NetworkCapabilities().apply {
             setCapability(NetworkCapabilities.NET_CAPABILITY_NOT_METERED, !metered)
diff --git a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
index 5e1699a..d7bbf50 100644
--- a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
+++ b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
@@ -266,7 +266,7 @@
         mServiceContext = new MockContext(context);
         when(mLocationPermissionChecker.checkCallersLocationPermission(
                 any(), any(), anyInt(), anyBoolean(), any())).thenReturn(true);
-        when(sWifiInfo.getCurrentNetworkKey()).thenReturn(TEST_WIFI_NETWORK_KEY);
+        when(sWifiInfo.getNetworkKey()).thenReturn(TEST_WIFI_NETWORK_KEY);
         mStatsDir = TestIoUtils.createTemporaryDirectory(getClass().getSimpleName());
 
         PowerManager powerManager = (PowerManager) mServiceContext.getSystemService(