Clear last connected wifi network when starting an invocation.

Bug: 17488684
Change-Id: I1efd5e8e153664dac9bd8b033dea20f77dfb722e
diff --git a/src/com/android/tradefed/device/ITestDevice.java b/src/com/android/tradefed/device/ITestDevice.java
index d659895..8c2810f 100644
--- a/src/com/android/tradefed/device/ITestDevice.java
+++ b/src/com/android/tradefed/device/ITestDevice.java
@@ -704,11 +704,18 @@
     public InputStreamSource getScreenshot() throws DeviceNotAvailableException;
 
     /**
+     * Clears the last connected wifi network. This should be called when starting a new invocation
+     * to avoid connecting to the wifi network used in the previous test after device reboots.
+     */
+    public void clearLastConnectedWifiNetwork();
+
+    /**
      * Connects to a wifi network.
      * <p/>
      * Turns on wifi and blocks until a successful connection is made to the specified wifi network.
      * Once a connection is made, the instance will try to restore the connection after every reboot
-     * until {@link ITestDevice#disconnectFromWifi()} is called.
+     * until {@link ITestDevice#disconnectFromWifi()} or
+     * {@link ITestDevice#clearLastConnectedWifiNetwork()} is called.
      *
      * @param wifiSsid the wifi ssid to connect to
      * @param wifiPsk PSK passphrase or null if unencrypted
diff --git a/src/com/android/tradefed/device/TestDevice.java b/src/com/android/tradefed/device/TestDevice.java
index e075343..0df517c 100644
--- a/src/com/android/tradefed/device/TestDevice.java
+++ b/src/com/android/tradefed/device/TestDevice.java
@@ -148,8 +148,8 @@
     private DeviceAllocationState mAllocationState = DeviceAllocationState.Unknown;
     private IDeviceMonitor mAllocationMonitor = null;
 
-    private String mWifiSsid = null;
-    private String mWifiPsk = null;
+    private String mLastConnectedWifiSsid = null;
+    private String mLastConnectedWifiPsk = null;
     private boolean mNetworkMonitorEnabled = false;
 
     /**
@@ -1697,12 +1697,21 @@
      * {@inheritDoc}
      */
     @Override
+    public void clearLastConnectedWifiNetwork() {
+        mLastConnectedWifiSsid = null;
+        mLastConnectedWifiPsk = null;
+    }
+
+    /**
+     * {@inheritDoc}
+     */
+    @Override
     public boolean connectToWifiNetwork(String wifiSsid, String wifiPsk)
             throws DeviceNotAvailableException {
 
         // Clears the last connected wifi network.
-        mWifiSsid = null;
-        mWifiPsk = null;
+        mLastConnectedWifiSsid = null;
+        mLastConnectedWifiPsk = null;
 
         IWifiHelper wifi = createWifiHelper();
         for (int i = 1; i <= mOptions.getWifiAttempts(); i++) {
@@ -1714,8 +1723,8 @@
                 CLog.i("Successfully connected to wifi network %s(%s) on %s",
                         wifiSsid, wifiInfo.get("bssid"), getSerialNumber());
 
-                mWifiSsid = wifiSsid;
-                mWifiPsk = wifiPsk;
+                mLastConnectedWifiSsid = wifiSsid;
+                mLastConnectedWifiPsk = wifiPsk;
 
                 return true;
             } else {
@@ -1803,8 +1812,8 @@
     public boolean disconnectFromWifi() throws DeviceNotAvailableException {
         CLog.i("Disconnecting from wifi on %s", getSerialNumber());
         // Clears the last connected wifi network.
-        mWifiSsid = null;
-        mWifiPsk = null;
+        mLastConnectedWifiSsid = null;
+        mLastConnectedWifiPsk = null;
 
         IWifiHelper wifi = createWifiHelper();
         return wifi.disconnectFromNetwork();
@@ -1944,10 +1953,10 @@
         if (mOptions.isDisableKeyguard()) {
             disableKeyguard();
         }
-        if (mWifiSsid != null) {
+        if (mLastConnectedWifiSsid != null) {
             // mWifiSsid is set to null if connection fails
-            final String wifiSsid = mWifiSsid;
-            if (!connectToWifiNetworkIfNeeded(mWifiSsid, mWifiPsk)) {
+            final String wifiSsid = mLastConnectedWifiSsid;
+            if (!connectToWifiNetworkIfNeeded(mLastConnectedWifiSsid, mLastConnectedWifiPsk)) {
                 throw new NetworkNotAvailableException(
                         String.format("Failed to connect to wifi network %s on %s after reboot",
                                 wifiSsid, getSerialNumber()));
diff --git a/src/com/android/tradefed/invoker/TestInvocation.java b/src/com/android/tradefed/invoker/TestInvocation.java
index 02d9740..d0f4f10 100644
--- a/src/com/android/tradefed/invoker/TestInvocation.java
+++ b/src/com/android/tradefed/invoker/TestInvocation.java
@@ -201,6 +201,7 @@
             mStatus = "fetching build";
             config.getLogOutput().init();
             getLogRegistry().registerLogger(config.getLogOutput());
+            device.clearLastConnectedWifiNetwork();
             device.setOptions(config.getDeviceOptions());
             if (config.getDeviceOptions().isLogcatCaptureEnabled()) {
                 device.startLogcat();
diff --git a/tests/src/com/android/tradefed/device/StubTestDevice.java b/tests/src/com/android/tradefed/device/StubTestDevice.java
index 28838d7..678ec49 100644
--- a/tests/src/com/android/tradefed/device/StubTestDevice.java
+++ b/tests/src/com/android/tradefed/device/StubTestDevice.java
@@ -21,11 +21,7 @@
 import com.android.ddmlib.testrunner.IRemoteAndroidTestRunner;
 import com.android.ddmlib.testrunner.ITestRunListener;
 import com.android.tradefed.build.IBuildInfo;
-import com.android.tradefed.device.DeviceNotAvailableException;
-import com.android.tradefed.device.IWifiHelper;
-import com.android.tradefed.log.LogUtil.CLog;
 import com.android.tradefed.result.InputStreamSource;
-import com.android.tradefed.targetprep.TargetSetupError;
 import com.android.tradefed.util.CommandResult;
 
 import java.io.File;
@@ -412,6 +408,11 @@
     }
 
     @Override
+    public void clearLastConnectedWifiNetwork() {
+        //ignore
+    }
+
+    @Override
     public boolean connectToWifiNetwork(String wifiSsid, String wifiPsk)
             throws DeviceNotAvailableException {
         // ignore
@@ -464,6 +465,7 @@
     /**
      * {@inheritDoc}
      */
+    @Override
     public boolean checkConnectivity() throws DeviceNotAvailableException {
         return false;
     }
diff --git a/tests/src/com/android/tradefed/invoker/TestInvocationTest.java b/tests/src/com/android/tradefed/invoker/TestInvocationTest.java
index ecfe742..4ffeba1 100644
--- a/tests/src/com/android/tradefed/invoker/TestInvocationTest.java
+++ b/tests/src/com/android/tradefed/invoker/TestInvocationTest.java
@@ -365,6 +365,7 @@
 
         EasyMock.expect(mMockBuildProvider.getBuild()).andReturn(mMockBuildInfo);
         resumeListener.invocationStarted(mMockBuildInfo);
+        mMockDevice.clearLastConnectedWifiNetwork();
         mMockDevice.setOptions((TestDeviceOptions)EasyMock.anyObject());
         mMockBuildInfo.setDeviceSerial(SERIAL);
         mMockDevice.startLogcat();
@@ -402,6 +403,7 @@
         mMockLogger.init();
         mMockLogSaver.invocationStarted(mMockBuildInfo);
         // now set resumed invocation expectations
+        mMockDevice.clearLastConnectedWifiNetwork();
         mMockDevice.setOptions((TestDeviceOptions)EasyMock.anyObject());
         mMockBuildInfo.setDeviceSerial(SERIAL);
         mMockDevice.startLogcat();
@@ -601,6 +603,7 @@
      * Set up expected calls that occur on every invoke, regardless of result
      */
     private void setupInvoke() {
+        mMockDevice.clearLastConnectedWifiNetwork();
         mMockDevice.setOptions((TestDeviceOptions)EasyMock.anyObject());
         mMockDevice.startLogcat();
         mMockDevice.stopLogcat();