xref: /aosp_15_r20/cts/tests/tests/wifi/src/android/net/wifi/cts/TestHelper.java (revision b7c941bb3fa97aba169d73cee0bed2de8ac964bf)
1 /*
2  * Copyright (C) 2021 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package android.net.wifi.cts;
18 
19 import static android.Manifest.permission.CONNECTIVITY_INTERNAL;
20 import static android.Manifest.permission.NETWORK_SETTINGS;
21 import static android.net.ConnectivityManager.NetworkCallback.FLAG_INCLUDE_LOCATION_INFO;
22 import static android.net.NetworkCapabilities.NET_CAPABILITY_INTERNET;
23 import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED;
24 import static android.net.NetworkCapabilities.NET_CAPABILITY_OEM_PAID;
25 import static android.net.NetworkCapabilities.NET_CAPABILITY_OEM_PRIVATE;
26 import static android.net.NetworkCapabilities.TRANSPORT_WIFI;
27 import static android.net.wifi.WifiManager.STATUS_LOCAL_ONLY_CONNECTION_FAILURE_UNKNOWN;
28 import static android.net.wifi.WifiManager.STATUS_LOCAL_ONLY_CONNECTION_FAILURE_USER_REJECT;
29 import static android.os.Process.myUid;
30 
31 import static com.google.common.truth.Truth.assertThat;
32 
33 import static org.junit.Assert.assertEquals;
34 import static org.junit.Assert.assertFalse;
35 import static org.junit.Assert.assertNotNull;
36 import static org.junit.Assert.assertNull;
37 import static org.junit.Assert.assertTrue;
38 import static org.junit.Assert.fail;
39 
40 import android.annotation.NonNull;
41 import android.app.UiAutomation;
42 import android.content.Context;
43 import android.net.ConnectivityManager;
44 import android.net.MacAddress;
45 import android.net.Network;
46 import android.net.NetworkCapabilities;
47 import android.net.NetworkRequest;
48 import android.net.wifi.ScanResult;
49 import android.net.wifi.WifiConfiguration;
50 import android.net.wifi.WifiInfo;
51 import android.net.wifi.WifiManager;
52 import android.net.wifi.WifiNetworkSpecifier;
53 import android.net.wifi.WifiNetworkSuggestion;
54 import android.os.Build;
55 import android.os.SystemClock;
56 import android.os.WorkSource;
57 import android.support.test.uiautomator.UiDevice;
58 import android.text.TextUtils;
59 import android.util.ArrayMap;
60 import android.util.Log;
61 
62 import androidx.test.platform.app.InstrumentationRegistry;
63 
64 import com.android.compatibility.common.util.ApiLevelUtil;
65 import com.android.compatibility.common.util.PollingCheck;
66 import com.android.wifi.flags.Flags;
67 
68 import java.util.ArrayList;
69 import java.util.Arrays;
70 import java.util.Collections;
71 import java.util.HashSet;
72 import java.util.List;
73 import java.util.Map;
74 import java.util.Optional;
75 import java.util.Set;
76 import java.util.concurrent.CountDownLatch;
77 import java.util.concurrent.Executors;
78 import java.util.concurrent.ScheduledExecutorService;
79 import java.util.concurrent.TimeUnit;
80 
81 /**
82  * Class to hold helper methods that are repeated across wifi CTS tests.
83  */
84 public class TestHelper {
85     private static final String TAG = "WifiTestHelper";
86 
87     private final Context mContext;
88     private final WifiManager mWifiManager;
89     private final ConnectivityManager mConnectivityManager;
90     private final UiDevice mUiDevice;
91 
92     private static final int DURATION_MILLIS = 10_000;
93     private static final int DURATION_NETWORK_CONNECTION_MILLIS = 40_000;
94     private static final int DURATION_SCREEN_TOGGLE_MILLIS = 2000;
95     private static final int DURATION_UI_INTERACTION_MILLIS = 25_000;
96     private static final int SCAN_RETRY_CNT_TO_FIND_MATCHING_BSSID = 5;
97     private static List<ScanResult> sScanResults = null;
98 
TestHelper(@onNull Context context, @NonNull UiDevice uiDevice)99     public TestHelper(@NonNull Context context, @NonNull UiDevice uiDevice) {
100         mContext = context;
101         mWifiManager = context.getSystemService(WifiManager.class);
102         mConnectivityManager = context.getSystemService(ConnectivityManager.class);
103         mUiDevice = uiDevice;
104     }
105 
turnScreenOn()106     public void turnScreenOn() throws Exception {
107         mUiDevice.executeShellCommand("input keyevent KEYCODE_WAKEUP");
108         mUiDevice.executeShellCommand("wm dismiss-keyguard");
109         // Since the screen on/off intent is ordered, they will not be sent right now.
110         Thread.sleep(DURATION_SCREEN_TOGGLE_MILLIS);
111     }
112 
turnScreenOff()113     public void turnScreenOff() throws Exception {
114         mUiDevice.executeShellCommand("input keyevent KEYCODE_SLEEP");
115         // Since the screen on/off intent is ordered, they will not be sent right now.
116         Thread.sleep(DURATION_SCREEN_TOGGLE_MILLIS);
117     }
118 
119     private static class TestScanResultsCallback extends WifiManager.ScanResultsCallback {
120         private final CountDownLatch mCountDownLatch;
121         public boolean onAvailableCalled = false;
122 
TestScanResultsCallback()123         TestScanResultsCallback() {
124             mCountDownLatch = new CountDownLatch(1);
125         }
126 
await()127         public void await() throws InterruptedException {
128             mCountDownLatch.await(DURATION_MILLIS, TimeUnit.MILLISECONDS);
129         }
130 
131         @Override
onScanResultsAvailable()132         public void onScanResultsAvailable() {
133             onAvailableCalled = true;
134             mCountDownLatch.countDown();
135         }
136     }
137 
138     /**
139      * Find the first saved network available in the scan results with specific capabilities
140      * support.
141      *
142      * @param wifiManager WifiManager service
143      * @param savedNetworks List of saved networks on the device.
144      * @param capabilities Network capabilities to be matched. This parameter is ignored if set
145      *                     to 0. See AP_CAPABILITIES_BIT_XXX for the supported capabilities.
146      * @return WifiConfiguration for the network
147      */
findFirstAvailableSavedNetwork(@onNull WifiManager wifiManager, @NonNull List<WifiConfiguration> savedNetworks, long capabilities)148     public static WifiConfiguration findFirstAvailableSavedNetwork(@NonNull WifiManager wifiManager,
149             @NonNull List<WifiConfiguration> savedNetworks,
150             long capabilities) {
151         if (savedNetworks.isEmpty()) return null;
152         List<WifiConfiguration> matchingNetworks = new ArrayList<>();
153         Map<Integer, List<WifiConfiguration>> networksMap =
154                 findMatchingSavedNetworksWithBssidByBand(wifiManager, savedNetworks,
155                         capabilities, 1);
156         for (List<WifiConfiguration> configs : networksMap.values()) {
157             matchingNetworks.addAll(configs);
158         }
159         if (matchingNetworks.isEmpty()) return null;
160         return matchingNetworks.get(0);
161     }
162 
163     /**
164      * Loops through all the saved networks available in the scan results. Returns a list of
165      * WifiConfiguration with the matching bssid filled in {@link WifiConfiguration#BSSID}.
166      *
167      * Note:
168      * a) If there are more than 2 networks with the same SSID, but different credential type, then
169      * this matching may pick the wrong one.
170      *
171      * @param wifiManager   WifiManager service
172      * @param savedNetworks List of saved networks on the device.
173      * @return List of WifiConfiguration with matching bssid.
174      */
findMatchingSavedNetworksWithBssid( @onNull WifiManager wifiManager, @NonNull List<WifiConfiguration> savedNetworks, int numberOfApRequested)175     public static List<WifiConfiguration> findMatchingSavedNetworksWithBssid(
176             @NonNull WifiManager wifiManager, @NonNull List<WifiConfiguration> savedNetworks,
177             int numberOfApRequested) {
178         if (savedNetworks.isEmpty()) return Collections.emptyList();
179         List<WifiConfiguration> matchingNetworksWithBssids = new ArrayList<>();
180         Map<Integer, List<WifiConfiguration>> networksMap =
181                 findMatchingSavedNetworksWithBssidByBand(wifiManager, savedNetworks,
182                         0, numberOfApRequested);
183         for (List<WifiConfiguration> configs : networksMap.values()) {
184             matchingNetworksWithBssids.addAll(configs);
185         }
186         return matchingNetworksWithBssids;
187     }
188 
189     public static final long AP_CAPABILITY_BIT_WIFI7 = 1 << 0;
190     public static final long AP_CAPABILITY_BIT_TWT_RESPONDER = 1 << 1;
191 
192     /**
193      * Check whether scan result matches with the capabilities provided.
194      *
195      * @param scanResult Scan result
196      * @param capabilities Capabilities to be matched. See AP_CAPABILITY_BIT_XXX for the
197      *                     available bits.
198      * @return true if the scan result matches or capabilities is 0 , otherwise false.
199      */
isMatchedScanResult(ScanResult scanResult, long capabilities)200     public static boolean isMatchedScanResult(ScanResult scanResult, long capabilities) {
201         if (capabilities == 0) return true;
202         if ((capabilities & AP_CAPABILITY_BIT_WIFI7) == AP_CAPABILITY_BIT_WIFI7 && (
203                 scanResult.getWifiStandard()
204                         != ScanResult.WIFI_STANDARD_11BE)) {
205             return false;
206         }
207         if ((capabilities & AP_CAPABILITY_BIT_TWT_RESPONDER) == AP_CAPABILITY_BIT_TWT_RESPONDER
208                 && !scanResult.isTwtResponder()) {
209             return false;
210         }
211         return true;
212     }
213 
214     /**
215      * Loops through all the saved networks available in the scan results. Returns a map of lists of
216      * WifiConfiguration with the matching bssid filled in {@link WifiConfiguration#BSSID}.
217      *
218      * Note:
219      * a) If there are more than 2 networks with the same SSID, but different credential type, then
220      * this matching may pick the wrong one.
221      *
222      * @param wifiManager   WifiManager service
223      * @param savedNetworks List of saved networks on the device.
224      * @param  capabilities AP capabilities to be matched. See AP_CAPABILITY_BIT_XXX for the
225      *                     supported capabilities. This parameter is ignored if set to 0.
226      *
227      * @param numberOfApRequested Number of APs requested
228      * @return Map from band to the list of WifiConfiguration with matching bssid.
229      */
findMatchingSavedNetworksWithBssidByBand( @onNull WifiManager wifiManager, @NonNull List<WifiConfiguration> savedNetworks, long capabilities, int numberOfApRequested)230     public static Map<Integer, List<WifiConfiguration>> findMatchingSavedNetworksWithBssidByBand(
231             @NonNull WifiManager wifiManager, @NonNull List<WifiConfiguration> savedNetworks,
232             long capabilities, int numberOfApRequested) {
233         if (savedNetworks.isEmpty()) return Collections.emptyMap();
234         Set<String> bssidSet = new HashSet<>();
235         Map<Integer, List<WifiConfiguration>> matchingNetworksWithBssids = new ArrayMap<>();
236         for (int i = 0; i < SCAN_RETRY_CNT_TO_FIND_MATCHING_BSSID; i++) {
237             int count = 0;
238             // Trigger a scan to get fresh scan results.
239             TestScanResultsCallback scanResultsCallback = new TestScanResultsCallback();
240             try {
241                 wifiManager.registerScanResultsCallback(
242                         Executors.newSingleThreadExecutor(), scanResultsCallback);
243                 wifiManager.startScan(new WorkSource(myUid()));
244                 // now wait for callback
245                 scanResultsCallback.await();
246             } catch (InterruptedException e) {
247             } finally {
248                 wifiManager.unregisterScanResultsCallback(scanResultsCallback);
249             }
250             sScanResults = wifiManager.getScanResults();
251             if (sScanResults == null || sScanResults.isEmpty()) continue;
252             for (ScanResult scanResult : sScanResults) {
253                 if (!isMatchedScanResult(scanResult, capabilities) || bssidSet.contains(
254                         scanResult.BSSID)) {
255                     continue;
256                 }
257                 WifiConfiguration matchingNetwork = savedNetworks.stream()
258                         .filter(network -> TextUtils.equals(
259                                 scanResult.SSID, WifiInfo.sanitizeSsid(network.SSID)))
260                         .findAny()
261                         .orElse(null);
262                 if (matchingNetwork != null) {
263                     // make a copy in case we have 2 bssid's for the same network.
264                     WifiConfiguration matchingNetworkCopy = new WifiConfiguration(matchingNetwork);
265                     matchingNetworkCopy.BSSID = scanResult.BSSID;
266                     bssidSet.add(scanResult.BSSID);
267                     List<WifiConfiguration> bandConfigs =
268                             matchingNetworksWithBssids.computeIfAbsent(
269                                     scanResult.getBand(), k -> new ArrayList<>());
270                     bandConfigs.add(matchingNetworkCopy);
271                 }
272             }
273             if (bssidSet.size() >= numberOfApRequested
274                     && !matchingNetworksWithBssids.isEmpty()) break;
275         }
276         return matchingNetworksWithBssids;
277     }
278 
279     /**
280      * Convert the provided saved network to a corresponding suggestion builder.
281      */
282     public static WifiNetworkSuggestion.Builder
createSuggestionBuilderWithCredentialFromSavedNetworkWithBssid( @onNull WifiConfiguration network)283             createSuggestionBuilderWithCredentialFromSavedNetworkWithBssid(
284             @NonNull WifiConfiguration network) {
285         WifiNetworkSuggestion.Builder suggestionBuilder = new WifiNetworkSuggestion.Builder()
286                 .setSsid(WifiInfo.sanitizeSsid(network.SSID))
287                 .setBssid(MacAddress.fromString(network.BSSID));
288         if (network.preSharedKey != null) {
289             if (network.allowedKeyManagement.get(WifiConfiguration.KeyMgmt.WPA_PSK)) {
290                 suggestionBuilder.setWpa2Passphrase(WifiInfo.sanitizeSsid(network.preSharedKey));
291             } else if (network.allowedKeyManagement.get(WifiConfiguration.KeyMgmt.SAE)) {
292                 suggestionBuilder.setWpa3Passphrase(WifiInfo.sanitizeSsid(network.preSharedKey));
293             } else {
294                 fail("Unsupported security type found in saved networks");
295             }
296         } else if (network.allowedKeyManagement.get(WifiConfiguration.KeyMgmt.OWE)) {
297             suggestionBuilder.setIsEnhancedOpen(true);
298         } else if (!network.allowedKeyManagement.get(WifiConfiguration.KeyMgmt.NONE)) {
299             fail("Unsupported security type found in saved networks");
300         }
301         suggestionBuilder.setIsHiddenSsid(network.hiddenSSID);
302         return suggestionBuilder;
303     }
304 
305 
306     /**
307      * Convert the provided saved network to a corresponding specifier builder.
308      */
createSpecifierBuilderWithCredentialFromSavedNetwork( @onNull WifiConfiguration network, boolean useChannel)309     public static WifiNetworkSpecifier.Builder createSpecifierBuilderWithCredentialFromSavedNetwork(
310             @NonNull WifiConfiguration network, boolean useChannel) {
311         WifiNetworkSpecifier.Builder specifierBuilder = new WifiNetworkSpecifier.Builder()
312                 .setSsid(WifiInfo.sanitizeSsid(network.SSID));
313         if (network.preSharedKey != null) {
314             if (network.allowedKeyManagement.get(WifiConfiguration.KeyMgmt.WPA_PSK)) {
315                 specifierBuilder.setWpa2Passphrase(WifiInfo.sanitizeSsid(network.preSharedKey));
316             } else if (network.allowedKeyManagement.get(WifiConfiguration.KeyMgmt.SAE)) {
317                 specifierBuilder.setWpa3Passphrase(WifiInfo.sanitizeSsid(network.preSharedKey));
318             } else {
319                 fail("Unsupported security type found in saved networks");
320             }
321         } else if (network.allowedKeyManagement.get(WifiConfiguration.KeyMgmt.OWE)) {
322             specifierBuilder.setIsEnhancedOpen(true);
323         } else if (!network.allowedKeyManagement.get(WifiConfiguration.KeyMgmt.NONE)) {
324             fail("Unsupported security type found in saved networks");
325         }
326         specifierBuilder.setIsHiddenSsid(network.hiddenSSID);
327         if (sScanResults != null && useChannel) {
328             Optional<ScanResult> matchedResult = sScanResults
329                     .stream()
330                     .filter(scanResult -> TextUtils.equals(scanResult.SSID,
331                             WifiInfo.sanitizeSsid(network.SSID))
332                             && TextUtils.equals(scanResult.BSSID, network.BSSID)).findAny();
333             matchedResult.ifPresent(
334                     scanResult -> specifierBuilder.setPreferredChannelsFrequenciesMhz(
335                             new int[]{scanResult.frequency}));
336         }
337         return specifierBuilder;
338     }
339 
340     /**
341      * Convert the provided saved network to a corresponding specifier builder.
342      */
343     public static WifiNetworkSpecifier.Builder
createSpecifierBuilderWithCredentialFromSavedNetworkWithBssid( @onNull WifiConfiguration network)344             createSpecifierBuilderWithCredentialFromSavedNetworkWithBssid(
345             @NonNull WifiConfiguration network) {
346         return createSpecifierBuilderWithCredentialFromSavedNetwork(network, false)
347                 .setBssid(MacAddress.fromString(network.BSSID));
348     }
349 
350     private static class TestLocalOnlyListener implements WifiManager
351             .LocalOnlyConnectionFailureListener {
352         private CountDownLatch mBlocker;
353         public boolean onFailureCalled = false;
354         public int failureReason = STATUS_LOCAL_ONLY_CONNECTION_FAILURE_UNKNOWN;
TestLocalOnlyListener()355         TestLocalOnlyListener() {
356             mBlocker = new CountDownLatch(1);
357         }
358 
359         @Override
onConnectionFailed( @ndroidx.annotation.NonNull WifiNetworkSpecifier wifiNetworkSpecifier, int failureReason)360         public void onConnectionFailed(
361                 @androidx.annotation.NonNull WifiNetworkSpecifier wifiNetworkSpecifier,
362                 int failureReason) {
363             mBlocker.countDown();
364             onFailureCalled = true;
365             this.failureReason = failureReason;
366         }
367 
await(long timeout)368         public boolean await(long timeout) throws Exception {
369             return mBlocker.await(timeout, TimeUnit.MILLISECONDS);
370         }
371 
372     }
373 
374     public static class TestNetworkCallback extends ConnectivityManager.NetworkCallback {
375         private CountDownLatch mBlocker;
376         public boolean onAvailableCalled = false;
377         public boolean onUnavailableCalled = false;
378         public boolean onLostCalled = false;
379         public boolean onLosingCalled = false;
380         public NetworkCapabilities networkCapabilities;
381 
TestNetworkCallback()382         TestNetworkCallback() {
383             mBlocker = new CountDownLatch(1);
384         }
385 
TestNetworkCallback(int flags)386         TestNetworkCallback(int flags) {
387             super(flags);
388             mBlocker = new CountDownLatch(1);
389         }
390 
await(long timeout)391         public boolean await(long timeout) throws Exception {
392             return mBlocker.await(timeout, TimeUnit.MILLISECONDS);
393         }
394 
395         @Override
onAvailable(Network network)396         public void onAvailable(Network network) {
397             Log.i(TAG, "onAvailable " + network);
398             onAvailableCalled = true;
399         }
400 
401         @Override
onCapabilitiesChanged(Network network, NetworkCapabilities networkCapabilities)402         public void onCapabilitiesChanged(Network network,
403                 NetworkCapabilities networkCapabilities) {
404             Log.i(TAG, "onCapabilitiesChanged " + network);
405             this.networkCapabilities = networkCapabilities;
406             mBlocker.countDown();
407         }
408 
409         @Override
onUnavailable()410         public void onUnavailable() {
411             Log.i(TAG, "onUnavailable ");
412             onUnavailableCalled = true;
413             mBlocker.countDown();
414         }
415 
416         @Override
onLosing(Network network, int maxMsToLive)417         public void onLosing(Network network, int maxMsToLive) {
418             Log.i(TAG, "onLosing + " + maxMsToLive);
419             onLosingCalled = true;
420             mBlocker.countDown();
421         }
422 
423         @Override
onLost(Network network)424         public void onLost(Network network) {
425             Log.i(TAG, "onLost ");
426             onLostCalled = true;
427             mBlocker.countDown();
428         }
429 
waitForAnyCallback(int timeout)430         boolean waitForAnyCallback(int timeout) {
431             try {
432                 boolean noTimeout = mBlocker.await(timeout, TimeUnit.MILLISECONDS);
433                 mBlocker = new CountDownLatch(1);
434                 return noTimeout;
435             } catch (InterruptedException e) {
436                 return false;
437             }
438         }
439     }
440 
createTestNetworkCallback()441     private static TestNetworkCallback createTestNetworkCallback() {
442         if (ApiLevelUtil.isAtLeast(Build.VERSION_CODES.S)) {
443             // flags for NetworkCallback only introduced in S.
444             return new TestNetworkCallback(FLAG_INCLUDE_LOCATION_INFO);
445         } else {
446             return new TestNetworkCallback();
447         }
448     }
449 
450     @NonNull
getWifiInfo(@onNull NetworkCapabilities networkCapabilities)451     private WifiInfo getWifiInfo(@NonNull NetworkCapabilities networkCapabilities) {
452         if (ApiLevelUtil.isAtLeast(Build.VERSION_CODES.S)) {
453             // WifiInfo in transport info, only available in S.
454             return (WifiInfo) networkCapabilities.getTransportInfo();
455         } else {
456             return mWifiManager.getConnectionInfo();
457         }
458     }
459 
assertConnectionEquals(@onNull WifiConfiguration network, @NonNull WifiInfo wifiInfo)460     private static void assertConnectionEquals(@NonNull WifiConfiguration network,
461             @NonNull WifiInfo wifiInfo) {
462         assertThat(network.SSID).isEqualTo(wifiInfo.getSSID());
463         assertThat(network.BSSID).isEqualTo(wifiInfo.getBSSID());
464     }
465 
466     private static class TestActionListener implements WifiManager.ActionListener {
467         private final CountDownLatch mCountDownLatch;
468         public boolean onSuccessCalled = false;
469         public boolean onFailedCalled = false;
470 
TestActionListener(CountDownLatch countDownLatch)471         TestActionListener(CountDownLatch countDownLatch) {
472             mCountDownLatch = countDownLatch;
473         }
474 
475         @Override
onSuccess()476         public void onSuccess() {
477             onSuccessCalled = true;
478             mCountDownLatch.countDown();
479         }
480 
481         @Override
onFailure(int reason)482         public void onFailure(int reason) {
483             onFailedCalled = true;
484             mCountDownLatch.countDown();
485         }
486     }
487 
488     /**
489      * Triggers connection to one of the saved networks using {@link WifiManager#connect(
490      * WifiConfiguration, WifiManager.ActionListener)}
491      *
492      * @param network saved network from the device to use for the connection.
493      *
494      * @return NetworkCallback used for the connection (can be used by client to release the
495      * connection.
496      */
testConnectionFlowWithConnect( @onNull WifiConfiguration network)497     public ConnectivityManager.NetworkCallback testConnectionFlowWithConnect(
498             @NonNull WifiConfiguration network) throws Exception {
499         CountDownLatch countDownLatchAl = new CountDownLatch(1);
500         TestActionListener actionListener = new TestActionListener(countDownLatchAl);
501         TestNetworkCallback testNetworkCallback = createTestNetworkCallback();
502         UiAutomation uiAutomation = InstrumentationRegistry.getInstrumentation().getUiAutomation();
503         try {
504             uiAutomation.adoptShellPermissionIdentity();
505             // File a callback for wifi network.
506             mConnectivityManager.registerNetworkCallback(
507                     new NetworkRequest.Builder()
508                             .addTransportType(NetworkCapabilities.TRANSPORT_WIFI)
509                             .addCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET)
510                             // Needed to ensure that the restricted concurrent connection does not
511                             // match this request.
512                             .addForbiddenCapability(NET_CAPABILITY_OEM_PAID)
513                             .addForbiddenCapability(NET_CAPABILITY_OEM_PRIVATE)
514                             .build(),
515                     testNetworkCallback);
516             // Trigger the connection.
517             mWifiManager.connect(network, actionListener);
518             // now wait for action listener callback
519             assertThat(countDownLatchAl.await(
520                     DURATION_NETWORK_CONNECTION_MILLIS, TimeUnit.MILLISECONDS)).isTrue();
521             // check if we got the success callback
522             assertThat(actionListener.onSuccessCalled).isTrue();
523 
524             // Wait for connection to complete & ensure we are connected to the saved network.
525             assertThat(testNetworkCallback.waitForAnyCallback(DURATION_NETWORK_CONNECTION_MILLIS))
526                     .isTrue();
527             assertThat(testNetworkCallback.onAvailableCalled).isTrue();
528             final WifiInfo wifiInfo = getWifiInfo(testNetworkCallback.networkCapabilities);
529             assertConnectionEquals(network, wifiInfo);
530             if (ApiLevelUtil.isAtLeast(Build.VERSION_CODES.S)) {
531                 // User connections should always be primary.
532                 assertThat(wifiInfo.isPrimary()).isTrue();
533             }
534         } catch (Throwable e /* catch assertions & exceptions */) {
535             // Unregister the network callback in case of any failure (since we don't end up
536             // returning the network callback to the caller).
537             try {
538                 mConnectivityManager.unregisterNetworkCallback(testNetworkCallback);
539             } catch (IllegalArgumentException ie) { }
540             throw e;
541         } finally {
542             uiAutomation.dropShellPermissionIdentity();
543         }
544         return testNetworkCallback;
545     }
546 
547     /**
548      * Tests the entire connection success flow using the provided suggestion.
549      *
550      * Note: The caller needs to invoke this after acquiring shell identity.
551      *
552      * @param network saved network from the device to use for the connection.
553      * @param suggestion suggestion to use for the connection.
554      * @param executorService Excutor service to run scan periodically (to trigger connection).
555      * @param restrictedNetworkCapabilities Whether this connection should be restricted with
556      *                                    the provided capability.
557      *
558      * @param isRestricted whether the suggestion is for a restricted network
559      * @return NetworkCallback used for the connection (can be used by client to release the
560      * connection.
561      */
testConnectionFlowWithSuggestionWithShellIdentity( WifiConfiguration network, WifiNetworkSuggestion suggestion, @NonNull ScheduledExecutorService executorService, @NonNull Set<Integer> restrictedNetworkCapabilities, boolean isRestricted)562     public ConnectivityManager.NetworkCallback testConnectionFlowWithSuggestionWithShellIdentity(
563             WifiConfiguration network, WifiNetworkSuggestion suggestion,
564             @NonNull ScheduledExecutorService executorService,
565             @NonNull Set<Integer> restrictedNetworkCapabilities,
566             boolean isRestricted) throws Exception {
567         return testConnectionFlowWithSuggestionInternal(
568                 network, suggestion, executorService, restrictedNetworkCapabilities, true,
569                 isRestricted);
570     }
571 
572     /**
573      * Tests the entire connection success flow using the provided suggestion.
574      *
575      * Note: The helper method drops the shell identity, so don't use this if the caller already
576      * adopted shell identity.
577      *
578      * @param network saved network from the device to use for the connection.
579      * @param suggestion suggestion to use for the connection.
580      * @param executorService Excutor service to run scan periodically (to trigger connection).
581      * @param restrictedNetworkCapabilities Whether this connection should be restricted with
582      *                                    the provided capability.
583      *
584      * @param isRestricted whether the suggestion is for a restricted network
585      * @return NetworkCallback used for the connection (can be used by client to release the
586      * connection.
587      */
testConnectionFlowWithSuggestion( WifiConfiguration network, WifiNetworkSuggestion suggestion, @NonNull ScheduledExecutorService executorService, @NonNull Set<Integer> restrictedNetworkCapabilities, boolean isRestricted)588     public ConnectivityManager.NetworkCallback testConnectionFlowWithSuggestion(
589             WifiConfiguration network, WifiNetworkSuggestion suggestion,
590             @NonNull ScheduledExecutorService executorService,
591             @NonNull Set<Integer> restrictedNetworkCapabilities,
592             boolean isRestricted) throws Exception {
593         final UiAutomation uiAutomation =
594                 InstrumentationRegistry.getInstrumentation().getUiAutomation();
595         try {
596             uiAutomation.adoptShellPermissionIdentity(NETWORK_SETTINGS, CONNECTIVITY_INTERNAL);
597             return testConnectionFlowWithSuggestionWithShellIdentity(
598                     network, suggestion, executorService, restrictedNetworkCapabilities,
599                     isRestricted);
600         } finally {
601             uiAutomation.dropShellPermissionIdentity();
602         }
603     }
604 
605     /**
606      * Tests the connection failure flow using the provided suggestion.
607      *
608      * @param network saved network from the device to use for the connection.
609      * @param suggestion suggestion to use for the connection.
610      * @param executorService Excutor service to run scan periodically (to trigger connection).
611      * @param restrictedNetworkCapabilities Whether this connection should be restricted with
612      *                                    the provided capability.
613      *
614      * @return NetworkCallback used for the connection (can be used by client to release the
615      * connection.
616      */
testConnectionFailureFlowWithSuggestion( WifiConfiguration network, WifiNetworkSuggestion suggestion, @NonNull ScheduledExecutorService executorService, @NonNull Set<Integer> restrictedNetworkCapabilities)617     public ConnectivityManager.NetworkCallback testConnectionFailureFlowWithSuggestion(
618             WifiConfiguration network, WifiNetworkSuggestion suggestion,
619             @NonNull ScheduledExecutorService executorService,
620             @NonNull Set<Integer> restrictedNetworkCapabilities) throws Exception {
621         final UiAutomation uiAutomation =
622                 InstrumentationRegistry.getInstrumentation().getUiAutomation();
623         try {
624             uiAutomation.adoptShellPermissionIdentity(NETWORK_SETTINGS, CONNECTIVITY_INTERNAL);
625             return testConnectionFlowWithSuggestionInternal(
626                     network, suggestion, executorService, restrictedNetworkCapabilities, false,
627                     false/* restrictedNetwork */);
628         } finally {
629             uiAutomation.dropShellPermissionIdentity();
630         }
631     }
632 
633     /**
634      * Tests the entire connection success/failure flow using the provided suggestion.
635      *
636      * @param network saved network from the device to use for the connection.
637      * @param suggestion suggestion to use for the connection.
638      * @param executorService Excutor service to run scan periodically (to trigger connection).
639      * @param restrictedNetworkCapabilities Whether this connection should be restricted with
640      *                                    the provided capability.
641      * @param expectConnectionSuccess Whether to expect connection success or not.
642      *
643      * @param isRestricted whether the suggestion is for a restricted network
644      * @return NetworkCallback used for the connection (can be used by client to release the
645      * connection.
646      */
testConnectionFlowWithSuggestionInternal( WifiConfiguration network, WifiNetworkSuggestion suggestion, @NonNull ScheduledExecutorService executorService, @NonNull Set<Integer> restrictedNetworkCapabilities, boolean expectConnectionSuccess, boolean isRestricted)647     private ConnectivityManager.NetworkCallback testConnectionFlowWithSuggestionInternal(
648             WifiConfiguration network, WifiNetworkSuggestion suggestion,
649             @NonNull ScheduledExecutorService executorService,
650             @NonNull Set<Integer> restrictedNetworkCapabilities,
651             boolean expectConnectionSuccess, boolean isRestricted) throws Exception {
652         // File the network request & wait for the callback.
653         TestNetworkCallback testNetworkCallback = createTestNetworkCallback();
654         try {
655             // File a request for restricted (oem paid) wifi network.
656             NetworkRequest.Builder nrBuilder = new NetworkRequest.Builder()
657                     .addTransportType(TRANSPORT_WIFI)
658                     .addCapability(NET_CAPABILITY_INTERNET);
659             if (restrictedNetworkCapabilities.isEmpty() && !isRestricted) {
660                 // If not a restricted connection, a network callback is sufficient.
661                 mConnectivityManager.registerNetworkCallback(
662                         nrBuilder.build(), testNetworkCallback);
663             } else {
664                 for (Integer restrictedNetworkCapability : restrictedNetworkCapabilities) {
665                     nrBuilder.addCapability(restrictedNetworkCapability);
666                 }
667                 nrBuilder.removeCapability(NET_CAPABILITY_NOT_RESTRICTED);
668                 mConnectivityManager.requestNetwork(nrBuilder.build(), testNetworkCallback);
669             }
670             // Add wifi network suggestion.
671             assertThat(mWifiManager.addNetworkSuggestions(Arrays.asList(suggestion)))
672                     .isEqualTo(WifiManager.STATUS_NETWORK_SUGGESTIONS_SUCCESS);
673             // Wait for the request to reach the wifi stack before kick-start periodic scans.
674             Thread.sleep(100);
675             // Step: Trigger scans periodically to trigger network selection quicker.
676             executorService.scheduleAtFixedRate(() -> {
677                 if (!mWifiManager.startScan()) {
678                     Log.w(TAG, "Failed to trigger scan");
679                 }
680             }, 0, DURATION_MILLIS, TimeUnit.MILLISECONDS);
681             if (expectConnectionSuccess) {
682                 // now wait for connection to complete and wait for callback
683                 assertThat(testNetworkCallback
684                         .waitForAnyCallback(DURATION_NETWORK_CONNECTION_MILLIS)).isTrue();
685                 assertThat(testNetworkCallback.onAvailableCalled).isTrue();
686                 final WifiInfo wifiInfo = getWifiInfo(testNetworkCallback.networkCapabilities);
687                 assertConnectionEquals(network, wifiInfo);
688                 assertThat(wifiInfo.isTrusted()).isTrue();
689                 assertThat(wifiInfo.isRestricted()).isEqualTo(isRestricted);
690                 WifiInfo redact = wifiInfo
691                         .makeCopy(NetworkCapabilities.REDACT_FOR_ACCESS_FINE_LOCATION);
692                 assertThat(wifiInfo.getInformationElements()).isNotNull();
693                 assertThat(redact.getInformationElements()).isNull();
694                 assertThat(redact.getApplicableRedactions()).isEqualTo(
695                         NetworkCapabilities.REDACT_FOR_ACCESS_FINE_LOCATION
696                                 | NetworkCapabilities.REDACT_FOR_LOCAL_MAC_ADDRESS
697                                 | NetworkCapabilities.REDACT_FOR_NETWORK_SETTINGS);
698                 if (ApiLevelUtil.isAtLeast(Build.VERSION_CODES.S)) {
699                     // If STA concurrency for restricted connection is supported, this should not
700                     // be the primary connection.
701                     if (!restrictedNetworkCapabilities.isEmpty()
702                             && mWifiManager.isStaConcurrencyForRestrictedConnectionsSupported()) {
703                         assertThat(wifiInfo.isPrimary()).isFalse();
704                     } else {
705                         assertThat(wifiInfo.isPrimary()).isTrue();
706                     }
707                 }
708             } else {
709                 // now wait for connection to timeout.
710                 assertThat(testNetworkCallback
711                         .waitForAnyCallback(DURATION_NETWORK_CONNECTION_MILLIS)).isFalse();
712             }
713         } catch (Throwable e /* catch assertions & exceptions */) {
714             try {
715                 mConnectivityManager.unregisterNetworkCallback(testNetworkCallback);
716             } catch (IllegalArgumentException ie) { }
717             throw e;
718         } finally {
719             executorService.shutdown();
720         }
721         return testNetworkCallback;
722     }
723 
724     private static class TestNetworkRequestMatchCallback implements
725             WifiManager.NetworkRequestMatchCallback {
726         public boolean onRegistrationCalled = false;
727         public boolean onAbortCalled = false;
728         public boolean onMatchCalled = false;
729         public boolean onConnectSuccessCalled = false;
730         public boolean onConnectFailureCalled = false;
731         private CountDownLatch mBlocker;
732 
733         public WifiManager.NetworkRequestUserSelectionCallback userSelectionCallback = null;
734         public List<ScanResult> matchedScanResults = null;
735 
TestNetworkRequestMatchCallback()736         TestNetworkRequestMatchCallback() {
737             mBlocker = new CountDownLatch(1);
738         }
739 
waitForAnyCallback(int timeout)740         public boolean waitForAnyCallback(int timeout) {
741             try {
742                 boolean noTimeout = mBlocker.await(timeout, TimeUnit.MILLISECONDS);
743                 mBlocker = new CountDownLatch(1);
744                 return noTimeout;
745             } catch (InterruptedException e) {
746                 return false;
747             }
748         }
749 
750         @Override
onUserSelectionCallbackRegistration( WifiManager.NetworkRequestUserSelectionCallback userSelectionCallback)751         public void onUserSelectionCallbackRegistration(
752                 WifiManager.NetworkRequestUserSelectionCallback userSelectionCallback) {
753             Log.d(TAG, "onUserSelectionCallbackRegistration");
754             onRegistrationCalled = true;
755             this.userSelectionCallback = userSelectionCallback;
756             mBlocker.countDown();
757         }
758 
759         @Override
onAbort()760         public void onAbort() {
761             Log.d(TAG, "onAbort");
762             onAbortCalled = true;
763             mBlocker.countDown();
764         }
765 
766         @Override
onMatch(List<ScanResult> scanResults)767         public void onMatch(List<ScanResult> scanResults) {
768             Log.d(TAG, "onMatch");
769             // This can be invoked multiple times. So, ignore after the first one to avoid
770             // disturbing the rest of the test sequence.
771             if (onMatchCalled) return;
772             onMatchCalled = true;
773             matchedScanResults = scanResults;
774             mBlocker.countDown();
775         }
776 
777         @Override
onUserSelectionConnectSuccess(WifiConfiguration config)778         public void onUserSelectionConnectSuccess(WifiConfiguration config) {
779             Log.d(TAG, "onUserSelectionConnectSuccess");
780             onConnectSuccessCalled = true;
781             mBlocker.countDown();
782         }
783 
784         @Override
onUserSelectionConnectFailure(WifiConfiguration config)785         public void onUserSelectionConnectFailure(WifiConfiguration config) {
786             Log.d(TAG, "onUserSelectionConnectFailure");
787             onConnectFailureCalled = true;
788             mBlocker.countDown();
789         }
790     }
791 
handleUiInteractions(WifiConfiguration network, boolean shouldUserReject)792     private void handleUiInteractions(WifiConfiguration network, boolean shouldUserReject) {
793         TestNetworkRequestMatchCallback networkRequestMatchCallback =
794                 new TestNetworkRequestMatchCallback();
795         mWifiManager.registerNetworkRequestMatchCallback(
796                 Executors.newSingleThreadExecutor(), networkRequestMatchCallback);
797         long start = SystemClock.elapsedRealtime();
798         try {
799             // 1. Wait for registration callback.
800             while (!networkRequestMatchCallback.onRegistrationCalled
801                     && SystemClock.elapsedRealtime() - start < DURATION_UI_INTERACTION_MILLIS) {
802                 networkRequestMatchCallback.waitForAnyCallback(1000);
803                 if (networkRequestMatchCallback.onAbortCalled) {
804                     networkRequestMatchCallback.onAbortCalled = false;
805                     mWifiManager.registerNetworkRequestMatchCallback(
806                             Executors.newSingleThreadExecutor(), networkRequestMatchCallback);
807                 }
808             }
809             assertThat(networkRequestMatchCallback.onRegistrationCalled).isTrue();
810             assertThat(networkRequestMatchCallback.userSelectionCallback).isNotNull();
811 
812             // 2. Wait for matching scan results
813             networkRequestMatchCallback.waitForAnyCallback(DURATION_UI_INTERACTION_MILLIS);
814             assertThat(networkRequestMatchCallback.onMatchCalled).isTrue();
815             assertThat(networkRequestMatchCallback.matchedScanResults).isNotNull();
816             assertThat(networkRequestMatchCallback.matchedScanResults.size()).isAtLeast(1);
817 
818             // 3. Trigger connection to one of the matched networks or reject the request.
819             if (shouldUserReject) {
820                 networkRequestMatchCallback.userSelectionCallback.reject();
821             } else {
822                 networkRequestMatchCallback.userSelectionCallback.select(network);
823             }
824 
825             // 4. Wait for connection success or abort.
826             networkRequestMatchCallback.waitForAnyCallback(DURATION_UI_INTERACTION_MILLIS);
827             if (shouldUserReject) {
828                 assertThat(networkRequestMatchCallback.onAbortCalled).isTrue();
829             } else {
830                 assertThat(networkRequestMatchCallback.onConnectSuccessCalled).isTrue();
831             }
832         } finally {
833             mWifiManager.unregisterNetworkRequestMatchCallback(networkRequestMatchCallback);
834         }
835     }
836 
837     /**
838      * Tests the entire connection flow using the provided specifier,
839      *
840      * Note: The caller needs to invoke this after acquiring shell identity.
841      *
842      * @param specifier Specifier to use for network request.
843      * @param shouldUserReject Whether to simulate user rejection or not.
844      *
845      * @return NetworkCallback used for the connection (can be used by client to release the
846      * connection.
847      */
testConnectionFlowWithSpecifierWithShellIdentity( WifiConfiguration network, WifiNetworkSpecifier specifier, boolean shouldUserReject)848     public ConnectivityManager.NetworkCallback testConnectionFlowWithSpecifierWithShellIdentity(
849             WifiConfiguration network, WifiNetworkSpecifier specifier, boolean shouldUserReject)
850             throws Exception {
851         // File the network request & wait for the callback.
852         TestNetworkCallback testNetworkCallback = createTestNetworkCallback();
853         TestLocalOnlyListener localOnlyListener = new TestLocalOnlyListener();
854         mWifiManager.addLocalOnlyConnectionFailureListener(Executors.newSingleThreadExecutor(),
855                 localOnlyListener);
856         try {
857             // File a request for wifi network.
858             mConnectivityManager.requestNetwork(
859                     new NetworkRequest.Builder()
860                             .addTransportType(NetworkCapabilities.TRANSPORT_WIFI)
861                             .removeCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET)
862                             .setNetworkSpecifier(specifier)
863                             .build(),
864                     testNetworkCallback);
865             handleUiInteractions(network, shouldUserReject);
866             // now wait for callback
867             assertThat(testNetworkCallback.waitForAnyCallback(DURATION_NETWORK_CONNECTION_MILLIS))
868                     .isTrue();
869             if (shouldUserReject) {
870                 if (Flags.localOnlyConnectionOptimization()) {
871                     assertThat(localOnlyListener.await(DURATION_MILLIS)).isTrue();
872                     assertThat(localOnlyListener.failureReason)
873                             .isEqualTo(STATUS_LOCAL_ONLY_CONNECTION_FAILURE_USER_REJECT);
874                 }
875                 assertThat(testNetworkCallback.onUnavailableCalled).isTrue();
876             } else {
877                 assertThat(localOnlyListener.await(DURATION_MILLIS)).isFalse();
878                 assertThat(testNetworkCallback.onAvailableCalled).isTrue();
879                 final WifiInfo wifiInfo = getWifiInfo(testNetworkCallback.networkCapabilities);
880                 assertConnectionEquals(network, wifiInfo);
881                 if (ApiLevelUtil.isAtLeast(Build.VERSION_CODES.S)) {
882                     // If STA concurrency for local only connection is supported, this should not
883                     // be the primary connection.
884                     if (mWifiManager.isStaConcurrencyForLocalOnlyConnectionsSupported()) {
885                         assertThat(wifiInfo.isPrimary()).isFalse();
886                         assertConnectionEquals(network, mWifiManager.getConnectionInfo());
887                     } else {
888                         assertThat(wifiInfo.isPrimary()).isTrue();
889                     }
890                 }
891                 assertThat(localOnlyListener.onFailureCalled).isFalse();
892             }
893         } catch (Throwable e /* catch assertions & exceptions */) {
894             try {
895                 mConnectivityManager.unregisterNetworkCallback(testNetworkCallback);
896             } catch (IllegalArgumentException ie) { }
897             throw e;
898         }
899         mWifiManager.removeLocalOnlyConnectionFailureListener(localOnlyListener);
900         return testNetworkCallback;
901     }
902 
903     /**
904      * Tests the entire connection flow using the provided specifier.
905      *
906      * Note: The helper method drops the shell identity, so don't use this if the caller already
907      * adopted shell identity.
908      *
909      * @param specifier Specifier to use for network request.
910      * @param shouldUserReject Whether to simulate user rejection or not.
911      *
912      * @return NetworkCallback used for the connection (can be used by client to release the
913      * connection.
914      */
testConnectionFlowWithSpecifier( WifiConfiguration network, WifiNetworkSpecifier specifier, boolean shouldUserReject)915     public ConnectivityManager.NetworkCallback testConnectionFlowWithSpecifier(
916             WifiConfiguration network, WifiNetworkSpecifier specifier, boolean shouldUserReject)
917             throws Exception {
918         final UiAutomation uiAutomation =
919                 InstrumentationRegistry.getInstrumentation().getUiAutomation();
920         try {
921             uiAutomation.adoptShellPermissionIdentity(NETWORK_SETTINGS);
922             return testConnectionFlowWithSpecifierWithShellIdentity(
923                     network, specifier, shouldUserReject);
924         } finally {
925             uiAutomation.dropShellPermissionIdentity();
926         }
927     }
928 
929     /**
930      * Returns the number of wifi connections visible at the networking layer.
931      */
getNumWifiConnections()932     public long getNumWifiConnections() {
933         Network[] networks = mConnectivityManager.getAllNetworks();
934         return Arrays.stream(networks)
935                 .filter(n -> mConnectivityManager.getNetworkCapabilities(n) != null
936                         && mConnectivityManager.getNetworkCapabilities(n)
937                                 .hasTransport(TRANSPORT_WIFI))
938                 .count();
939     }
940 
941     /**
942      * Registers a network callback for internet connectivity via wifi and asserts that a network
943      * is available within {@link #DURATION_NETWORK_CONNECTION_MILLIS}.
944      *
945      * @throws Exception
946      */
assertWifiInternetConnectionAvailable()947     public void assertWifiInternetConnectionAvailable() throws Exception {
948         TestNetworkCallback testNetworkCallback = createTestNetworkCallback();
949         try {
950             // File a callback for wifi network.
951             NetworkRequest.Builder builder = new NetworkRequest.Builder()
952                     .addTransportType(NetworkCapabilities.TRANSPORT_WIFI)
953                     .addCapability(NET_CAPABILITY_INTERNET);
954             if (ApiLevelUtil.isAtLeast(Build.VERSION_CODES.S)) {
955                 // Needed to ensure that the restricted concurrent connection does not
956                 // match this request.
957                 builder.addForbiddenCapability(NET_CAPABILITY_OEM_PAID)
958                         .addForbiddenCapability(NET_CAPABILITY_OEM_PRIVATE);
959             }
960             mConnectivityManager.registerNetworkCallback(builder.build(), testNetworkCallback);
961             // Wait for connection to complete & ensure we are connected to some network capable
962             // of providing internet access.
963             assertThat(testNetworkCallback.waitForAnyCallback(DURATION_NETWORK_CONNECTION_MILLIS))
964                     .isTrue();
965             assertThat(testNetworkCallback.onAvailableCalled).isTrue();
966         } finally {
967             mConnectivityManager.unregisterNetworkCallback(testNetworkCallback);
968         }
969     }
970 
getBandFromFrequency(final int freqMHz)971     public static int getBandFromFrequency(final int freqMHz) {
972         if (freqMHz < 1000) {
973             return ScanResult.UNSPECIFIED;
974         } else if (freqMHz < 4000) { // getFrequency is in WifiInfo.FREQUENCY_UNITS = MHz
975             return ScanResult.WIFI_BAND_24_GHZ;
976         } else if (freqMHz < 5900) {
977             // 5GHz band stops at 5885MHz, 6GHz band starts at 5955. See android.net.wifi.ScanResult
978             return ScanResult.WIFI_BAND_5_GHZ;
979         } else if (freqMHz < 10_000) {
980             return ScanResult.WIFI_BAND_6_GHZ;
981         } else if (freqMHz < 71_000) {
982             // 60 GHz band stops at 70_200
983             return ScanResult.WIFI_BAND_60_GHZ;
984         } else {
985             return ScanResult.UNSPECIFIED;
986         }
987     }
988 
989     /**
990      * Create a network request for specified band in a network specifier.
991      */
createNetworkRequestForInternet(int band)992     private NetworkRequest createNetworkRequestForInternet(int band) {
993         final NetworkRequest networkRequest = new NetworkRequest.Builder()
994                 .clearCapabilities()
995                 .addCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET)
996                 .addTransportType(TRANSPORT_WIFI)
997                 .setNetworkSpecifier(new WifiNetworkSpecifier.Builder()
998                         .setBand(band).build())
999                 .build();
1000         return networkRequest;
1001     }
1002 
1003     /**
1004      * Check if a wifi network info is as expected for multi internet connections.
1005      * @return the WifiInfo of the network.
1006      */
checkWifiNetworkInfo(TestNetworkCallback testNetworkCallback, int band)1007     private WifiInfo checkWifiNetworkInfo(TestNetworkCallback testNetworkCallback,
1008             int band) {
1009         if (testNetworkCallback.networkCapabilities == null) {
1010             return null;
1011         }
1012         WifiInfo wifiInfo = getWifiInfo(testNetworkCallback.networkCapabilities);
1013         assertTrue(wifiInfo.isTrusted());
1014         assertFalse(wifiInfo.isRestricted());
1015         WifiInfo redact = wifiInfo
1016                 .makeCopy(NetworkCapabilities.REDACT_FOR_ACCESS_FINE_LOCATION);
1017         assertNotNull(wifiInfo.getInformationElements());
1018         assertNull(redact.getInformationElements());
1019         assertEquals(NetworkCapabilities.REDACT_FOR_ACCESS_FINE_LOCATION
1020                         | NetworkCapabilities.REDACT_FOR_LOCAL_MAC_ADDRESS
1021                         | NetworkCapabilities.REDACT_FOR_NETWORK_SETTINGS,
1022                         redact.getApplicableRedactions());
1023         assertEquals(band, getBandFromFrequency(wifiInfo.getFrequency()));
1024 
1025         return wifiInfo;
1026     }
1027 
1028     /**
1029      * Tests the entire connection success/failure flow using the provided suggestion.
1030      *
1031      * @param executorService Excutor service to run scan periodically (to trigger connection).
1032      * @param expectConnectionSuccess Whether to expect connection success or not.
1033      */
testMultiInternetConnectionFlowInternal( @onNull ScheduledExecutorService executorService, boolean expectConnectionSuccess)1034     private void testMultiInternetConnectionFlowInternal(
1035             @NonNull ScheduledExecutorService executorService,
1036             boolean expectConnectionSuccess) throws Exception {
1037         // File the network request & wait for the callback.
1038         TestNetworkCallback testNetworkCallback2G = createTestNetworkCallback();
1039         TestNetworkCallback testNetworkCallback5G = createTestNetworkCallback();
1040         TestNetworkCallback testNetworkCallback5GLow = createTestNetworkCallback();
1041         TestNetworkCallback testNetworkCallback5GHigh = createTestNetworkCallback();
1042         final NetworkRequest networkRequest2G = createNetworkRequestForInternet(
1043                 ScanResult.WIFI_BAND_24_GHZ);
1044         final NetworkRequest networkRequest5G = createNetworkRequestForInternet(
1045                 ScanResult.WIFI_BAND_5_GHZ);
1046         final NetworkRequest networkRequest5GLow = createNetworkRequestForInternet(
1047                 1 << 29 /*WIFI_BAND_INDEX_5_GHZ_LOW*/);
1048         final NetworkRequest networkRequest5GHigh = createNetworkRequestForInternet(
1049                 1 << 30 /*WIFI_BAND_INDEX_5_GHZ_HIGH*/);
1050          // Make sure wifi is connected to primary after wifi enabled with saved network.
1051         PollingCheck.check("Wifi not connected", DURATION_NETWORK_CONNECTION_MILLIS,
1052                 () -> getNumWifiConnections() > 0);
1053         try {
1054             // Request both 2G and 5G wifi networks.
1055             mConnectivityManager.requestNetwork(networkRequest2G, testNetworkCallback2G);
1056             mConnectivityManager.requestNetwork(networkRequest5G, testNetworkCallback5G);
1057             mConnectivityManager.requestNetwork(networkRequest5GLow, testNetworkCallback5GLow);
1058             mConnectivityManager.requestNetwork(networkRequest5GHigh, testNetworkCallback5GHigh);
1059             // Wait for the request to reach the wifi stack before kick-start periodic scans.
1060             Thread.sleep(200);
1061             boolean band2gFound = false;
1062             boolean band5gFound = false;
1063             boolean band5gLowFound = false;
1064             boolean band5gHighFound = false;
1065             // now wait for connection to complete and wait for callback
1066             WifiInfo primaryInfo = null;
1067             WifiInfo secondaryInfo = null;
1068             if (testNetworkCallback2G.await(DURATION_NETWORK_CONNECTION_MILLIS)) {
1069                 WifiInfo info2g = checkWifiNetworkInfo(testNetworkCallback2G,
1070                         ScanResult.WIFI_BAND_24_GHZ);
1071                 if (info2g != null) {
1072                     if (info2g.isPrimary()) {
1073                         primaryInfo = info2g;
1074                     } else {
1075                         secondaryInfo = info2g;
1076                     }
1077                     band2gFound = true;
1078                 }
1079             }
1080             if (testNetworkCallback5G.await(DURATION_NETWORK_CONNECTION_MILLIS)) {
1081                 WifiInfo info5g = checkWifiNetworkInfo(testNetworkCallback5G,
1082                         ScanResult.WIFI_BAND_5_GHZ);
1083                 if (info5g != null) {
1084                     if (info5g.isPrimary()) {
1085                         primaryInfo = info5g;
1086                     } else {
1087                         secondaryInfo = info5g;
1088                     }
1089                     band5gFound = true;
1090                 }
1091             }
1092             if (testNetworkCallback5GLow.await(DURATION_NETWORK_CONNECTION_MILLIS)) {
1093                 WifiInfo info5g = checkWifiNetworkInfo(testNetworkCallback5GLow,
1094                         ScanResult.WIFI_BAND_5_GHZ);
1095                 if (info5g != null) {
1096                     if (info5g.isPrimary()) {
1097                         primaryInfo = info5g;
1098                     } else {
1099                         secondaryInfo = info5g;
1100                     }
1101                     band5gLowFound = true;
1102                 }
1103             }
1104             if (testNetworkCallback5GHigh.await(DURATION_NETWORK_CONNECTION_MILLIS)) {
1105                 WifiInfo info5g = checkWifiNetworkInfo(testNetworkCallback5GHigh,
1106                         ScanResult.WIFI_BAND_5_GHZ);
1107                 if (info5g != null) {
1108                     if (info5g.isPrimary()) {
1109                         primaryInfo = info5g;
1110                     } else {
1111                         secondaryInfo = info5g;
1112                     }
1113                     band5gHighFound = true;
1114                 }
1115             }
1116             if (expectConnectionSuccess) {
1117                 // Ensure both primary and non-primary networks are created.
1118                 if (band2gFound) {
1119                     // Expect a 2.4Ghz connection and a 5Ghz connection
1120                     assertTrue("Network not found on 2g", band2gFound);
1121                     assertTrue("Network not found on 5g", band5gFound);
1122                     assertFalse("Network unavailable on 2g",
1123                             testNetworkCallback2G.onUnavailableCalled);
1124                     assertFalse("Network unavailable on 5g",
1125                             testNetworkCallback5G.onUnavailableCalled);
1126                 } else {
1127                     // If there's no 2.4Ghz connection, then expect a 5Ghz low and a 5Ghz high
1128                     // connection
1129                     assertTrue("Network not found on 5g", band5gFound);
1130                     assertTrue("Network not found on secondary 5g",
1131                             band5gLowFound || band5gHighFound);
1132                     assertFalse("Network unavailable on 5g",
1133                             testNetworkCallback5G.onUnavailableCalled);
1134                     assertFalse("Network unavailable on secondary 5G",
1135                             testNetworkCallback5GLow.onUnavailableCalled
1136                                     || testNetworkCallback5GHigh.onUnavailableCalled);
1137                 }
1138                 assertNotNull("No primary network info", primaryInfo);
1139                 assertNotNull("No secondary network info", secondaryInfo);
1140                 assertFalse("Primary and secondary networks are same",
1141                         primaryInfo.equals(secondaryInfo));
1142                 // Ensure that there are 2 wifi connections available for apps.
1143                 assertEquals("Expecting 2 Wifi networks", 2, getNumWifiConnections());
1144                 // Check if the networks meets the expected requested multi internet state
1145                 int mode = mWifiManager.getStaConcurrencyForMultiInternetMode();
1146                 if (mode == mWifiManager.WIFI_MULTI_INTERNET_MODE_MULTI_AP) {
1147                     // Multi AP state allows connecting to same network or multi APs in other
1148                     // networks, with different BSSIDs.
1149                     assertFalse("Can not connect to same bssid" + primaryInfo.getBSSID()
1150                             + " / " + secondaryInfo.getBSSID(),
1151                             TextUtils.equals(primaryInfo.getBSSID(), secondaryInfo.getBSSID()));
1152                 } else if (mode == mWifiManager.WIFI_MULTI_INTERNET_MODE_DBS_AP) {
1153                     assertTrue("NETWORK_DBS mode can only connect to the same SSID but got "
1154                             + primaryInfo.getSSID() + " / " + secondaryInfo.getSSID(),
1155                             TextUtils.equals(primaryInfo.getSSID(), secondaryInfo.getSSID()));
1156                     assertEquals("NETWORK_DBS mode can only connect to the same network Id but got"
1157                             + primaryInfo.getNetworkId() + " / " + secondaryInfo.getNetworkId(),
1158                             primaryInfo.getNetworkId(), secondaryInfo.getNetworkId());
1159                     assertEquals("NETWORK_DBS mode can only connect to same security type but got"
1160                             + primaryInfo.getCurrentSecurityType() + " / "
1161                             + secondaryInfo.getCurrentSecurityType(),
1162                             primaryInfo.getCurrentSecurityType(),
1163                             secondaryInfo.getCurrentSecurityType());
1164                 } else {
1165                     fail("Invalid multi internet mode " + mode);
1166                 }
1167             } else {
1168                 // Ensure no band specified wifi connection is created.
1169                 assertTrue(testNetworkCallback2G.onUnavailableCalled
1170                         || testNetworkCallback5G.onUnavailableCalled);
1171                 // Only one wifi network
1172                 assertEquals("There should be only one wifi network but got "
1173                         + getNumWifiConnections(), 1, getNumWifiConnections());
1174             }
1175         } finally {
1176             mConnectivityManager.unregisterNetworkCallback(testNetworkCallback2G);
1177             mConnectivityManager.unregisterNetworkCallback(testNetworkCallback5G);
1178             mConnectivityManager.unregisterNetworkCallback(testNetworkCallback5GLow);
1179             mConnectivityManager.unregisterNetworkCallback(testNetworkCallback5GHigh);
1180             executorService.shutdown();
1181         }
1182     }
1183 
1184     /**
1185      * Tests the entire connection success flow using the provided suggestion.
1186      *
1187      * Note: The caller needs to invoke this after acquiring shell identity.
1188      *
1189      * @param executorService Excutor service to run scan periodically (to trigger connection).
1190      * @param expectConnectionSuccess Whether to expect connection success or not.
1191      */
testMultiInternetConnectionFlowWithShellIdentity( @onNull ScheduledExecutorService executorService, boolean expectConnectionSuccess)1192     public void testMultiInternetConnectionFlowWithShellIdentity(
1193             @NonNull ScheduledExecutorService executorService,
1194             boolean expectConnectionSuccess) throws Exception {
1195         final UiAutomation uiAutomation =
1196                 InstrumentationRegistry.getInstrumentation().getUiAutomation();
1197         try {
1198             uiAutomation.adoptShellPermissionIdentity(NETWORK_SETTINGS, CONNECTIVITY_INTERNAL);
1199             testMultiInternetConnectionFlowInternal(
1200                     executorService, expectConnectionSuccess);
1201         } finally {
1202             uiAutomation.dropShellPermissionIdentity();
1203         }
1204     }
1205 }
1206