Merge "Fix flake in testConnectivityDiagnosticsCallbackOnConnectivityReported" am: 1d2ad89424 am: 5c7e4c179b

Original change: https://android-review.googlesource.com/c/platform/packages/modules/Connectivity/+/1755630

Change-Id: I32db3adb9f1a9c95b99ace428044e331df780821
diff --git a/tests/unit/java/com/android/server/ConnectivityServiceTest.java b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
index 9d778a0..75f7126 100644
--- a/tests/unit/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
@@ -10337,10 +10337,7 @@
         mService.registerConnectivityDiagnosticsCallback(
                 mConnectivityDiagnosticsCallback, request, mContext.getPackageName());
 
-        // Block until all other events are done processing.
-        HandlerUtils.waitForIdle(mCsHandlerThread, TIMEOUT_MS);
-
-        verify(mConnectivityDiagnosticsCallback)
+        verify(mConnectivityDiagnosticsCallback, timeout(TIMEOUT_MS))
                 .onConnectivityReportAvailable(argThat(report -> {
                     return INTERFACE_NAME.equals(report.getLinkProperties().getInterfaceName())
                             && report.getNetworkCapabilities().hasTransport(TRANSPORT_CELLULAR);
@@ -10388,12 +10385,10 @@
             throws Exception {
         setUpConnectivityDiagnosticsCallback();
 
-        // Block until all other events are done processing.
-        HandlerUtils.waitForIdle(mCsHandlerThread, TIMEOUT_MS);
-
         // Verify onConnectivityReport fired
-        verify(mConnectivityDiagnosticsCallback).onConnectivityReportAvailable(
-                argThat(report -> areConnDiagCapsRedacted(report.getNetworkCapabilities())));
+        verify(mConnectivityDiagnosticsCallback, timeout(TIMEOUT_MS))
+                .onConnectivityReportAvailable(argThat(report ->
+                        areConnDiagCapsRedacted(report.getNetworkCapabilities())));
     }
 
     @Test
@@ -10404,11 +10399,8 @@
         // cellular network agent
         mCellNetworkAgent.notifyDataStallSuspected();
 
-        // Block until all other events are done processing.
-        HandlerUtils.waitForIdle(mCsHandlerThread, TIMEOUT_MS);
-
         // Verify onDataStallSuspected fired
-        verify(mConnectivityDiagnosticsCallback).onDataStallSuspected(
+        verify(mConnectivityDiagnosticsCallback, timeout(TIMEOUT_MS)).onDataStallSuspected(
                 argThat(report -> areConnDiagCapsRedacted(report.getNetworkCapabilities())));
     }
 
@@ -10423,23 +10415,19 @@
         final boolean hasConnectivity = true;
         mService.reportNetworkConnectivity(n, hasConnectivity);
 
-        // Block until all other events are done processing.
-        HandlerUtils.waitForIdle(mCsHandlerThread, TIMEOUT_MS);
-
         // Verify onNetworkConnectivityReported fired
-        verify(mConnectivityDiagnosticsCallback)
+        verify(mConnectivityDiagnosticsCallback, timeout(TIMEOUT_MS))
                 .onNetworkConnectivityReported(eq(n), eq(hasConnectivity));
-        verify(mConnectivityDiagnosticsCallback).onConnectivityReportAvailable(
-                argThat(report -> areConnDiagCapsRedacted(report.getNetworkCapabilities())));
+        verify(mConnectivityDiagnosticsCallback, timeout(TIMEOUT_MS))
+                .onConnectivityReportAvailable(
+                        argThat(report ->
+                                areConnDiagCapsRedacted(report.getNetworkCapabilities())));
 
         final boolean noConnectivity = false;
         mService.reportNetworkConnectivity(n, noConnectivity);
 
-        // Block until all other events are done processing.
-        HandlerUtils.waitForIdle(mCsHandlerThread, TIMEOUT_MS);
-
         // Wait for onNetworkConnectivityReported to fire
-        verify(mConnectivityDiagnosticsCallback)
+        verify(mConnectivityDiagnosticsCallback, timeout(TIMEOUT_MS))
                 .onNetworkConnectivityReported(eq(n), eq(noConnectivity));
 
         // Also expect a ConnectivityReport after NetworkMonitor asynchronously re-validates
@@ -10477,9 +10465,6 @@
         final boolean noConnectivity = false;
         doAsUid(Process.myUid() + 1, () -> mService.reportNetworkConnectivity(n, noConnectivity));
 
-        // Block until all other events are done processing.
-        HandlerUtils.waitForIdle(mCsHandlerThread, TIMEOUT_MS);
-
         // Wait for onNetworkConnectivityReported to fire
         verify(mConnectivityDiagnosticsCallback, timeout(TIMEOUT_MS))
                 .onNetworkConnectivityReported(eq(n), eq(noConnectivity));