From 40f958b90d0c661d9cdbe4e84be3d4a73a45fddb Mon Sep 17 00:00:00 2001
From: Yuyang Huang <yuyanghuang@google.com>
Date: Thu, 10 Aug 2023 16:30:09 +0900
Subject: [PATCH] Fix flaky test: testNsdManager_registerOffloadEngine

Fix flaky test by only catch relevant the AddOrUpdateEvent and
RemoveEvent.

Bug: 294685851
Test: atest CtsNetTestCases
Change-Id: Ia12e51a6527f1e23610fdb0bea1068ebb5b80a73
---
 tests/cts/net/src/android/net/cts/NsdManagerTest.kt | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/tests/cts/net/src/android/net/cts/NsdManagerTest.kt b/tests/cts/net/src/android/net/cts/NsdManagerTest.kt
index 7731ac99f5..e4ee8de031 100644
--- a/tests/cts/net/src/android/net/cts/NsdManagerTest.kt
+++ b/tests/cts/net/src/android/net/cts/NsdManagerTest.kt
@@ -910,12 +910,16 @@ class NsdManagerTest {
         val record = NsdRegistrationRecord()
         nsdManager.registerService(si, NsdManager.PROTOCOL_DNS_SD, record)
         val addOrUpdateEvent = offloadEngine
-            .expectCallback<TestNsdOffloadEngine.OffloadEvent.AddOrUpdateEvent>()
+            .expectCallbackEventually<TestNsdOffloadEngine.OffloadEvent.AddOrUpdateEvent> {
+                it.info.key.serviceName == serviceName
+            }
         checkOffloadServiceInfo(addOrUpdateEvent.info)
 
         nsdManager.unregisterService(record)
         val unregisterEvent = offloadEngine
-            .expectCallback<TestNsdOffloadEngine.OffloadEvent.RemoveEvent>()
+            .expectCallbackEventually<TestNsdOffloadEngine.OffloadEvent.RemoveEvent> {
+                it.info.key.serviceName == serviceName
+            }
         checkOffloadServiceInfo(unregisterEvent.info)
 
         runAsShell(NETWORK_SETTINGS) {