diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index 7eef2d5..bbb4b7b 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -112,6 +112,7 @@
 import android.net.RouteInfo;
 import android.net.StringNetworkSpecifier;
 import android.net.UidRange;
+import android.net.VpnService;
 import android.net.captiveportal.CaptivePortalProbeResult;
 import android.net.metrics.IpConnectivityLog;
 import android.net.util.MultinetworkPolicyTracker;
@@ -134,6 +135,7 @@
 import android.util.ArraySet;
 import android.util.Log;
 
+import com.android.internal.net.VpnConfig;
 import com.android.internal.util.ArrayUtils;
 import com.android.internal.util.WakeupMessage;
 import com.android.internal.util.test.BroadcastInterceptingContext;
@@ -197,13 +199,13 @@
     private MockNetworkAgent mWiFiNetworkAgent;
     private MockNetworkAgent mCellNetworkAgent;
     private MockNetworkAgent mEthernetNetworkAgent;
+    private MockVpn mMockVpn;
     private Context mContext;
 
     @Mock IpConnectivityMetrics.Logger mMetricsService;
     @Mock DefaultNetworkMetrics mDefaultNetworkMetrics;
     @Mock INetworkManagementService mNetworkManagementService;
     @Mock INetworkStatsService mStatsService;
-    @Mock Vpn mMockVpn;
 
     private ArgumentCaptor<String[]> mStringArrayCaptor = ArgumentCaptor.forClass(String[].class);
 
@@ -479,6 +481,14 @@
             mNetworkAgent.sendNetworkCapabilities(mNetworkCapabilities);
         }
 
+        public void setNetworkCapabilities(NetworkCapabilities nc,
+                boolean sendToConnectivityService) {
+            mNetworkCapabilities.set(nc);
+            if (sendToConnectivityService) {
+                mNetworkAgent.sendNetworkCapabilities(mNetworkCapabilities);
+            }
+        }
+
         public void connectWithoutInternet() {
             mNetworkInfo.setDetailedState(DetailedState.CONNECTED, null, null);
             mNetworkAgent.sendNetworkInfo(mNetworkInfo);
@@ -594,6 +604,10 @@
             return mRedirectUrl;
         }
 
+        public NetworkAgent getNetworkAgent() {
+            return mNetworkAgent;
+        }
+
         public NetworkCapabilities getNetworkCapabilities() {
             return mNetworkCapabilities;
         }
@@ -726,6 +740,87 @@
         }
     }
 
+    private static Looper startHandlerThreadAndReturnLooper() {
+        final HandlerThread handlerThread = new HandlerThread("MockVpnThread");
+        handlerThread.start();
+        return handlerThread.getLooper();
+    }
+
+    private class MockVpn extends Vpn {
+        // TODO : the interactions between this mock and the mock network agent are too
+        // hard to get right at this moment, because it's unclear in which case which
+        // target needs to get a method call or both, and in what order. It's because
+        // MockNetworkAgent wants to manage its own NetworkCapabilities, but the Vpn
+        // parent class of MockVpn agent wants that responsibility.
+        // That being said inside the test it should be possible to make the interactions
+        // harder to get wrong with precise speccing, judicious comments, helper methods
+        // and a few sprinkled assertions.
+
+        private boolean mConnected = false;
+        // Careful ! This is different from mNetworkAgent, because MockNetworkAgent does
+        // not inherit from NetworkAgent.
+        private MockNetworkAgent mMockNetworkAgent;
+
+        public MockVpn(int userId) {
+            super(startHandlerThreadAndReturnLooper(), mServiceContext, mNetworkManagementService,
+                    userId);
+        }
+
+        public void setNetworkAgent(MockNetworkAgent agent) {
+            waitForIdle(agent, TIMEOUT_MS);
+            mMockNetworkAgent = agent;
+            mNetworkAgent = agent.getNetworkAgent();
+            mNetworkCapabilities.set(agent.getNetworkCapabilities());
+        }
+
+        public void setUids(Set<UidRange> uids) {
+            mNetworkCapabilities.setUids(uids);
+            updateCapabilities();
+        }
+
+        @Override
+        public int getNetId() {
+            return mMockNetworkAgent.getNetwork().netId;
+        }
+
+        @Override
+        public boolean appliesToUid(int uid) {
+            return mConnected;  // Trickery to simplify testing.
+        }
+
+        @Override
+        protected boolean isCallerEstablishedOwnerLocked() {
+            return mConnected;  // Similar trickery
+        }
+
+        public void connect() {
+            mNetworkCapabilities.set(mMockNetworkAgent.getNetworkCapabilities());
+            mConnected = true;
+            mConfig = new VpnConfig();
+        }
+
+        @Override
+        public void updateCapabilities() {
+            if (!mConnected) return;
+            super.updateCapabilities();
+            // Because super.updateCapabilities will update the capabilities of the agent but not
+            // the mock agent, the mock agent needs to know about them.
+            copyCapabilitiesToNetworkAgent();
+        }
+
+        private void copyCapabilitiesToNetworkAgent() {
+            if (null != mMockNetworkAgent) {
+                mMockNetworkAgent.setNetworkCapabilities(mNetworkCapabilities,
+                        false /* sendToConnectivityService */);
+            }
+        }
+
+        public void disconnect() {
+            mConnected = false;
+            mConfig = null;
+        }
+    }
+
     private class FakeWakeupMessage extends WakeupMessage {
         private static final int UNREASONABLY_LONG_WAIT = 1000;
 
@@ -894,10 +989,12 @@
 
         public void mockVpn(int uid) {
             synchronized (mVpns) {
+                int userId = UserHandle.getUserId(uid);
+                mMockVpn = new MockVpn(userId);
                 // This has no effect unless the VPN is actually connected, because things like
                 // getActiveNetworkForUidInternal call getNetworkAgentInfoForNetId on the VPN
                 // netId, and check if that network is actually connected.
-                mVpns.put(UserHandle.getUserId(Process.myUid()), mMockVpn);
+                mVpns.put(userId, mMockVpn);
             }
         }
 
@@ -927,7 +1024,6 @@
 
         MockitoAnnotations.initMocks(this);
         when(mMetricsService.defaultNetworkMetrics()).thenReturn(mDefaultNetworkMetrics);
-        when(mMockVpn.appliesToUid(Process.myUid())).thenReturn(true);
 
         // InstrumentationTestRunner prepares a looper, but AndroidJUnitRunner does not.
         // http://b/25897652 .
@@ -1549,7 +1645,8 @@
 
         void expectCapabilitiesLike(Predicate<NetworkCapabilities> fn, MockNetworkAgent agent) {
             CallbackInfo cbi = expectCallback(CallbackState.NETWORK_CAPABILITIES, agent);
-            assertTrue(fn.test((NetworkCapabilities) cbi.arg));
+            assertTrue("Received capabilities don't match expectations : " + cbi.arg,
+                    fn.test((NetworkCapabilities) cbi.arg));
         }
 
         void assertNoCallback() {
@@ -2577,9 +2674,10 @@
         final MockNetworkAgent vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN);
         final ArraySet<UidRange> ranges = new ArraySet<>();
         ranges.add(new UidRange(uid, uid));
-        when(mMockVpn.getNetId()).thenReturn(vpnNetworkAgent.getNetwork().netId);
-        vpnNetworkAgent.setUids(ranges);
+        mMockVpn.setNetworkAgent(vpnNetworkAgent);
+        mMockVpn.setUids(ranges);
         vpnNetworkAgent.connect(true);
+        mMockVpn.connect();
         defaultNetworkCallback.expectAvailableThenValidatedCallbacks(vpnNetworkAgent);
         assertEquals(defaultNetworkCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
@@ -4108,9 +4206,10 @@
         final MockNetworkAgent vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN);
         final ArraySet<UidRange> ranges = new ArraySet<>();
         ranges.add(new UidRange(uid, uid));
-        when(mMockVpn.getNetId()).thenReturn(vpnNetworkAgent.getNetwork().netId);
-        vpnNetworkAgent.setUids(ranges);
+        mMockVpn.setNetworkAgent(vpnNetworkAgent);
+        mMockVpn.setUids(ranges);
         vpnNetworkAgent.connect(false);
+        mMockVpn.connect();
 
         genericNetworkCallback.expectAvailableCallbacksUnvalidated(vpnNetworkAgent);
         genericNotVpnNetworkCallback.assertNoCallback();
@@ -4142,7 +4241,7 @@
         defaultCallback.expectCallback(CallbackState.NETWORK_CAPABILITIES, vpnNetworkAgent);
 
         ranges.add(new UidRange(uid, uid));
-        vpnNetworkAgent.setUids(ranges);
+        mMockVpn.setUids(ranges);
 
         genericNetworkCallback.expectAvailableCallbacksValidated(vpnNetworkAgent);
         genericNotVpnNetworkCallback.assertNoCallback();
@@ -4192,9 +4291,10 @@
         MockNetworkAgent vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN);
         final ArraySet<UidRange> ranges = new ArraySet<>();
         ranges.add(new UidRange(uid, uid));
-        when(mMockVpn.getNetId()).thenReturn(vpnNetworkAgent.getNetwork().netId);
-        vpnNetworkAgent.setUids(ranges);
+        mMockVpn.setNetworkAgent(vpnNetworkAgent);
+        mMockVpn.setUids(ranges);
         vpnNetworkAgent.connect(true /* validated */, false /* hasInternet */);
+        mMockVpn.connect();
 
         defaultCallback.assertNoCallback();
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
@@ -4203,9 +4303,10 @@
         defaultCallback.assertNoCallback();
 
         vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN);
-        when(mMockVpn.getNetId()).thenReturn(vpnNetworkAgent.getNetwork().netId);
-        vpnNetworkAgent.setUids(ranges);
+        mMockVpn.setNetworkAgent(vpnNetworkAgent);
+        mMockVpn.setUids(ranges);
         vpnNetworkAgent.connect(true /* validated */, true /* hasInternet */);
+        mMockVpn.connect();
         defaultCallback.expectAvailableThenValidatedCallbacks(vpnNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
@@ -4214,13 +4315,115 @@
         defaultCallback.expectAvailableCallbacksValidated(mWiFiNetworkAgent);
 
         vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN);
-        when(mMockVpn.getNetId()).thenReturn(vpnNetworkAgent.getNetwork().netId);
         ranges.clear();
-        vpnNetworkAgent.setUids(ranges);
-
+        mMockVpn.setNetworkAgent(vpnNetworkAgent);
+        mMockVpn.setUids(ranges);
         vpnNetworkAgent.connect(false /* validated */, true /* hasInternet */);
+        mMockVpn.connect();
         defaultCallback.assertNoCallback();
 
         mCm.unregisterNetworkCallback(defaultCallback);
     }
+
+    @Test
+    public void testVpnSetUnderlyingNetworks() {
+        final int uid = Process.myUid();
+
+        final TestNetworkCallback vpnNetworkCallback = new TestNetworkCallback();
+        final NetworkRequest vpnNetworkRequest = new NetworkRequest.Builder()
+                .removeCapability(NET_CAPABILITY_NOT_VPN)
+                .addTransportType(TRANSPORT_VPN)
+                .build();
+        NetworkCapabilities nc;
+        mCm.registerNetworkCallback(vpnNetworkRequest, vpnNetworkCallback);
+        vpnNetworkCallback.assertNoCallback();
+
+        final MockNetworkAgent vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN);
+        final ArraySet<UidRange> ranges = new ArraySet<>();
+        ranges.add(new UidRange(uid, uid));
+        mMockVpn.setNetworkAgent(vpnNetworkAgent);
+        mMockVpn.connect();
+        mMockVpn.setUids(ranges);
+        vpnNetworkAgent.connect(true /* validated */, false /* hasInternet */);
+
+        vpnNetworkCallback.expectAvailableThenValidatedCallbacks(vpnNetworkAgent);
+        nc = mCm.getNetworkCapabilities(vpnNetworkAgent.getNetwork());
+        assertTrue(nc.hasTransport(TRANSPORT_VPN));
+        assertFalse(nc.hasTransport(TRANSPORT_CELLULAR));
+        assertFalse(nc.hasTransport(TRANSPORT_WIFI));
+        // For safety reasons a VPN without underlying networks is considered metered.
+        assertFalse(nc.hasCapability(NET_CAPABILITY_NOT_METERED));
+
+        // Connect cell and use it as an underlying network.
+        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent.connect(true);
+
+        mService.setUnderlyingNetworksForVpn(
+                new Network[] { mCellNetworkAgent.getNetwork() });
+
+        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+                && caps.hasTransport(TRANSPORT_CELLULAR) && !caps.hasTransport(TRANSPORT_WIFI)
+                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED),
+                vpnNetworkAgent);
+
+        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent.addCapability(NET_CAPABILITY_NOT_METERED);
+        mWiFiNetworkAgent.connect(true);
+
+        mService.setUnderlyingNetworksForVpn(
+                new Network[] { mCellNetworkAgent.getNetwork(), mWiFiNetworkAgent.getNetwork() });
+
+        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+                && caps.hasTransport(TRANSPORT_CELLULAR) && caps.hasTransport(TRANSPORT_WIFI)
+                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED),
+                vpnNetworkAgent);
+
+        // Don't disconnect, but note the VPN is not using wifi any more.
+        mService.setUnderlyingNetworksForVpn(
+                new Network[] { mCellNetworkAgent.getNetwork() });
+
+        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+                && caps.hasTransport(TRANSPORT_CELLULAR) && !caps.hasTransport(TRANSPORT_WIFI)
+                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED),
+                vpnNetworkAgent);
+
+        // Use Wifi but not cell. Note the VPN is now unmetered.
+        mService.setUnderlyingNetworksForVpn(
+                new Network[] { mWiFiNetworkAgent.getNetwork() });
+
+        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+                && !caps.hasTransport(TRANSPORT_CELLULAR) && caps.hasTransport(TRANSPORT_WIFI)
+                && caps.hasCapability(NET_CAPABILITY_NOT_METERED),
+                vpnNetworkAgent);
+
+        // Use both again.
+        mService.setUnderlyingNetworksForVpn(
+                new Network[] { mCellNetworkAgent.getNetwork(), mWiFiNetworkAgent.getNetwork() });
+
+        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+                && caps.hasTransport(TRANSPORT_CELLULAR) && caps.hasTransport(TRANSPORT_WIFI)
+                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED),
+                vpnNetworkAgent);
+
+        if (false) { // TODO : reactivate this ; in the current state, the below fail.
+            // Disconnect cell. Receive update without even removing the dead network from the
+            // underlying networks – it's dead anyway. Not metered any more.
+            mCellNetworkAgent.disconnect();
+            vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+                    && !caps.hasTransport(TRANSPORT_CELLULAR) && caps.hasTransport(TRANSPORT_WIFI)
+                    && caps.hasCapability(NET_CAPABILITY_NOT_METERED),
+                    vpnNetworkAgent);
+
+            // Disconnect wifi too. No underlying networks should mean this is now metered,
+            // unfortunately a discrepancy in the current implementation has this unmetered.
+            // TODO : fix this.
+            mWiFiNetworkAgent.disconnect();
+            vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+                    && !caps.hasTransport(TRANSPORT_CELLULAR) && !caps.hasTransport(TRANSPORT_WIFI)
+                    && caps.hasCapability(NET_CAPABILITY_NOT_METERED),
+                    vpnNetworkAgent);
+        }
+
+        mMockVpn.disconnect();
+    }
 }