Create a recoverer that uses allocated device to reset usb

Test: unit tests
Bug: 154662049
Change-Id: I20307cb48baa85b1587e39d304bbcceccdf8972b
diff --git a/res/config/usb-reset.xml b/res/config/usb-reset.xml
new file mode 100644
index 0000000..93c3511
--- /dev/null
+++ b/res/config/usb-reset.xml
@@ -0,0 +1,22 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2020 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+<configuration description="A fake test that resets the device USB and reboot it.">
+
+    <option name="device" value="true" />
+    <option name="test-tag" value="UsbResetTest" />
+    <test class="com.android.tradefed.testtype.UsbResetTest" />
+    <logger class="com.android.tradefed.log.FileLogger" />
+</configuration>
diff --git a/src/com/android/tradefed/device/recovery/UsbResetMultiDeviceRecovery.java b/src/com/android/tradefed/device/recovery/UsbResetMultiDeviceRecovery.java
index 3aefab7..6877d24 100644
--- a/src/com/android/tradefed/device/recovery/UsbResetMultiDeviceRecovery.java
+++ b/src/com/android/tradefed/device/recovery/UsbResetMultiDeviceRecovery.java
@@ -46,6 +46,11 @@
     @Option(name = "disable", description = "Disable the device recoverer.")
     private boolean mDisable = false;
 
+    @Option(
+            name = "only-reset-unmanaged",
+            description = "Only reset the device that are not currently managed by Tradefed.")
+    private boolean mOnlyResetUnmanaged = false;
+
     private String mFastbootPath = "fastboot";
 
     @Override
@@ -84,6 +89,9 @@
             // Perform a USB port reset on the remaining devices
             for (String serial : deviceSerials) {
                 try (UsbDevice device = usb.getDevice(serial)) {
+                    if (mOnlyResetUnmanaged && managedDeviceMap.containsKey(serial)) {
+                        continue;
+                    }
                     if (device == null) {
                         CLog.w("Device '%s' not found during USB reset.", serial);
                         continue;
diff --git a/src/com/android/tradefed/device/recovery/UsbResetRunConfigRecovery.java b/src/com/android/tradefed/device/recovery/UsbResetRunConfigRecovery.java
new file mode 100644
index 0000000..28999d5
--- /dev/null
+++ b/src/com/android/tradefed/device/recovery/UsbResetRunConfigRecovery.java
@@ -0,0 +1,30 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.tradefed.device.recovery;
+
+import com.android.tradefed.config.OptionClass;
+import com.android.tradefed.device.IManagedTestDevice;
+import com.android.tradefed.device.TestDeviceState;
+
+/** Allow to trigger a command to reset the USB of a device */
+@OptionClass(alias = "usb-reset-recovery")
+public class UsbResetRunConfigRecovery extends RunConfigDeviceRecovery {
+
+    @Override
+    public boolean shouldSkip(IManagedTestDevice device) {
+        return TestDeviceState.FASTBOOT.equals(device.getDeviceState());
+    }
+}
diff --git a/src/com/android/tradefed/testtype/DeviceBatteryLevelChecker.java b/src/com/android/tradefed/testtype/DeviceBatteryLevelChecker.java
index 14f8e78..7cf18f5 100644
--- a/src/com/android/tradefed/testtype/DeviceBatteryLevelChecker.java
+++ b/src/com/android/tradefed/testtype/DeviceBatteryLevelChecker.java
@@ -25,7 +25,6 @@
 import com.android.tradefed.invoker.TestInformation;
 import com.android.tradefed.log.LogUtil.CLog;
 import com.android.tradefed.result.ITestInvocationListener;
-import com.android.tradefed.targetprep.ITargetPreparer;
 import com.android.tradefed.util.IRunUtil;
 import com.android.tradefed.util.RunUtil;
 import com.android.tradefed.util.TimeUtil;
@@ -33,7 +32,7 @@
 import org.junit.Assert;
 
 /**
- * An {@link ITargetPreparer} that checks for a minimum battery charge, and waits for the battery to
+ * An {@link IRemoteTest} that checks for a minimum battery charge, and waits for the battery to
  * reach a second charging threshold if the minimum charge isn't present.
  */
 @OptionClass(alias = "battery-checker")
diff --git a/src/com/android/tradefed/testtype/UsbResetTest.java b/src/com/android/tradefed/testtype/UsbResetTest.java
new file mode 100644
index 0000000..1be0792
--- /dev/null
+++ b/src/com/android/tradefed/testtype/UsbResetTest.java
@@ -0,0 +1,62 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.tradefed.testtype;
+
+import com.android.helper.aoa.UsbDevice;
+import com.android.helper.aoa.UsbHelper;
+import com.android.tradefed.config.OptionClass;
+import com.android.tradefed.device.DeviceNotAvailableException;
+import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.invoker.TestInformation;
+import com.android.tradefed.log.LogUtil.CLog;
+import com.android.tradefed.result.ITestInvocationListener;
+
+import com.google.common.annotations.VisibleForTesting;
+
+/**
+ * An {@link IRemoteTest} that reset the device USB and checks whether the device comes back online
+ * afterwards.
+ */
+@OptionClass(alias = "usb-reset-test")
+public class UsbResetTest implements IRemoteTest {
+
+    @Override
+    public void run(TestInformation testInfo, ITestInvocationListener listener)
+            throws DeviceNotAvailableException {
+        ITestDevice device = testInfo.getDevice();
+        try (UsbHelper usb = getUsbHelper()) {
+            String serial = device.getSerialNumber();
+            try (UsbDevice usbDevice = usb.getDevice(serial)) {
+                if (usbDevice == null) {
+                    throw new DeviceNotAvailableException(
+                            String.format("Device '%s' not found during USB reset.", serial),
+                            serial);
+                } else {
+                    CLog.d("Resetting USB port for device '%s'", serial);
+                    usbDevice.reset();
+                    // If device fails to reboot it will throw an exception and be left unavailable
+                    // again.
+                    device.reboot();
+                }
+            }
+        }
+    }
+
+    @VisibleForTesting
+    UsbHelper getUsbHelper() {
+        return new UsbHelper();
+    }
+}
diff --git a/tests/src/com/android/tradefed/UnitTests.java b/tests/src/com/android/tradefed/UnitTests.java
index ed1ba81..3dabdfc 100644
--- a/tests/src/com/android/tradefed/UnitTests.java
+++ b/tests/src/com/android/tradefed/UnitTests.java
@@ -280,6 +280,7 @@
 import com.android.tradefed.testtype.PythonUnitTestResultParserTest;
 import com.android.tradefed.testtype.PythonUnitTestRunnerTest;
 import com.android.tradefed.testtype.TfTestLauncherTest;
+import com.android.tradefed.testtype.UsbResetTestTest;
 import com.android.tradefed.testtype.binary.ExecutableHostTestTest;
 import com.android.tradefed.testtype.binary.ExecutableTargetTestTest;
 import com.android.tradefed.testtype.host.CoverageMeasurementForwarderTest;
@@ -750,6 +751,7 @@
     PythonUnitTestResultParserTest.class,
     PythonUnitTestRunnerTest.class,
     TfTestLauncherTest.class,
+    UsbResetTestTest.class,
 
     // testtype/binary
     ExecutableHostTestTest.class,
diff --git a/tests/src/com/android/tradefed/testtype/UsbResetTestTest.java b/tests/src/com/android/tradefed/testtype/UsbResetTestTest.java
new file mode 100644
index 0000000..6d166e3
--- /dev/null
+++ b/tests/src/com/android/tradefed/testtype/UsbResetTestTest.java
@@ -0,0 +1,87 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.tradefed.testtype;
+
+import static org.junit.Assert.fail;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+
+import com.android.helper.aoa.UsbDevice;
+import com.android.helper.aoa.UsbHelper;
+import com.android.tradefed.config.ConfigurationDef;
+import com.android.tradefed.device.DeviceNotAvailableException;
+import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.invoker.IInvocationContext;
+import com.android.tradefed.invoker.InvocationContext;
+import com.android.tradefed.invoker.TestInformation;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
+
+/** Unit tests for {@link UsbResetTest}. */
+@RunWith(JUnit4.class)
+public class UsbResetTestTest {
+
+    private UsbResetTest mTest;
+    private TestInformation mTestInfo;
+    private ITestDevice mDevice;
+    private UsbHelper mUsb;
+
+    @Before
+    public void setUp() {
+        mUsb = Mockito.mock(UsbHelper.class);
+        mDevice = Mockito.mock(ITestDevice.class);
+        IInvocationContext context = new InvocationContext();
+        context.addAllocatedDevice(ConfigurationDef.DEFAULT_DEVICE_NAME, mDevice);
+        mTestInfo = TestInformation.newBuilder().setInvocationContext(context).build();
+        mTest =
+                new UsbResetTest() {
+                    @Override
+                    UsbHelper getUsbHelper() {
+                        return mUsb;
+                    }
+                };
+    }
+
+    @Test
+    public void testReset() throws DeviceNotAvailableException {
+        UsbDevice usbDevice = Mockito.mock(UsbDevice.class);
+        doReturn("serial").when(mDevice).getSerialNumber();
+        doReturn(usbDevice).when(mUsb).getDevice("serial");
+        mTest.run(mTestInfo, null);
+
+        verify(usbDevice).reset();
+        verify(mDevice).reboot();
+    }
+
+    @Test
+    public void testReset_noDevice() throws DeviceNotAvailableException {
+        doReturn("serial").when(mDevice).getSerialNumber();
+        doReturn(null).when(mUsb).getDevice("serial");
+        try {
+            mTest.run(mTestInfo, null);
+            fail("Should have thrown an exception");
+        } catch (DeviceNotAvailableException expected) {
+            // expected
+        }
+
+        verify(mDevice, never()).reboot();
+    }
+}