Test: Replace wait_until() with wait_for()

Refactor waitForPrivateDnsValidation() to use wait_for().
The predicate is used to check whether the waiting should be
continued or not. Since waitForNat64Prefix() and popDnsEvent()
will be changed to use condition variable, writing this way
would make it easier to change them and avoid unnecessary awakenings.

Bug: 310108475
Test: resolv_integration_test passed
Change-Id: I2282450a6f3626cbb749030154b581554df7b174

2 files changed, 6 insertions, 22 deletions

diff --git a/tests/dns_metrics_listener/dns_metrics_listener.cpp b/tests/dns_metrics_listener/dns_metrics_listener.cpp
index 1715118d..e6cdef88 100644
--- a/tests/dns_metrics_listener/dns_metrics_listener.cpp
+++ b/tests/dns_metrics_listener/dns_metrics_listener.cpp
@@ -95,18 +95,10 @@ bool DnsMetricsListener::waitForNat64Prefix(ExpectNat64PrefixStatus status, mill
 
 bool DnsMetricsListener::waitForPrivateDnsValidation(const std::string& serverAddr,
                                                      const bool validated) {
-    const auto now = std::chrono::steady_clock::now();
-
     std::unique_lock lock(mMutex);
-    ScopedLockAssertion assume_lock(mMutex);
-
-    // onPrivateDnsValidationEvent() might already be invoked. Search for the record first.
-    do {
-        if (findAndRemoveValidationRecord({mNetId, serverAddr}, validated)) return true;
-    } while (mCv.wait_until(lock, now + kEventTimeoutMs) != std::cv_status::timeout);
-
-    // Timeout.
-    return false;
+    return mCv.wait_for(lock, kEventTimeoutMs, [&]() REQUIRES(mMutex) {
+        return findAndRemoveValidationRecord({mNetId, serverAddr}, validated);
+    });
 }
 
 bool DnsMetricsListener::findAndRemoveValidationRecord(const ServerKey& key, const bool value) {
diff --git a/tests/unsolicited_listener/unsolicited_event_listener.cpp b/tests/unsolicited_listener/unsolicited_event_listener.cpp
index 20edae13..367df972 100644
--- a/tests/unsolicited_listener/unsolicited_event_listener.cpp
+++ b/tests/unsolicited_listener/unsolicited_event_listener.cpp
@@ -68,18 +68,10 @@ constexpr milliseconds kRetryIntervalMs{20};
 
 bool UnsolicitedEventListener::waitForPrivateDnsValidation(const std::string& serverAddr,
                                                            int validation, int protocol) {
-    const auto now = std::chrono::steady_clock::now();
-
     std::unique_lock lock(mMutex);
-    ScopedLockAssertion assume_lock(mMutex);
-
-    // onPrivateDnsValidationEvent() might already be invoked. Search for the record first.
-    do {
-        if (findAndRemoveValidationRecord({mNetId, serverAddr, protocol}, validation)) return true;
-    } while (mCv.wait_until(lock, now + kEventTimeoutMs) != std::cv_status::timeout);
-
-    // Timeout.
-    return false;
+    return mCv.wait_for(lock, kEventTimeoutMs, [&]() REQUIRES(mMutex) {
+        return findAndRemoveValidationRecord({mNetId, serverAddr, protocol}, validation);
+    });
 }
 
 bool UnsolicitedEventListener::findAndRemoveValidationRecord(const ServerKey& key, int value) {