From 78bd66003a714d5c814383c73104ae2dac4af74b Mon Sep 17 00:00:00 2001
From: Paul Hu <paulhu@google.com>
Date: Thu, 26 Sep 2024 05:56:28 +0000
Subject: [PATCH] Fix the test flake in testRemoveServicesAfterSocketDestroyed

The removeServices() is running on a different handler, which
might take some time to complete. Therefore, the verification
should include a timeout to wait longer and prevent test flake.

Bug: 367429622
Test: atest MdnsDiscoveryManagerTests#testRemoveServicesAfter\
      SocketDestroyed --iteration
Change-Id: I79626bad680a62993c3e7d6b356acc10755257e3
---
 .../server/connectivity/mdns/MdnsDiscoveryManagerTests.java  | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java
index d801fbae0b..ab2fb99de8 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java
@@ -26,6 +26,7 @@ import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.eq;
 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 static org.mockito.Mockito.when;
@@ -499,7 +500,7 @@ public class MdnsDiscoveryManagerTests {
         verify(executorProvider).shutdownExecutorService(mockExecutorService);
         verify(mockServiceTypeClientType1Network1).stopSendAndReceive(mockListenerOne);
         verify(socketClient).stopDiscovery();
-        verify(mockServiceCache).removeServices(cacheKey);
+        verify(mockServiceCache, timeout(DEFAULT_TIMEOUT)).removeServices(cacheKey);
     }
 
     @Test
@@ -523,7 +524,7 @@ public class MdnsDiscoveryManagerTests {
         runOnHandler(() -> callback.onSocketDestroyed(SOCKET_KEY_NETWORK_1));
         verify(mockServiceTypeClientType1Network1).notifySocketDestroyed();
         verify(executorProvider).shutdownExecutorService(mockExecutorService);
-        verify(mockServiceCache).removeServices(cacheKey);
+        verify(mockServiceCache, timeout(DEFAULT_TIMEOUT)).removeServices(cacheKey);
     }
 
     private MdnsPacket createMdnsPacket(String serviceType) {