Merge "Fix NetworkRequestReceivedRemoved test" into main am: 8088170283

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

Change-Id: Icecfdc678aedb5b9c8587f4f02129b4327c4275a
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>
diff --git a/tests/unit/java/com/android/server/connectivity/NetworkRequestStateStatsMetricsTest.java b/tests/unit/java/com/android/server/connectivity/NetworkRequestStateStatsMetricsTest.java
index e8b62f8..8dc0528 100644
--- a/tests/unit/java/com/android/server/connectivity/NetworkRequestStateStatsMetricsTest.java
+++ b/tests/unit/java/com/android/server/connectivity/NetworkRequestStateStatsMetricsTest.java
@@ -27,7 +27,6 @@
 import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.timeout;
-import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
 import android.net.NetworkCapabilities;
@@ -99,11 +98,10 @@
         // This call will be used to calculate NR received time
         Mockito.when(mNRStateInfoDeps.getElapsedRealtime()).thenReturn(nrStartTime);
         mNetworkRequestStateStatsMetrics.onNetworkRequestReceived(NOT_METERED_WIFI_NETWORK_REQUEST);
-        HandlerUtils.waitForIdle(mHandlerThread, TIMEOUT_MS);
 
         ArgumentCaptor<NetworkRequestStateInfo> networkRequestStateInfoCaptor =
                 ArgumentCaptor.forClass(NetworkRequestStateInfo.class);
-        verify(mNRStateStatsDeps, times(1))
+        verify(mNRStateStatsDeps, timeout(TIMEOUT_MS))
                 .writeStats(networkRequestStateInfoCaptor.capture());
 
         NetworkRequestStateInfo nrStateInfoSent = networkRequestStateInfoCaptor.getValue();
@@ -120,9 +118,8 @@
         // This call will be used to calculate NR removed time
         Mockito.when(mNRStateInfoDeps.getElapsedRealtime()).thenReturn(nrEndTime);
         mNetworkRequestStateStatsMetrics.onNetworkRequestRemoved(NOT_METERED_WIFI_NETWORK_REQUEST);
-        HandlerUtils.waitForIdle(mHandlerThread, TIMEOUT_MS);
 
-        verify(mNRStateStatsDeps, times(1))
+        verify(mNRStateStatsDeps, timeout(TIMEOUT_MS))
                 .writeStats(networkRequestStateInfoCaptor.capture());
 
         nrStateInfoSent = networkRequestStateInfoCaptor.getValue();
@@ -147,8 +144,7 @@
     @Test
     public void testNoMessagesWhenNetworkRequestReceived() {
         mNetworkRequestStateStatsMetrics.onNetworkRequestReceived(NOT_METERED_WIFI_NETWORK_REQUEST);
-        HandlerUtils.waitForIdle(mHandlerThread, TIMEOUT_MS);
-        verify(mNRStateStatsDeps, times(1))
+        verify(mNRStateStatsDeps, timeout(TIMEOUT_MS))
                 .writeStats(any(NetworkRequestStateInfo.class));
 
         clearInvocations(mNRStateStatsDeps);
@@ -184,7 +180,7 @@
         // Check only first MAX_QUEUED_REQUESTS events are logged.
         ArgumentCaptor<NetworkRequestStateInfo> networkRequestStateInfoCaptor =
                 ArgumentCaptor.forClass(NetworkRequestStateInfo.class);
-        verify(mNRStateStatsDeps, timeout(100).times(
+        verify(mNRStateStatsDeps, timeout(TIMEOUT_MS).times(
                 NetworkRequestStateStatsMetrics.MAX_QUEUED_REQUESTS))
                 .writeStats(networkRequestStateInfoCaptor.capture());
         for (int i = 0; i < NetworkRequestStateStatsMetrics.MAX_QUEUED_REQUESTS; i++) {