1 /* 2 * Copyright (C) 2020 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 com.android.frameworks.core.batterystatsloadtests; 18 19 import static org.junit.Assert.assertEquals; 20 21 import android.app.Instrumentation; 22 import android.content.Context; 23 import android.net.ConnectivityManager; 24 import android.net.Network; 25 import android.net.NetworkCapabilities; 26 import android.net.NetworkRequest; 27 import android.net.wifi.WifiManager; 28 import android.os.Handler; 29 import android.os.Looper; 30 import android.os.Message; 31 32 import androidx.test.platform.app.InstrumentationRegistry; 33 34 import com.android.compatibility.common.util.SystemUtil; 35 36 import org.junit.rules.TestRule; 37 import org.junit.runner.Description; 38 import org.junit.runners.model.Statement; 39 40 import java.util.concurrent.CountDownLatch; 41 import java.util.concurrent.TimeUnit; 42 43 public class ConnectivitySetupRule implements TestRule { 44 45 private final boolean mWifiEnabled; 46 private final ConnectivityManager mConnectivityManager; 47 private final WifiManager mWifiManager; 48 private boolean mInitialWifiState; 49 ConnectivitySetupRule(boolean wifiEnabled)50 public ConnectivitySetupRule(boolean wifiEnabled) { 51 mWifiEnabled = wifiEnabled; 52 53 Instrumentation instrumentation = InstrumentationRegistry.getInstrumentation(); 54 Context context = instrumentation.getContext(); 55 mConnectivityManager = context.getSystemService(ConnectivityManager.class); 56 mWifiManager = context.getSystemService(WifiManager.class); 57 } 58 59 @Override apply(Statement base, Description description)60 public Statement apply(Statement base, Description description) { 61 return new Statement() { 62 @Override 63 public void evaluate() throws Throwable { 64 try { 65 mInitialWifiState = isWiFiConnected(); 66 setWiFiState(mWifiEnabled); 67 base.evaluate(); 68 } finally { 69 setWiFiState(mInitialWifiState); 70 } 71 } 72 }; 73 } 74 75 private void setWiFiState(final boolean enable) throws InterruptedException { 76 boolean wiFiConnected = isWiFiConnected(); 77 if (enable == wiFiConnected) { 78 return; 79 } 80 81 NetworkTracker tracker = new NetworkTracker(!mWifiEnabled); 82 mConnectivityManager.registerNetworkCallback( 83 new NetworkRequest.Builder() 84 .addCapability(NetworkCapabilities.NET_CAPABILITY_NOT_METERED).build(), 85 tracker); 86 87 if (enable) { 88 SystemUtil.runShellCommand("svc wifi enable"); 89 //noinspection deprecation 90 SystemUtil.runWithShellPermissionIdentity(mWifiManager::reconnect, 91 android.Manifest.permission.NETWORK_SETTINGS); 92 } else { 93 SystemUtil.runShellCommand("svc wifi disable"); 94 } 95 96 tracker.waitForExpectedState(); 97 98 assertEquals("Wifi must be " + (enable ? "connected to" : "disconnected from") 99 + " an access point for this test.", enable, isWiFiConnected()); 100 101 mConnectivityManager.unregisterNetworkCallback(tracker); 102 } 103 104 private boolean isWiFiConnected() { 105 return mWifiManager.isWifiEnabled() && mConnectivityManager.getActiveNetwork() != null 106 && !mConnectivityManager.isActiveNetworkMetered(); 107 } 108 109 private class NetworkTracker extends ConnectivityManager.NetworkCallback { 110 private static final int MSG_CHECK_ACTIVE_NETWORK = 1; 111 112 private final CountDownLatch mReceiveLatch = new CountDownLatch(1); 113 114 private final boolean mExpectedMetered; 115 116 private final Handler mHandler = new Handler(Looper.getMainLooper()) { 117 @Override 118 public void handleMessage(Message msg) { 119 if (msg.what == MSG_CHECK_ACTIVE_NETWORK) { 120 checkActiveNetwork(); 121 } 122 } 123 }; 124 125 private NetworkTracker(boolean expectedMetered) { 126 mExpectedMetered = expectedMetered; 127 } 128 129 @Override 130 public void onAvailable(Network network) { 131 checkActiveNetwork(); 132 } 133 134 @Override 135 public void onLost(Network network) { 136 checkActiveNetwork(); 137 } 138 139 boolean waitForExpectedState() throws InterruptedException { 140 checkActiveNetwork(); 141 return mReceiveLatch.await(60, TimeUnit.SECONDS); 142 } 143 144 private void checkActiveNetwork() { 145 if (mReceiveLatch.getCount() == 0) { 146 return; 147 } 148 149 if (mConnectivityManager.getActiveNetwork() != null 150 && mConnectivityManager.isActiveNetworkMetered() == mExpectedMetered) { 151 mReceiveLatch.countDown(); 152 } else { 153 mHandler.sendEmptyMessageDelayed(MSG_CHECK_ACTIVE_NETWORK, 5000); 154 } 155 } 156 } 157 } 158