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