diff --git a/services/core/java/com/android/server/ConnectivityService.java b/services/core/java/com/android/server/ConnectivityService.java
index 9012fbc..7541833 100644
--- a/services/core/java/com/android/server/ConnectivityService.java
+++ b/services/core/java/com/android/server/ConnectivityService.java
@@ -3479,6 +3479,9 @@
             }
         }
         nai.clearLingerState();
+        // TODO: this loop, and the mLegacyTypeTracker.remove just below it, seem redundant given
+        // there's a full rematch right after. Currently, deleting it breaks tests that check for
+        // the default network disconnecting. Find out why, fix the rematch code, and delete this.
         if (nai.isSatisfyingRequest(mDefaultRequest.requestId)) {
             mDefaultNetworkNai = null;
             updateDataActivityTracking(null /* newNetwork */, nai);
@@ -4995,6 +4998,8 @@
 
     @Override
     public boolean updateLockdownVpn() {
+        // Allow the system UID for the system server and for Settings.
+        // Also, for unit tests, allow the process that ConnectivityService is running in.
         if (mDeps.getCallingUid() != Process.SYSTEM_UID
                 && Binder.getCallingPid() != Process.myPid()) {
             logw("Lockdown VPN only available to system process or AID_SYSTEM");
diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index bca4986..37307a4 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -1389,9 +1389,6 @@
         verify(mNetworkPolicyManager).registerListener(policyListenerCaptor.capture());
         mPolicyListener = policyListenerCaptor.getValue();
 
-        mServiceContext.setPermission(
-                Manifest.permission.CONTROL_VPN, PERMISSION_GRANTED);
-
         // Create local CM before sending system ready so that we can answer
         // getSystemService() correctly.
         mCm = new WrappedConnectivityManager(InstrumentationRegistry.getContext(), mService);
@@ -1588,6 +1585,7 @@
         }
 
         public void expectNoBroadcast(int timeoutMs) throws Exception {
+            waitForIdle();
             try {
                 final Intent intent = get(timeoutMs, TimeUnit.MILLISECONDS);
                 fail("Unexpected broadcast: " + intent.getAction() + " " + intent.getExtras());
@@ -5906,7 +5904,8 @@
         mCm.registerDefaultNetworkCallback(callback);
 
         // Connect a VPN.
-        mMockVpn.establishForMyUid(false, true, false);
+        mMockVpn.establishForMyUid(false /* validated */, true /* hasInternet */,
+                false /* isStrictMode */);
         callback.expectAvailableCallbacksUnvalidated(mMockVpn);
 
         // Connect cellular data.
@@ -5966,6 +5965,7 @@
         // The same bug happens in the opposite direction: the VPN's capabilities correctly have
         // NOT_SUSPENDED, but the VPN's NetworkInfo is in state SUSPENDED.
         mCellNetworkAgent.resume();
+        callback.assertNoCallback();
         mWiFiNetworkAgent.disconnect();
         callback.expectCapabilitiesThat(mMockVpn,
                 nc -> nc.hasCapability(NET_CAPABILITY_NOT_SUSPENDED)
@@ -6573,9 +6573,13 @@
 
     @Test
     public void testLockdownVpnWithRestrictedProfiles() throws Exception {
-        // NETWORK_SETTINGS is necessary to see the UID ranges in NetworkCapabilities.
+        // For ConnectivityService#setAlwaysOnVpnPackage.
         mServiceContext.setPermission(
                 Manifest.permission.CONTROL_ALWAYS_ON_VPN, PERMISSION_GRANTED);
+        // For call Vpn#setAlwaysOnPackage.
+        mServiceContext.setPermission(
+                Manifest.permission.CONTROL_VPN, PERMISSION_GRANTED);
+        // Necessary to see the UID ranges in NetworkCapabilities.
         mServiceContext.setPermission(
                 Manifest.permission.NETWORK_SETTINGS, PERMISSION_GRANTED);
 
@@ -7138,6 +7142,9 @@
 
     @Test
     public void testLegacyLockdownVpn() throws Exception {
+        mServiceContext.setPermission(
+                Manifest.permission.CONTROL_VPN, PERMISSION_GRANTED);
+
         final NetworkRequest request = new NetworkRequest.Builder().clearCapabilities().build();
         final TestNetworkCallback callback = new TestNetworkCallback();
         mCm.registerNetworkCallback(request, callback);
@@ -7255,11 +7262,14 @@
         mMockVpn.expectStopVpnRunnerPrivileged();
         mMockVpn.expectStartLegacyVpnRunner();
 
-        // TODO: why is wifi not blocked? Is this because something calls prepare()?
+        // TODO: why is wifi not blocked? Is it because when this callback is sent, the VPN is still
+        // connected, so the network is not considered blocked by the lockdown UID ranges? But the
+        // fact that a VPN is connected should only result in the VPN itself being unblocked, not
+        // any other network. Bug in isUidBlockedByVpn?
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        callback.expectCapabilitiesThat(mMockVpn, (nc) -> nc.hasTransport(TRANSPORT_WIFI));
-        defaultCallback.expectCapabilitiesThat(mMockVpn, (nc) -> nc.hasTransport(TRANSPORT_WIFI));
+        callback.expectCapabilitiesThat(mMockVpn, nc -> nc.hasTransport(TRANSPORT_WIFI));
         callback.expectCallback(CallbackEntry.LOST, mMockVpn);
+        defaultCallback.expectCapabilitiesThat(mMockVpn, nc -> nc.hasTransport(TRANSPORT_WIFI));
         defaultCallback.expectCallback(CallbackEntry.LOST, mMockVpn);
         defaultCallback.expectAvailableCallbacksUnvalidatedAndBlocked(mWiFiNetworkAgent);
 
@@ -7302,7 +7312,7 @@
         mWiFiNetworkAgent.disconnect();
         callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
         b1.expectBroadcast();
-        callback.expectCapabilitiesThat(mMockVpn, (nc) -> !nc.hasTransport(TRANSPORT_WIFI));
+        callback.expectCapabilitiesThat(mMockVpn, nc -> !nc.hasTransport(TRANSPORT_WIFI));
         b2 = expectConnectivityAction(TYPE_VPN, DetailedState.DISCONNECTED);
         mMockVpn.expectStopVpnRunnerPrivileged();
         callback.expectCallback(CallbackEntry.LOST, mMockVpn);