Fix flaky test - testPartialConnectivity

This timing issue is happened when test case already verified
NetworkMonitor#setAcceptPartialConnectivity() before
NetworkAgent send the CONNECTED state to ConnectivityService.
This test case expects that the
NetworkMonitor#setAcceptPartialConnectivity() will be invoked
when network state is CONNECTED in
ConnectivityService#updateNetworkInfo().

NetworkAgent communicates with ConnectivityService by
AsyncChannel. Before they connect to each other, the thread of
NetworkAgent is keeping idle, so waitForIdle cannot handle this
kind of case. Although they have connected to each other, but
they don't communicate with each other, waitForIdle() still
cannot handle this case. It means waitForIdle() can only handle
the case that NetworkAgent already connected to
ConnectivityService and NetworkAgent is communicating with
ConnectivityService. So, that's why waitForIdle() sometimes
doesn't work in this case, and the most simplest and right way
to fix this case is to wait the available callback then verify
the method.

Bug: 140306320
Test: atest FrameworksNetTests: \
      ConnectivityServiceTest#testPartialConnectivity \
      --generate-new-metrics 500

Change-Id: I64bfecbfe3c5776bd5c347a821f9127e79539c48
diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index c63bf42..2d5700b 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -2188,7 +2188,6 @@
     }
 
     @Test
-    @FlakyTest(bugId = 140306320)
     public void testPartialConnectivity() throws Exception {
         // Register network callback.
         NetworkRequest request = new NetworkRequest.Builder()
@@ -2261,9 +2260,8 @@
         // If user accepted partial connectivity network before,
         // NetworkMonitor#setAcceptPartialConnectivity() will be called in
         // ConnectivityService#updateNetworkInfo().
-        waitForIdle();
-        verify(mWiFiNetworkAgent.mNetworkMonitor, times(1)).setAcceptPartialConnectivity();
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
+        verify(mWiFiNetworkAgent.mNetworkMonitor, times(1)).setAcceptPartialConnectivity();
         callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         nc = callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         assertFalse(nc.hasCapability(NET_CAPABILITY_PARTIAL_CONNECTIVITY));
@@ -2283,9 +2281,8 @@
         // If user accepted partial connectivity network before,
         // NetworkMonitor#setAcceptPartialConnectivity() will be called in
         // ConnectivityService#updateNetworkInfo().
-        waitForIdle();
-        verify(mWiFiNetworkAgent.mNetworkMonitor, times(1)).setAcceptPartialConnectivity();
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
+        verify(mWiFiNetworkAgent.mNetworkMonitor, times(1)).setAcceptPartialConnectivity();
         callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         callback.expectCapabilitiesWith(NET_CAPABILITY_PARTIAL_CONNECTIVITY, mWiFiNetworkAgent);
@@ -2308,9 +2305,8 @@
         // valid, because ConnectivityService calls setAcceptPartialConnectivity before it calls
         // notifyNetworkConnected.
         mWiFiNetworkAgent.connectWithPartialValidConnectivity();
-        waitForIdle();
-        verify(mWiFiNetworkAgent.mNetworkMonitor, times(1)).setAcceptPartialConnectivity();
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
+        verify(mWiFiNetworkAgent.mNetworkMonitor, times(1)).setAcceptPartialConnectivity();
         callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         callback.expectCapabilitiesWith(
                 NET_CAPABILITY_PARTIAL_CONNECTIVITY | NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);