From ad61ab46f7862c7111720ec8081449fe5f674774 Mon Sep 17 00:00:00 2001
From: David Su <dysu@google.com>
Date: Thu, 12 Mar 2020 16:16:50 -0700
Subject: [PATCH] CTS WifiManagerTest: Enable scanning if not enabled

Previously some Wifi CTS tests would fail if
scanning is not enabled. Now we have a new API to
enable scanning, so call that if scanning is
disabled.

Bug: 146661308
Test: atest android.net.wifi.cts.WifiManagerTest
Change-Id: I9cc7e54302c7849ce403dbf6f14256a4de3a421d
---
 .../android/net/wifi/cts/WifiManagerTest.java | 61 ++++++++++++-------
 1 file changed, 39 insertions(+), 22 deletions(-)

diff --git a/tests/cts/net/src/android/net/wifi/cts/WifiManagerTest.java b/tests/cts/net/src/android/net/wifi/cts/WifiManagerTest.java
index f4c20e3072..870cf72955 100644
--- a/tests/cts/net/src/android/net/wifi/cts/WifiManagerTest.java
+++ b/tests/cts/net/src/android/net/wifi/cts/WifiManagerTest.java
@@ -66,6 +66,7 @@ import androidx.test.platform.app.InstrumentationRegistry;
 import com.android.compatibility.common.util.PollingCheck;
 import com.android.compatibility.common.util.ShellIdentityUtils;
 import com.android.compatibility.common.util.SystemUtil;
+import com.android.compatibility.common.util.ThrowingRunnable;
 
 import java.io.BufferedReader;
 import java.io.IOException;
@@ -76,6 +77,7 @@ import java.net.URL;
 import java.util.List;
 import java.util.Locale;
 import java.util.Objects;
+import java.util.concurrent.Callable;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
@@ -1067,6 +1069,23 @@ public class WifiManagerTest extends AndroidTestCase {
         }
     }
 
+    private void runWithScanningEnabled(ThrowingRunnable r) throws Exception {
+        boolean wasScanEnabledForTest = false;
+        if (!mWifiManager.isScanAlwaysAvailable()) {
+            ShellIdentityUtils.invokeWithShellPermissions(
+                    () -> mWifiManager.setScanAlwaysAvailable(true));
+            wasScanEnabledForTest = true;
+        }
+        try {
+            r.run();
+        } finally {
+            if (wasScanEnabledForTest) {
+                ShellIdentityUtils.invokeWithShellPermissions(
+                        () -> mWifiManager.setScanAlwaysAvailable(false));
+            }
+        }
+    }
+
     /**
      * Verify that Wi-Fi scanning is not turned off when the screen turns off while wifi is disabled
      * but location is on.
@@ -1085,17 +1104,16 @@ public class WifiManagerTest extends AndroidTestCase {
             fail("Please enable location for this test - since Marshmallow WiFi scan results are"
                     + " empty when location is disabled!");
         }
-        if(!mWifiManager.isScanAlwaysAvailable()) {
-            fail("Please enable Wi-Fi scanning for this test!");
-        }
-        setWifiEnabled(false);
-        turnScreenOn();
-        assertWifiScanningIsOn();
-        // Toggle screen and verify Wi-Fi scanning is still on.
-        turnScreenOff();
-        assertWifiScanningIsOn();
-        turnScreenOn();
-        assertWifiScanningIsOn();
+        runWithScanningEnabled(() -> {
+            setWifiEnabled(false);
+            turnScreenOn();
+            assertWifiScanningIsOn();
+            // Toggle screen and verify Wi-Fi scanning is still on.
+            turnScreenOff();
+            assertWifiScanningIsOn();
+            turnScreenOn();
+            assertWifiScanningIsOn();
+        });
     }
 
     /**
@@ -1115,17 +1133,16 @@ public class WifiManagerTest extends AndroidTestCase {
             fail("Please enable location for this test - since Marshmallow WiFi scan results are"
                     + " empty when location is disabled!");
         }
-        if(!mWifiManager.isScanAlwaysAvailable()) {
-            fail("Please enable Wi-Fi scanning for this test!");
-        }
-        setWifiEnabled(true);
-        turnScreenOn();
-        assertWifiScanningIsOn();
-        // Toggle screen and verify Wi-Fi scanning is still on.
-        turnScreenOff();
-        assertWifiScanningIsOn();
-        turnScreenOn();
-        assertWifiScanningIsOn();
+        runWithScanningEnabled(() -> {
+            setWifiEnabled(true);
+            turnScreenOn();
+            assertWifiScanningIsOn();
+            // Toggle screen and verify Wi-Fi scanning is still on.
+            turnScreenOff();
+            assertWifiScanningIsOn();
+            turnScreenOn();
+            assertWifiScanningIsOn();
+        });
     }
 
     /**