Merge "[TF] add ExecutableTargetTest to run binary test in device."
diff --git a/invocation_interfaces/com/android/tradefed/invoker/ExecutionFiles.java b/invocation_interfaces/com/android/tradefed/invoker/ExecutionFiles.java
index 55e0a33..43b4d16 100644
--- a/invocation_interfaces/com/android/tradefed/invoker/ExecutionFiles.java
+++ b/invocation_interfaces/com/android/tradefed/invoker/ExecutionFiles.java
@@ -121,6 +121,18 @@
         return this;
     }
 
+    /**
+     * Copies all of the mappings from the specified map to this map.
+     *
+     * @param copyFrom original {@link ExecutionFiles} to copy from.
+     * @return The final mapping
+     */
+    public ExecutionFiles putAll(ExecutionFiles copyFrom) {
+        mFiles.putAll(copyFrom.getAll());
+        mShouldNotDelete.addAll(copyFrom.mShouldNotDelete);
+        return this;
+    }
+
     /** Returns whether or not the map of properties is empty. */
     public boolean isEmpty() {
         return mFiles.isEmpty();
diff --git a/invocation_interfaces/com/android/tradefed/invoker/TestInformation.java b/invocation_interfaces/com/android/tradefed/invoker/TestInformation.java
index 5bebbea..a2b1d7f 100644
--- a/invocation_interfaces/com/android/tradefed/invoker/TestInformation.java
+++ b/invocation_interfaces/com/android/tradefed/invoker/TestInformation.java
@@ -60,7 +60,7 @@
         mDependenciesFolder = invocationInfo.mDependenciesFolder;
         if (copyExecFile) {
             mExecutionFiles = new ExecutionFiles();
-            mExecutionFiles.putAll(invocationInfo.executionFiles().getAll());
+            mExecutionFiles.putAll(invocationInfo.executionFiles());
         } else {
             mExecutionFiles = invocationInfo.mExecutionFiles;
         }
diff --git a/src/com/android/tradefed/command/CommandScheduler.java b/src/com/android/tradefed/command/CommandScheduler.java
index 35afc13..64636a1 100644
--- a/src/com/android/tradefed/command/CommandScheduler.java
+++ b/src/com/android/tradefed/command/CommandScheduler.java
@@ -808,13 +808,13 @@
         if (e instanceof DeviceUnresponsiveException) {
             ITestDevice badDevice =
                     context.getDeviceBySerial(((DeviceUnresponsiveException) e).getSerial());
-            if (badDevice != null) {
+            if (badDevice != null && !(badDevice.getIDevice() instanceof StubDevice)) {
                 deviceStates.put(badDevice, FreeDeviceState.UNRESPONSIVE);
             }
         } else if (e instanceof DeviceNotAvailableException) {
             ITestDevice badDevice =
                     context.getDeviceBySerial(((DeviceNotAvailableException) e).getSerial());
-            if (badDevice != null) {
+            if (badDevice != null && !(badDevice.getIDevice() instanceof StubDevice)) {
                 deviceStates.put(badDevice, FreeDeviceState.UNAVAILABLE);
             }
         }
diff --git a/src/com/android/tradefed/invoker/TestInvocation.java b/src/com/android/tradefed/invoker/TestInvocation.java
index 2bf69ee..fdd0a4b 100644
--- a/src/com/android/tradefed/invoker/TestInvocation.java
+++ b/src/com/android/tradefed/invoker/TestInvocation.java
@@ -343,23 +343,8 @@
             // Track the timestamp when we are done with devices
             addInvocationMetric(
                     InvocationMetricKey.DEVICE_DONE_TIMESTAMP, System.currentTimeMillis());
-            // Capture the FreeDeviceState of the primary device
             Map<ITestDevice, FreeDeviceState> devicesStates =
-                    CommandScheduler.createReleaseMap(context, exception);
-            if (devicesStates.size() >= 1) {
-                addInvocationMetric(
-                        InvocationMetricKey.DEVICE_RELEASE_STATE,
-                        devicesStates.values().iterator().next().toString());
-            }
-            int countLost = 0;
-            for (FreeDeviceState fds : devicesStates.values()) {
-                if (FreeDeviceState.UNAVAILABLE.equals(fds)) {
-                    countLost++;
-                }
-            }
-            if (countLost > 0) {
-                addInvocationMetric(InvocationMetricKey.DEVICE_LOST_DETECTED, countLost);
-            }
+                    handleAndLogReleaseState(context, exception);
             if (config.getCommandOptions().earlyDeviceRelease()) {
                 for (IScheduledInvocationListener scheduleListener : mSchedulerListeners) {
                     scheduleListener.releaseDevices(context, devicesStates);
@@ -1052,6 +1037,32 @@
         testInfo.executionFiles().clearFiles();
     }
 
+    private Map<ITestDevice, FreeDeviceState> handleAndLogReleaseState(
+            IInvocationContext context, Throwable exception) {
+        // Capture the FreeDeviceState of the primary device
+        Map<ITestDevice, FreeDeviceState> devicesStates =
+                CommandScheduler.createReleaseMap(context, exception);
+        if (devicesStates.size() >= 1) {
+            addInvocationMetric(
+                    InvocationMetricKey.DEVICE_RELEASE_STATE,
+                    devicesStates.values().iterator().next().toString());
+        }
+        // TODO: Add Handling of virtual devices
+        int countPhysicalLost = 0;
+        for (Entry<ITestDevice, FreeDeviceState> fds : devicesStates.entrySet()) {
+            if (fds.getKey().getIDevice() instanceof StubDevice) {
+                continue;
+            }
+            if (FreeDeviceState.UNAVAILABLE.equals(fds.getValue())) {
+                countPhysicalLost++;
+            }
+        }
+        if (countPhysicalLost > 0) {
+            addInvocationMetric(InvocationMetricKey.DEVICE_LOST_DETECTED, countPhysicalLost);
+        }
+        return devicesStates;
+    }
+
     /** Helper Thread that ensures host_log is reported in case of killed JVM */
     private class ReportHostLog extends Thread {
 
diff --git a/test_framework/com/android/tradefed/testtype/HostTest.java b/test_framework/com/android/tradefed/testtype/HostTest.java
index 51e9acf..f9aad4b 100644
--- a/test_framework/com/android/tradefed/testtype/HostTest.java
+++ b/test_framework/com/android/tradefed/testtype/HostTest.java
@@ -39,6 +39,7 @@
 import com.android.tradefed.testtype.host.PrettyTestEventLogger;
 import com.android.tradefed.testtype.junit4.CarryDnaeError;
 import com.android.tradefed.testtype.junit4.JUnit4ResultForwarder;
+import com.android.tradefed.testtype.suite.ModuleDefinition;
 import com.android.tradefed.util.FileUtil;
 import com.android.tradefed.util.JUnit4TestFilter;
 import com.android.tradefed.util.StreamUtil;
@@ -68,6 +69,7 @@
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
+import java.net.MalformedURLException;
 import java.net.URL;
 import java.net.URLClassLoader;
 import java.util.ArrayDeque;
@@ -890,12 +892,45 @@
             if (classNames.contains(className)) {
                 continue;
             }
+            IllegalArgumentException initialError = null;
             try {
                 classes.add(Class.forName(className, true, getClassLoader()));
                 classNames.add(className);
             } catch (ClassNotFoundException e) {
-                throw new IllegalArgumentException(String.format("Could not load Test class %s",
-                        className), e);
+                initialError =
+                        new IllegalArgumentException(
+                                String.format("Could not load Test class %s", className), e);
+            }
+            if (initialError != null) {
+                // Fallback search a jar for the module under tests if any.
+                String moduleName =
+                        mTestInfo
+                                .getContext()
+                                .getAttributes()
+                                .getUniqueMap()
+                                .get(ModuleDefinition.MODULE_NAME);
+                if (moduleName != null) {
+                    try {
+                        File f = getJarFile(moduleName + ".jar", mTestInfo);
+                        URL[] urls = {f.toURI().toURL()};
+                        URLClassLoader cl = URLClassLoader.newInstance(urls);
+                        mJUnit4JarFiles.add(f);
+                        Class<?> cls = cl.loadClass(className);
+                        classes.add(cls);
+                        classNames.add(className);
+                        initialError = null;
+                    } catch (FileNotFoundException
+                            | MalformedURLException
+                            | ClassNotFoundException fallbackSearch) {
+                        CLog.e(
+                                "Fallback search for a jar containing '%s' didn't work."
+                                        + "Consider using --jar option directly instead of using --class",
+                                className);
+                    }
+                }
+            }
+            if (initialError != null) {
+                throw initialError;
             }
         }
         // Inspect for the jar files
diff --git a/tests/src/com/android/tradefed/testtype/JarHostTestTest.java b/tests/src/com/android/tradefed/testtype/JarHostTestTest.java
index f8fbc1a..e7f1bca 100644
--- a/tests/src/com/android/tradefed/testtype/JarHostTestTest.java
+++ b/tests/src/com/android/tradefed/testtype/JarHostTestTest.java
@@ -33,6 +33,7 @@
 import com.android.tradefed.result.ITestInvocationListener;
 import com.android.tradefed.result.TestDescription;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner.TestMetrics;
+import com.android.tradefed.testtype.suite.ModuleDefinition;
 import com.android.tradefed.util.FileUtil;
 import com.android.tradefed.util.proto.TfMetricProtoUtil;
 
@@ -69,18 +70,6 @@
     private File mTestDir = null;
     private ITestInvocationListener mListener;
 
-    /** More testable version of {@link HostTest} */
-    public static class HostTestable extends HostTest {
-
-        public static File mTestDir;
-
-        public HostTestable() {}
-
-        public HostTestable(File testDir) {
-            mTestDir = testDir;
-        }
-    }
-
     @Before
     public void setUp() throws Exception {
         mTest = new HostTest();
@@ -145,7 +134,7 @@
     @Test
     public void testSplit_withJar() throws Exception {
         File testJar = getJarResource(TEST_JAR1, mTestDir);
-        mTest = new HostTestLoader(mTestDir, testJar);
+        mTest = new HostTestLoader(testJar);
         mTest.setBuild(mStubBuildInfo);
         ITestDevice device = EasyMock.createNiceMock(ITestDevice.class);
         mTest.setDevice(device);
@@ -225,7 +214,7 @@
     @Test
     public void testSplit_countWithFilter() throws Exception {
         File testJar = getJarResource(TEST_JAR1, mTestDir);
-        mTest = new HostTestLoader(mTestDir, testJar);
+        mTest = new HostTestLoader(testJar);
         mTest.setBuild(mStubBuildInfo);
         ITestDevice device = EasyMock.createNiceMock(ITestDevice.class);
         mTest.setDevice(device);
@@ -242,14 +231,13 @@
     /**
      * Testable version of {@link HostTest} that allows adding jar to classpath for testing purpose.
      */
-    public static class HostTestLoader extends HostTestable {
+    public static class HostTestLoader extends HostTest {
 
         private static File mTestJar;
 
         public HostTestLoader() {}
 
-        public HostTestLoader(File testDir, File jar) {
-            super(testDir);
+        public HostTestLoader(File jar) {
             mTestJar = jar;
         }
 
@@ -326,7 +314,7 @@
     @Test
     public void testRunWithJar() throws Exception {
         File testJar = getJarResource(TEST_JAR2, mTestDir);
-        mTest = new HostTestLoader(mTestDir, testJar);
+        mTest = new HostTest();
         mTest.setBuild(mStubBuildInfo);
         ITestDevice device = EasyMock.createNiceMock(ITestDevice.class);
         mTest.setDevice(device);
@@ -352,4 +340,38 @@
         mTest.run(mTestInfo, mListener);
         EasyMock.verify(mListener);
     }
+
+    @Test
+    public void testRunWithClassFromExternalJar() throws Exception {
+        File testJar = getJarResource(TEST_JAR2, mTestDir);
+        mTestInfo
+                .getContext()
+                .addInvocationAttribute(
+                        ModuleDefinition.MODULE_NAME, FileUtil.getBaseName(testJar.getName()));
+        mTest = new HostTest();
+        mTest.setBuild(mStubBuildInfo);
+        ITestDevice device = EasyMock.createNiceMock(ITestDevice.class);
+        mTest.setDevice(device);
+        OptionSetter setter = new OptionSetter(mTest);
+        setter.setOptionValue("enable-pretty-logs", "false");
+        setter.setOptionValue("class", "com.android.tradefed.JUnit4TfUnitTest");
+        // full class count without sharding
+        mTest.setTestInformation(mTestInfo);
+        assertEquals(2, mTest.countTestCases());
+
+        mListener.testRunStarted("com.android.tradefed.JUnit4TfUnitTest", 2);
+        TestDescription testOne =
+                new TestDescription("com.android.tradefed.JUnit4TfUnitTest", "testOne");
+        TestDescription testTwo =
+                new TestDescription("com.android.tradefed.JUnit4TfUnitTest", "testTwo");
+        mListener.testStarted(testOne);
+        mListener.testEnded(EasyMock.eq(testOne), EasyMock.<HashMap<String, Metric>>anyObject());
+        mListener.testStarted(testTwo);
+        mListener.testEnded(EasyMock.eq(testTwo), EasyMock.<HashMap<String, Metric>>anyObject());
+        mListener.testRunEnded(EasyMock.anyLong(), EasyMock.<HashMap<String, Metric>>anyObject());
+
+        EasyMock.replay(mListener);
+        mTest.run(mTestInfo, mListener);
+        EasyMock.verify(mListener);
+    }
 }
diff --git a/tests/src/com/android/tradefed/testtype/python/PythonBinaryHostTestTest.java b/tests/src/com/android/tradefed/testtype/python/PythonBinaryHostTestTest.java
index 2c32f13..8cf07d0 100644
--- a/tests/src/com/android/tradefed/testtype/python/PythonBinaryHostTestTest.java
+++ b/tests/src/com/android/tradefed/testtype/python/PythonBinaryHostTestTest.java
@@ -579,7 +579,7 @@
                                 eq("--test-output-file"),
                                 capture(testOutputFilePath)))
                 .andStubAnswer(
-                        new IAnswer<>() {
+                        new IAnswer<CommandResult>() {
                             @Override
                             public CommandResult answer() {
                                 try {