diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index 1539dc3..ad0d0d5 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -322,6 +322,7 @@
     private static final String MOBILE_IFNAME = "test_rmnet_data0";
     private static final String WIFI_IFNAME = "test_wlan0";
     private static final String WIFI_WOL_IFNAME = "test_wlan_wol";
+    private static final String VPN_IFNAME = "tun10042";
     private static final String TEST_PACKAGE_NAME = "com.android.test.package";
     private static final String[] EMPTY_STRING_ARRAY = new String[0];
 
@@ -1046,12 +1047,14 @@
         public MockVpn(int userId) {
             super(startHandlerThreadAndReturnLooper(), mServiceContext, mNetworkManagementService,
                     userId, mock(KeyStore.class));
+            mConfig = new VpnConfig();
         }
 
         public void setNetworkAgent(TestNetworkAgentWrapper agent) {
             agent.waitForIdle(TIMEOUT_MS);
             mMockNetworkAgent = agent;
             mNetworkAgent = agent.getNetworkAgent();
+            mInterface = VPN_IFNAME;
             mNetworkCapabilities.set(agent.getNetworkCapabilities());
         }
 
@@ -1073,16 +1076,6 @@
         }
 
         @Override
-        public boolean appliesToUid(int uid) {
-            return mConnected;  // Trickery to simplify testing.
-        }
-
-        @Override
-        protected boolean isCallerEstablishedOwnerLocked() {
-            return mConnected;  // Similar trickery
-        }
-
-        @Override
         public int getActiveAppVpnType() {
             return mVpnType;
         }
@@ -1090,7 +1083,6 @@
         private void connect(boolean isAlwaysMetered) {
             mNetworkCapabilities.set(mMockNetworkAgent.getNetworkCapabilities());
             mConnected = true;
-            mConfig = new VpnConfig();
             mConfig.isMetered = isAlwaysMetered;
         }
 
@@ -1121,7 +1113,6 @@
 
         public void disconnect() {
             mConnected = false;
-            mConfig = null;
         }
 
         @Override
@@ -1134,18 +1125,6 @@
         private synchronized void setVpnInfo(VpnInfo vpnInfo) {
             mVpnInfo = vpnInfo;
         }
-
-        @Override
-        public synchronized Network[] getUnderlyingNetworks() {
-            if (mUnderlyingNetworks != null) return mUnderlyingNetworks;
-
-            return super.getUnderlyingNetworks();
-        }
-
-        /** Don't override behavior for {@link Vpn#setUnderlyingNetworks}. */
-        private synchronized void overrideUnderlyingNetworks(Network[] underlyingNetworks) {
-            mUnderlyingNetworks = underlyingNetworks;
-        }
     }
 
     private void mockVpn(int uid) {
@@ -5255,7 +5234,7 @@
         final Network wifiNetwork = new Network(mNetIdManager.peekNextNetId());
         mMockVpn.setNetworkAgent(vpnNetworkAgent);
         mMockVpn.setUids(ranges);
-        mMockVpn.setUnderlyingNetworks(new Network[]{wifiNetwork});
+        mService.setUnderlyingNetworksForVpn(new Network[]{wifiNetwork});
         vpnNetworkAgent.connect(false);
         mMockVpn.connect();
         callback.expectAvailableCallbacksUnvalidated(vpnNetworkAgent);
@@ -5269,8 +5248,17 @@
         mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         assertEquals(wifiNetwork, mWiFiNetworkAgent.getNetwork());
         mWiFiNetworkAgent.connect(false);
-        callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
+        // TODO: the callback for the VPN happens before any callbacks are called for the wifi
+        // network that has just connected. There appear to be two issues here:
+        // 1. The VPN code will accept an underlying network as soon as getNetworkCapabilities() for
+        //    it returns non-null (which happens very early, during handleRegisterNetworkAgent).
+        //    This is not correct because that that point the network is not connected and cannot
+        //    pass any traffic.
+        // 2. When a network connects, updateNetworkInfo propagates underlying network capabilities
+        //    before rematching networks.
+        // Given that this scenario can't really happen, this is probably fine for now.
         callback.expectCallback(CallbackEntry.NETWORK_CAPS_UPDATED, vpnNetworkAgent);
+        callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         assertTrue(mCm.getNetworkCapabilities(vpnNetworkAgent.getNetwork())
                 .hasTransport(TRANSPORT_VPN));
         assertTrue(mCm.getNetworkCapabilities(vpnNetworkAgent.getNetwork())
@@ -5334,7 +5322,7 @@
 
         vpnNetworkAgent.connect(false);
         mMockVpn.connect();
-        mMockVpn.setUnderlyingNetworks(new Network[0]);
+        mService.setUnderlyingNetworksForVpn(new Network[0]);
 
         genericNetworkCallback.expectAvailableCallbacksUnvalidated(vpnNetworkAgent);
         genericNotVpnNetworkCallback.assertNoCallback();
@@ -7298,16 +7286,21 @@
         // active
         final VpnInfo info = new VpnInfo();
         info.ownerUid = Process.myUid();
-        info.vpnIface = "interface";
+        info.vpnIface = VPN_IFNAME;
         mMockVpn.setVpnInfo(info);
-        mMockVpn.overrideUnderlyingNetworks(new Network[] {network});
+
+        final TestNetworkAgentWrapper vpnNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_VPN);
+        mMockVpn.setNetworkAgent(vpnNetworkAgent);
+        mMockVpn.connect();
+
+        assertTrue(mService.setUnderlyingNetworksForVpn(new Network[] {network}));
         assertTrue(
                 "Active VPN permission not applied",
                 mService.checkConnectivityDiagnosticsPermissions(
                         Process.myPid(), Process.myUid(), naiWithoutUid,
                         mContext.getOpPackageName()));
 
-        mMockVpn.overrideUnderlyingNetworks(null);
+        assertTrue(mService.setUnderlyingNetworksForVpn(null));
         assertFalse(
                 "VPN shouldn't receive callback on non-underlying network",
                 mService.checkConnectivityDiagnosticsPermissions(