Merge "Use waitForIdle() to wait for Threaded processing."
diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index 08defd2..8da1a5b 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -6446,14 +6446,16 @@
     public void testRegisterUnregisterConnectivityDiagnosticsCallback() throws Exception {
         final NetworkRequest wifiRequest =
                 new NetworkRequest.Builder().addTransportType(TRANSPORT_WIFI).build();
-
         when(mConnectivityDiagnosticsCallback.asBinder()).thenReturn(mIBinder);
 
         mService.registerConnectivityDiagnosticsCallback(
                 mConnectivityDiagnosticsCallback, wifiRequest, mContext.getPackageName());
 
-        verify(mIBinder, timeout(TIMEOUT_MS))
-                .linkToDeath(any(ConnectivityDiagnosticsCallbackInfo.class), anyInt());
+        // Block until all other events are done processing.
+        HandlerUtilsKt.waitForIdle(mCsHandlerThread, TIMEOUT_MS);
+
+        verify(mIBinder).linkToDeath(any(ConnectivityDiagnosticsCallbackInfo.class), anyInt());
+        verify(mConnectivityDiagnosticsCallback).asBinder();
         assertTrue(
                 mService.mConnectivityDiagnosticsCallbacks.containsKey(
                         mConnectivityDiagnosticsCallback));
@@ -6476,8 +6478,10 @@
         mService.registerConnectivityDiagnosticsCallback(
                 mConnectivityDiagnosticsCallback, wifiRequest, mContext.getPackageName());
 
-        verify(mIBinder, timeout(TIMEOUT_MS))
-                .linkToDeath(any(ConnectivityDiagnosticsCallbackInfo.class), anyInt());
+        // Block until all other events are done processing.
+        HandlerUtilsKt.waitForIdle(mCsHandlerThread, TIMEOUT_MS);
+
+        verify(mIBinder).linkToDeath(any(ConnectivityDiagnosticsCallbackInfo.class), anyInt());
         verify(mConnectivityDiagnosticsCallback).asBinder();
         assertTrue(
                 mService.mConnectivityDiagnosticsCallbacks.containsKey(
@@ -6635,8 +6639,11 @@
     public void testConnectivityDiagnosticsCallbackOnConnectivityReport() throws Exception {
         setUpConnectivityDiagnosticsCallback();
 
-        // Wait for onConnectivityReport to fire
-        verify(mConnectivityDiagnosticsCallback, timeout(TIMEOUT_MS))
+        // Block until all other events are done processing.
+        HandlerUtilsKt.waitForIdle(mCsHandlerThread, TIMEOUT_MS);
+
+        // Verify onConnectivityReport fired
+        verify(mConnectivityDiagnosticsCallback)
                 .onConnectivityReport(any(ConnectivityReport.class));
     }
 
@@ -6648,9 +6655,11 @@
         // cellular network agent
         mCellNetworkAgent.notifyDataStallSuspected();
 
-        // Wait for onDataStallSuspected to fire
-        verify(mConnectivityDiagnosticsCallback, timeout(TIMEOUT_MS))
-                .onDataStallSuspected(any(DataStallReport.class));
+        // Block until all other events are done processing.
+        HandlerUtilsKt.waitForIdle(mCsHandlerThread, TIMEOUT_MS);
+
+        // Verify onDataStallSuspected fired
+        verify(mConnectivityDiagnosticsCallback).onDataStallSuspected(any(DataStallReport.class));
     }
 
     @Test
@@ -6661,15 +6670,21 @@
         final boolean hasConnectivity = true;
         mService.reportNetworkConnectivity(n, hasConnectivity);
 
-        // Wait for onNetworkConnectivityReported to fire
-        verify(mConnectivityDiagnosticsCallback, timeout(TIMEOUT_MS))
+        // Block until all other events are done processing.
+        HandlerUtilsKt.waitForIdle(mCsHandlerThread, TIMEOUT_MS);
+
+        // Verify onNetworkConnectivityReported fired
+        verify(mConnectivityDiagnosticsCallback)
                 .onNetworkConnectivityReported(eq(n), eq(hasConnectivity));
 
         final boolean noConnectivity = false;
         mService.reportNetworkConnectivity(n, noConnectivity);
 
+        // Block until all other events are done processing.
+        HandlerUtilsKt.waitForIdle(mCsHandlerThread, TIMEOUT_MS);
+
         // Wait for onNetworkConnectivityReported to fire
-        verify(mConnectivityDiagnosticsCallback, timeout(TIMEOUT_MS))
+        verify(mConnectivityDiagnosticsCallback)
                 .onNetworkConnectivityReported(eq(n), eq(noConnectivity));
     }
 }