Deflake testNetworkBlockedStatusAlwaysOnVpn

This test flakes about 1% of the time on my device. This is due
to the test calling expectNetworkRejectNonSecureVpn() before the
handler thread has processed the call to setAlwaysOnVpnPackage.

Fix this by waiting for callbacks where it is possible to do so,
and by calling waitForIdle where it is not.

Test: test-only change
Change-Id: Ic5750d4fdf9e7e3654a2b6ade74da6bc5de18fb6
diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index 20413f6..1ca79ee 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -6853,6 +6853,7 @@
         final int userId = UserHandle.getUserId(uid);
         final ArrayList<String> allowList = new ArrayList<>();
         mService.setAlwaysOnVpnPackage(userId, ALWAYS_ON_PACKAGE, true /* lockdown */, allowList);
+        waitForIdle();
 
         UidRangeParcel firstHalf = new UidRangeParcel(1, VPN_UID - 1);
         UidRangeParcel secondHalf = new UidRangeParcel(VPN_UID + 1, 99999);
@@ -6874,10 +6875,10 @@
 
         // Disable lockdown, expect to see the network unblocked.
         mService.setAlwaysOnVpnPackage(userId, null, false /* lockdown */, allowList);
-        expectNetworkRejectNonSecureVpn(inOrder, false, firstHalf, secondHalf);
         callback.expectBlockedStatusCallback(false, mWiFiNetworkAgent);
         defaultCallback.expectBlockedStatusCallback(false, mWiFiNetworkAgent);
         vpnUidCallback.assertNoCallback();
+        expectNetworkRejectNonSecureVpn(inOrder, false, firstHalf, secondHalf);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetworkForUid(VPN_UID));
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         assertActiveNetworkInfo(TYPE_WIFI, DetailedState.CONNECTED);
@@ -6920,9 +6921,11 @@
         // Disable lockdown, remove our UID from the allowlist, and re-enable lockdown.
         // Everything should now be blocked.
         mService.setAlwaysOnVpnPackage(userId, null, false /* lockdown */, allowList);
+        waitForIdle();
         expectNetworkRejectNonSecureVpn(inOrder, false, piece1, piece2, piece3);
         allowList.clear();
         mService.setAlwaysOnVpnPackage(userId, ALWAYS_ON_PACKAGE, true /* lockdown */, allowList);
+        waitForIdle();
         expectNetworkRejectNonSecureVpn(inOrder, true, firstHalf, secondHalf);
         defaultCallback.expectBlockedStatusCallback(true, mWiFiNetworkAgent);
         assertBlockedCallbackInAnyOrder(callback, true, mWiFiNetworkAgent, mCellNetworkAgent);