diff --git a/tests/unit/java/com/android/server/connectivity/VpnTest.java b/tests/unit/java/com/android/server/connectivity/VpnTest.java
index 6f25d1b..041e4ea 100644
--- a/tests/unit/java/com/android/server/connectivity/VpnTest.java
+++ b/tests/unit/java/com/android/server/connectivity/VpnTest.java
@@ -92,7 +92,6 @@
 import android.net.LinkProperties;
 import android.net.LocalSocket;
 import android.net.Network;
-import android.net.NetworkAgent;
 import android.net.NetworkCapabilities;
 import android.net.NetworkInfo.DetailedState;
 import android.net.RouteInfo;
@@ -246,7 +245,7 @@
     @Mock private Vpn.SystemServices mSystemServices;
     @Mock private Vpn.IkeSessionWrapper mIkeSessionWrapper;
     @Mock private Vpn.Ikev2SessionCreator mIkev2SessionCreator;
-    @Mock private NetworkAgent mMockNetworkAgent;
+    @Mock private Vpn.VpnNetworkAgentWrapper mMockNetworkAgent;
     @Mock private ConnectivityManager mConnectivityManager;
     @Mock private IpSecService mIpSecService;
     @Mock private VpnProfileStore mVpnProfileStore;
@@ -870,7 +869,7 @@
     public void testRefreshPlatformVpnAppExclusionList_updatesExcludedUids() throws Exception {
         final Vpn vpn = prepareVpnForVerifyAppExclusionList();
         vpn.setAppExclusionList(TEST_VPN_PKG, Arrays.asList(PKGS));
-        verify(mMockNetworkAgent).sendNetworkCapabilities(any());
+        verify(mMockNetworkAgent).doSendNetworkCapabilities(any());
         assertEquals(Arrays.asList(PKGS), vpn.getAppExclusionList(TEST_VPN_PKG));
 
         reset(mMockNetworkAgent);
@@ -887,7 +886,7 @@
                 vpn.mNetworkCapabilities.getUids());
         ArgumentCaptor<NetworkCapabilities> ncCaptor =
                 ArgumentCaptor.forClass(NetworkCapabilities.class);
-        verify(mMockNetworkAgent).sendNetworkCapabilities(ncCaptor.capture());
+        verify(mMockNetworkAgent).doSendNetworkCapabilities(ncCaptor.capture());
         assertEquals(makeVpnUidRange(PRIMARY_USER.id, newExcludedUids),
                 ncCaptor.getValue().getUids());
 
@@ -902,7 +901,7 @@
         assertEquals(Arrays.asList(PKGS), vpn.getAppExclusionList(TEST_VPN_PKG));
         assertEquals(makeVpnUidRange(PRIMARY_USER.id, newExcludedUids),
                 vpn.mNetworkCapabilities.getUids());
-        verify(mMockNetworkAgent).sendNetworkCapabilities(ncCaptor.capture());
+        verify(mMockNetworkAgent).doSendNetworkCapabilities(ncCaptor.capture());
         assertEquals(makeVpnUidRange(PRIMARY_USER.id, newExcludedUids),
                 ncCaptor.getValue().getUids());
     }
@@ -1854,7 +1853,7 @@
                 Collections.singletonList(TEST_NETWORK_2),
                 vpnSnapShot.vpn.mNetworkCapabilities.getUnderlyingNetworks());
         verify(mMockNetworkAgent)
-                .setUnderlyingNetworks(Collections.singletonList(TEST_NETWORK_2));
+                .doSetUnderlyingNetworks(Collections.singletonList(TEST_NETWORK_2));
 
         vpnSnapShot.vpn.mVpnRunner.exitVpnRunner();
     }
@@ -1896,7 +1895,7 @@
                 Collections.singletonList(TEST_NETWORK_2),
                 vpnSnapShot.vpn.mNetworkCapabilities.getUnderlyingNetworks());
         verify(mMockNetworkAgent)
-                .setUnderlyingNetworks(Collections.singletonList(TEST_NETWORK_2));
+                .doSetUnderlyingNetworks(Collections.singletonList(TEST_NETWORK_2));
 
         vpnSnapShot.vpn.mVpnRunner.exitVpnRunner();
     }
@@ -1904,7 +1903,7 @@
     private void verifyHandlingNetworkLoss() throws Exception {
         final ArgumentCaptor<LinkProperties> lpCaptor =
                 ArgumentCaptor.forClass(LinkProperties.class);
-        verify(mMockNetworkAgent).sendLinkProperties(lpCaptor.capture());
+        verify(mMockNetworkAgent).doSendLinkProperties(lpCaptor.capture());
         final LinkProperties lp = lpCaptor.getValue();
 
         assertNull(lp.getInterfaceName());