Merge "Reduce DnsResolverTest flaky rate" am: f1562ba31c

Change-Id: I30389e9ab8e663dff9d5eb336b482b4eb323f559
diff --git a/tests/tests/net/src/android/net/cts/DnsResolverTest.java b/tests/tests/net/src/android/net/cts/DnsResolverTest.java
index c32a7a0..1cc49f9 100644
--- a/tests/tests/net/src/android/net/cts/DnsResolverTest.java
+++ b/tests/tests/net/src/android/net/cts/DnsResolverTest.java
@@ -86,7 +86,7 @@
     static final int CANCEL_RETRY_TIMES = 5;
     static final int QUERY_TIMES = 10;
     static final int NXDOMAIN = 3;
-    static final int PRIVATE_DNS_SETTING_TIMEOUT_MS = 2_000;
+    static final int PRIVATE_DNS_SETTING_TIMEOUT_MS = 6_000;
 
     private ContentResolver mCR;
     private ConnectivityManager mCM;
@@ -122,10 +122,15 @@
         mOldDnsSpecifier = Settings.Global.getString(mCR, Settings.Global.PRIVATE_DNS_SPECIFIER);
     }
 
-    private void restorePrivateDnsSetting() {
+    private void restorePrivateDnsSetting() throws InterruptedException {
         // restore private DNS setting
         Settings.Global.putString(mCR, Settings.Global.PRIVATE_DNS_MODE, mOldMode);
-        Settings.Global.putString(mCR, Settings.Global.PRIVATE_DNS_SPECIFIER, mOldDnsSpecifier);
+        if ("hostname".equals(mOldMode)) {
+            Settings.Global.putString(
+                mCR, Settings.Global.PRIVATE_DNS_SPECIFIER, mOldDnsSpecifier);
+            mCtsNetUtils.awaitPrivateDnsSetting("restorePrivateDnsSetting timeout",
+                    mCM.getActiveNetwork(), mOldDnsSpecifier, PRIVATE_DNS_SETTING_TIMEOUT_MS, true);
+        }
     }
 
     private static String byteArrayToHexString(byte[] bytes) {
@@ -203,6 +208,7 @@
         private final CancellationSignal mCancelSignal;
         private int mRcode;
         private DnsAnswer mDnsAnswer;
+        private String mErrorMsg = null;
 
         VerifyCancelCallback(@NonNull String msg, @Nullable CancellationSignal cancel) {
             mMsg = msg;
@@ -228,14 +234,18 @@
         @Override
         public void onAnswer(@NonNull byte[] answer, int rcode) {
             if (mCancelSignal != null && mCancelSignal.isCanceled()) {
-                fail(mMsg + " should not have returned any answers");
+                mErrorMsg = mMsg + " should not have returned any answers";
+                mLatch.countDown();
+                return;
             }
 
             mRcode = rcode;
             try {
                 mDnsAnswer = new DnsAnswer(answer);
             } catch (ParseException | DnsParseException e) {
-                fail(mMsg + e.getMessage());
+                mErrorMsg = mMsg + e.getMessage();
+                mLatch.countDown();
+                return;
             }
             Log.d(TAG, "Reported blob: " + byteArrayToHexString(answer));
             mLatch.countDown();
@@ -243,10 +253,12 @@
 
         @Override
         public void onError(@NonNull DnsResolver.DnsException error) {
-            fail(mMsg + error.getMessage());
+            mErrorMsg = mMsg + error.getMessage();
+            mLatch.countDown();
         }
 
         private void assertValidAnswer() {
+            assertNull(mErrorMsg);
             assertNotNull(mMsg + " No valid answer", mDnsAnswer);
             assertEquals(mMsg + " Unexpected error: reported rcode" + mRcode +
                     " blob's rcode " + mDnsAnswer.getRcode(), mRcode, mDnsAnswer.getRcode());
@@ -402,20 +414,18 @@
     public void doTestRawQueryNXDomainWithPrivateDns(Executor executor)
             throws InterruptedException {
         final String msg = "RawQuery " + TEST_NX_DOMAIN + " with private DNS";
-
         // Enable private DNS strict mode and set server to dns.google before doing NxDomain test.
         // b/144521720
         Settings.Global.putString(mCR, Settings.Global.PRIVATE_DNS_MODE, "hostname");
         Settings.Global.putString(mCR,
                 Settings.Global.PRIVATE_DNS_SPECIFIER, GOOGLE_PRIVATE_DNS_SERVER);
-
         for (Network network :  getTestableNetworks()) {
             final Network networkForPrivateDns =
                     (network != null) ? network : mCM.getActiveNetwork();
             assertNotNull("Can't find network to await private DNS on", networkForPrivateDns);
             mCtsNetUtils.awaitPrivateDnsSetting(msg + " wait private DNS setting timeout",
                     networkForPrivateDns, GOOGLE_PRIVATE_DNS_SERVER,
-                    PRIVATE_DNS_SETTING_TIMEOUT_MS);
+                    PRIVATE_DNS_SETTING_TIMEOUT_MS, true);
             final VerifyCancelCallback callback = new VerifyCancelCallback(msg);
             mDns.rawQuery(network, TEST_NX_DOMAIN, CLASS_IN, TYPE_AAAA, FLAG_NO_CACHE_LOOKUP,
                     executor, null, callback);
@@ -508,6 +518,7 @@
         private final String mMsg;
         private final List<InetAddress> mAnswers;
         private final CancellationSignal mCancelSignal;
+        private String mErrorMsg = null;
 
         VerifyCancelInetAddressCallback(@NonNull String msg, @Nullable CancellationSignal cancel) {
             this.mMsg = msg;
@@ -541,10 +552,16 @@
             return false;
         }
 
+        public void assertNoError() {
+            assertNull(mErrorMsg);
+        }
+
         @Override
         public void onAnswer(@NonNull List<InetAddress> answerList, int rcode) {
             if (mCancelSignal != null && mCancelSignal.isCanceled()) {
-                fail(mMsg + " should not have returned any answers");
+                mErrorMsg = mMsg + " should not have returned any answers";
+                mLatch.countDown();
+                return;
             }
             for (InetAddress addr : answerList) {
                 Log.d(TAG, "Reported addr: " + addr.toString());
@@ -556,7 +573,7 @@
 
         @Override
         public void onError(@NonNull DnsResolver.DnsException error) {
-            fail(mMsg + error.getMessage());
+            mErrorMsg = mMsg + error.getMessage();
         }
     }
 
@@ -601,6 +618,7 @@
 
             assertTrue(msg + " but no answer after " + TIMEOUT_MS + "ms.",
                     callback.waitForAnswer());
+            callback.assertNoError();
             assertTrue(msg + " returned 0 results", !callback.isAnswerEmpty());
         }
     }
@@ -644,6 +662,7 @@
 
             assertTrue(msg + " but no answer after " + TIMEOUT_MS + "ms.",
                     callback.waitForAnswer());
+            callback.assertNoError();
             assertTrue(msg + " returned 0 results", !callback.isAnswerEmpty());
             assertTrue(msg + " returned Ipv6 results", !callback.hasIpv6Answer());
         }
@@ -659,6 +678,7 @@
 
             assertTrue(msg + " but no answer after " + TIMEOUT_MS + "ms.",
                     callback.waitForAnswer());
+            callback.assertNoError();
             assertTrue(msg + " returned 0 results", !callback.isAnswerEmpty());
             assertTrue(msg + " returned Ipv4 results", !callback.hasIpv4Answer());
         }
@@ -671,7 +691,6 @@
         Settings.Global.putString(mCR, Settings.Global.PRIVATE_DNS_MODE, "hostname");
         Settings.Global.putString(mCR,
                 Settings.Global.PRIVATE_DNS_SPECIFIER, INVALID_PRIVATE_DNS_SERVER);
-
         final String msg = "Test PrivateDnsBypass " + TEST_DOMAIN;
         for (Network network : testNetworks) {
             // This test cannot be ran with null network because we need to explicitly pass a
@@ -680,7 +699,7 @@
 
             // wait for private DNS setting propagating
             mCtsNetUtils.awaitPrivateDnsSetting(msg + " wait private DNS setting timeout",
-                    network, INVALID_PRIVATE_DNS_SERVER, PRIVATE_DNS_SETTING_TIMEOUT_MS);
+                    network, INVALID_PRIVATE_DNS_SERVER, PRIVATE_DNS_SETTING_TIMEOUT_MS, false);
 
             final CountDownLatch latch = new CountDownLatch(1);
             final DnsResolver.Callback<List<InetAddress>> errorCallback =
@@ -712,6 +731,7 @@
 
             assertTrue(msg + " bypass private DNS round. No answer after " + TIMEOUT_MS + "ms.",
                     callback.waitForAnswer());
+            callback.assertNoError();
             assertTrue(msg + " returned 0 results", !callback.isAnswerEmpty());
 
             // To ensure private DNS bypass still work even if passing null network.
@@ -724,6 +744,7 @@
 
             assertTrue(msg + " with null network bypass private DNS round. No answer after " +
                     TIMEOUT_MS + "ms.", callbackWithNullNetwork.waitForAnswer());
+            callbackWithNullNetwork.assertNoError();
             assertTrue(msg + " with null network returned 0 results",
                     !callbackWithNullNetwork.isAnswerEmpty());
 
@@ -745,6 +766,7 @@
 
                 assertTrue(msg + " but no answer after " + TIMEOUT_MS + "ms.",
                         callback.waitForAnswer());
+                callback.assertNoError();
                 assertTrue(msg + " returned 0 results", !callback.isAnswerEmpty());
                 assertTrue(msg + " returned " + (queryV6 ? "Ipv4" : "Ipv6") + " results",
                         queryV6 ? !callback.hasIpv4Answer() : !callback.hasIpv6Answer());
diff --git a/tests/tests/net/src/android/net/cts/MultinetworkApiTest.java b/tests/tests/net/src/android/net/cts/MultinetworkApiTest.java
index 766c55e..f123187 100644
--- a/tests/tests/net/src/android/net/cts/MultinetworkApiTest.java
+++ b/tests/tests/net/src/android/net/cts/MultinetworkApiTest.java
@@ -245,7 +245,7 @@
             for (Network network : getTestableNetworks()) {
               // Wait for private DNS setting to propagate.
               mCtsNetUtils.awaitPrivateDnsSetting("NxDomain test wait private DNS setting timeout",
-                        network, GOOGLE_PRIVATE_DNS_SERVER, PRIVATE_DNS_SETTING_TIMEOUT_MS);
+                        network, GOOGLE_PRIVATE_DNS_SERVER, PRIVATE_DNS_SETTING_TIMEOUT_MS, true);
               runResNnxDomainCheck(network.getNetworkHandle());
             }
         } finally {
diff --git a/tests/tests/net/util/java/android/net/cts/util/CtsNetUtils.java b/tests/tests/net/util/java/android/net/cts/util/CtsNetUtils.java
index f0c34e3..6214f89 100644
--- a/tests/tests/net/util/java/android/net/cts/util/CtsNetUtils.java
+++ b/tests/tests/net/util/java/android/net/cts/util/CtsNetUtils.java
@@ -56,6 +56,7 @@
     private static final String TAG = CtsNetUtils.class.getSimpleName();
     private static final int DURATION = 10000;
     private static final int SOCKET_TIMEOUT_MS = 2000;
+    private static final int PRIVATE_DNS_PROBE_MS = 1_000;
 
     public static final int HTTP_PORT = 80;
     public static final String TEST_HOST = "connectivitycheck.gstatic.com";
@@ -246,12 +247,16 @@
     }
 
     public void awaitPrivateDnsSetting(@NonNull String msg, @NonNull Network network,
-            @NonNull String server, int timeoutMs) throws InterruptedException {
+            @NonNull String server, int timeoutMs,
+            boolean requiresValidatedServers) throws InterruptedException {
         CountDownLatch latch = new CountDownLatch(1);
         NetworkRequest request = new NetworkRequest.Builder().clearCapabilities().build();
         NetworkCallback callback = new NetworkCallback() {
             @Override
             public void onLinkPropertiesChanged(Network n, LinkProperties lp) {
+                if (requiresValidatedServers && lp.getValidatedPrivateDnsServers().isEmpty()) {
+                    return;
+                }
                 if (network.equals(n) && server.equals(lp.getPrivateDnsServerName())) {
                     latch.countDown();
                 }
@@ -260,6 +265,18 @@
         mCm.registerNetworkCallback(request, callback);
         assertTrue(msg, latch.await(timeoutMs, TimeUnit.MILLISECONDS));
         mCm.unregisterNetworkCallback(callback);
+        // Wait some time for NetworkMonitor's private DNS probe to complete. If we do not do
+        // this, then the test could complete before the NetworkMonitor private DNS probe
+        // completes. This would result in tearDown disabling private DNS, and the NetworkMonitor
+        // private DNS probe getting stuck because there are no longer any private DNS servers to
+        // query. This then results in the next test not being able to change the private DNS
+        // setting within the timeout, because the NetworkMonitor thread is blocked in the
+        // private DNS probe. There is no way to know when the probe has completed: because the
+        // network is likely already validated, there is no callback that we can listen to, so
+        // just sleep.
+        if (requiresValidatedServers) {
+            Thread.sleep(PRIVATE_DNS_PROBE_MS);
+        }
     }
 
     /**