Merge "release-request-8a6a1d17-d673-453b-8612-ea961499cd38-for-aosp-oreo-cts-release-4341710 snap-temp-L24800000103383547" into oreo-cts-release
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..50fc1d3
--- /dev/null
+++ b/README.md
@@ -0,0 +1,17 @@
+# Trade Federation (TF / tradefed)
+
+TF is a test harness used to drive Android automated testing. It runs on test hosts
+and monitors the connected devices, handling test scheduling & execution and device
+management.
+
+Other test harnesses like Compatibility Test Suite (CTS) and Vendor Test Suite
+(VTS) use TF as a basis and extend it for their particular needs.
+
+Building TF:
+  * source build/envsetup.sh
+  * tapas tradefed-all
+  * make -j8
+
+More information at:
+https://source.android.com/devices/tech/test_infra/tradefed/
+
diff --git a/error_prone_rules.mk b/error_prone_rules.mk
index 5d5e174..df0ba64 100644
--- a/error_prone_rules.mk
+++ b/error_prone_rules.mk
@@ -21,6 +21,7 @@
                           -Xep:EqualsIncompatibleType:ERROR \
                           -Xep:FormatString:ERROR \
                           -Xep:GetClassOnClass:ERROR \
+                          -Xep:IdentityBinaryExpression:ERROR \
                           -Xep:JUnit3TestNotRun:ERROR \
                           -Xep:JUnitAmbiguousTestClass:ERROR \
                           -Xep:MissingFail:ERROR \
diff --git a/prod-tests/src/com/android/ota/tests/SideloadOtaStabilityTest.java b/prod-tests/src/com/android/ota/tests/SideloadOtaStabilityTest.java
index 13db9ad..a993536 100644
--- a/prod-tests/src/com/android/ota/tests/SideloadOtaStabilityTest.java
+++ b/prod-tests/src/com/android/ota/tests/SideloadOtaStabilityTest.java
@@ -293,15 +293,15 @@
         listener.testStarted(test);
         try {
             mKmsgReceiver.start();
-            CLog.i("Pushing OTA package %s", otaBuild.getOtaPackageFile().getAbsolutePath());
-            Assert.assertTrue(mDevice.pushFile(otaBuild.getOtaPackageFile(), mPackageDataPath));
-            // this file needs to be uncrypted, since /data isn't mounted in recovery
-            // block.map should be empty since cache should be cleared
-            mDevice.pushString(mPackageDataPath + "\n", UNCRYPT_FILE_PATH);
-            // Flushing the file to flash.
-            mDevice.executeShellCommand("sync");
-
             try {
+                CLog.i("Pushing OTA package %s", otaBuild.getOtaPackageFile().getAbsolutePath());
+                Assert.assertTrue(mDevice.pushFile(otaBuild.getOtaPackageFile(), mPackageDataPath));
+                // this file needs to be uncrypted, since /data isn't mounted in recovery
+                // block.map should be empty since cache should be cleared
+                mDevice.pushString(mPackageDataPath + "\n", UNCRYPT_FILE_PATH);
+                // Flushing the file to flash.
+                mDevice.executeShellCommand("sync");
+
                 mUncryptDuration = doUncrypt(SocketFactory.getInstance(), listener);
                 metrics.put("uncrypt_duration", Long.toString(mUncryptDuration));
                 String installOtaCmd = String.format("--update_package=%s\n", BLOCK_MAP_PATH);
@@ -397,12 +397,14 @@
         double elapsedTime = 0;
         // last_log contains a timing metric in its last line, capture it here and return it
         // for the metrics map to report
-        if (lastLog == null || lastKmsg == null) {
-            CLog.w("Could not find last_log at directory %s, "
-                    + "or last_kmsg at directory %s", LOG_RECOV, LOG_KMSG);
-            return elapsedTime;
-        }
         try {
+            if (lastLog == null || lastKmsg == null) {
+                CLog.w(
+                        "Could not find last_log at directory %s, or last_kmsg at directory %s",
+                        LOG_RECOV, LOG_KMSG);
+                return elapsedTime;
+            }
+
             try {
                 String[] lastLogLines = StreamUtil.getStringFromSource(lastLog).split("\n");
                 String endLine = lastLogLines[lastLogLines.length - 1];
diff --git a/remote/src/com/android/tradefed/command/remote/DeviceDescriptor.java b/remote/src/com/android/tradefed/command/remote/DeviceDescriptor.java
index 8ca3e71..f3060ed 100644
--- a/remote/src/com/android/tradefed/command/remote/DeviceDescriptor.java
+++ b/remote/src/com/android/tradefed/command/remote/DeviceDescriptor.java
@@ -16,6 +16,7 @@
 
 package com.android.tradefed.command.remote;
 
+import com.android.ddmlib.IDevice;
 import com.android.ddmlib.IDevice.DeviceState;
 import com.android.tradefed.device.DeviceAllocationState;
 
@@ -37,6 +38,7 @@
     private final String mMacAddress;
     private final String mSimState;
     private final String mSimOperator;
+    private final IDevice mIDevice;
 
     public DeviceDescriptor(String serial, boolean isStubDevice, DeviceAllocationState state,
             String product, String productVariant, String sdkVersion, String buildId,
@@ -49,14 +51,38 @@
             String product, String productVariant, String sdkVersion, String buildId,
             String batteryLevel, String deviceClass, String macAddress, String simState,
             String simOperator) {
-        this(serial, isStubDevice, null, state, product, productVariant, sdkVersion, buildId,
-                batteryLevel, deviceClass, macAddress, simState, simOperator);
+        this(
+                serial,
+                isStubDevice,
+                null,
+                state,
+                product,
+                productVariant,
+                sdkVersion,
+                buildId,
+                batteryLevel,
+                deviceClass,
+                macAddress,
+                simState,
+                simOperator,
+                null);
     }
 
-    public DeviceDescriptor(String serial, boolean isStubDevice, DeviceState deviceState,
-            DeviceAllocationState state, String product, String productVariant, String sdkVersion,
-            String buildId, String batteryLevel, String deviceClass, String macAddress,
-            String simState, String simOperator) {
+    public DeviceDescriptor(
+            String serial,
+            boolean isStubDevice,
+            DeviceState deviceState,
+            DeviceAllocationState state,
+            String product,
+            String productVariant,
+            String sdkVersion,
+            String buildId,
+            String batteryLevel,
+            String deviceClass,
+            String macAddress,
+            String simState,
+            String simOperator,
+            IDevice idevice) {
         mSerial = serial;
         mIsStubDevice = isStubDevice;
         mDeviceState = deviceState;
@@ -70,6 +96,7 @@
         mMacAddress = macAddress;
         mSimState = simState;
         mSimOperator = simOperator;
+        mIDevice = idevice;
     }
 
     public String getSerial() {
@@ -127,6 +154,13 @@
         return mSimOperator;
     }
 
+    public String getProperty(String name) {
+        if (mIDevice == null) {
+            throw new UnsupportedOperationException("this descriptor does not have IDevice");
+        }
+        return mIDevice.getProperty(name);
+    }
+
     /**
      * Provides a description with serials, product and build id
      */
diff --git a/src/com/android/tradefed/build/AppBuildInfo.java b/src/com/android/tradefed/build/AppBuildInfo.java
index 12826ae..f1647a7 100644
--- a/src/com/android/tradefed/build/AppBuildInfo.java
+++ b/src/com/android/tradefed/build/AppBuildInfo.java
@@ -28,6 +28,7 @@
  */
 public class AppBuildInfo extends BuildInfo implements IAppBuildInfo {
 
+    private static final long serialVersionUID = BuildSerializedVersion.VERSION;
     private List<VersionedFile> mAppPackageFiles = new ArrayList<VersionedFile>();
 
     /**
diff --git a/src/com/android/tradefed/build/AppDeviceBuildInfo.java b/src/com/android/tradefed/build/AppDeviceBuildInfo.java
index dbda1e3..8c42eea 100644
--- a/src/com/android/tradefed/build/AppDeviceBuildInfo.java
+++ b/src/com/android/tradefed/build/AppDeviceBuildInfo.java
@@ -27,6 +27,7 @@
  */
 public class AppDeviceBuildInfo extends BuildInfo implements IDeviceBuildInfo, IAppBuildInfo {
 
+    private static final long serialVersionUID = BuildSerializedVersion.VERSION;
     private IDeviceBuildInfo mDeviceBuild;
     private IAppBuildInfo mAppBuildInfo;
 
diff --git a/src/com/android/tradefed/build/BootstrapBuildProvider.java b/src/com/android/tradefed/build/BootstrapBuildProvider.java
index 552e667..ed51648 100644
--- a/src/com/android/tradefed/build/BootstrapBuildProvider.java
+++ b/src/com/android/tradefed/build/BootstrapBuildProvider.java
@@ -21,8 +21,10 @@
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.log.LogUtil.CLog;
+import com.android.tradefed.util.FileUtil;
 
 import java.io.File;
+import java.io.IOException;
 
 /**
  * A {@link IDeviceBuildProvider} that bootstraps build info from the test device
@@ -64,6 +66,8 @@
     @Option(name="tests-dir", description="Path to top directory of expanded tests zip")
     private File mTestsDir = null;
 
+    private boolean mCreatedTestDir = false;
+
     @Override
     public IBuildInfo getBuild() throws BuildRetrievalError {
         throw new UnsupportedOperationException("Call getBuild(ITestDevice)");
@@ -77,7 +81,10 @@
 
     @Override
     public void cleanUp(IBuildInfo info) {
-        // no op
+        // If we created the tests dir, we delete it.
+        if (mCreatedTestDir) {
+            FileUtil.recursiveDelete(((IDeviceBuildInfo) info).getTestsDir());
+        }
     }
 
     @Override
@@ -103,6 +110,16 @@
         if (mTestsDir != null && mTestsDir.isDirectory()) {
             info.setFile("testsdir", mTestsDir, buildId);
         }
+        // Avoid tests dir being null, by creating a temporary dir.
+        if (mTestsDir == null) {
+            mCreatedTestDir = true;
+            try {
+                mTestsDir = FileUtil.createTempDir("bootstrap-test-dir");
+            } catch (IOException e) {
+                throw new BuildRetrievalError(e.getMessage(), e);
+            }
+            ((IDeviceBuildInfo) info).setTestsDir(mTestsDir, "1");
+        }
         return info;
     }
 }
diff --git a/src/com/android/tradefed/build/BuildInfo.java b/src/com/android/tradefed/build/BuildInfo.java
index c466da5..f334cb1 100644
--- a/src/com/android/tradefed/build/BuildInfo.java
+++ b/src/com/android/tradefed/build/BuildInfo.java
@@ -35,6 +35,7 @@
  * with a {@link ITestDevice}.
  */
 public class BuildInfo implements IBuildInfo {
+    private static final long serialVersionUID = BuildSerializedVersion.VERSION;
     private static final String BUILD_ALIAS_KEY = "build_alias";
 
     private String mBuildId = UNKNOWN_BUILD_ID;
diff --git a/src/com/android/tradefed/build/BuildSerializedVersion.java b/src/com/android/tradefed/build/BuildSerializedVersion.java
new file mode 100644
index 0000000..373cc79
--- /dev/null
+++ b/src/com/android/tradefed/build/BuildSerializedVersion.java
@@ -0,0 +1,26 @@
+/*
+ * Copyright (C) 2017 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.build;
+
+/**
+ * Class that contains the current serialization version of all {@link IBuildInfo}. This allows to
+ * synchronize all the build info version of serialization and update them all together if a non
+ * compatible change is made.
+ */
+public class BuildSerializedVersion {
+    public static final long VERSION = 1L;
+}
diff --git a/src/com/android/tradefed/build/DeviceBuildInfo.java b/src/com/android/tradefed/build/DeviceBuildInfo.java
index c9053ef..d982c31 100644
--- a/src/com/android/tradefed/build/DeviceBuildInfo.java
+++ b/src/com/android/tradefed/build/DeviceBuildInfo.java
@@ -24,6 +24,7 @@
  */
 public class DeviceBuildInfo extends BuildInfo implements IDeviceBuildInfo {
 
+    private static final long serialVersionUID = BuildSerializedVersion.VERSION;
     private static final String DEVICE_IMAGE_NAME = "device";
     private static final String USERDATA_IMAGE_NAME = "userdata";
     private static final String TESTDIR_IMAGE_NAME = "testsdir";
diff --git a/src/com/android/tradefed/build/DeviceFolderBuildInfo.java b/src/com/android/tradefed/build/DeviceFolderBuildInfo.java
index 0fa4abb..843c081 100644
--- a/src/com/android/tradefed/build/DeviceFolderBuildInfo.java
+++ b/src/com/android/tradefed/build/DeviceFolderBuildInfo.java
@@ -26,6 +26,7 @@
  */
 public class DeviceFolderBuildInfo extends BuildInfo implements IDeviceBuildInfo, IFolderBuildInfo {
 
+    private static final long serialVersionUID = BuildSerializedVersion.VERSION;
     private IDeviceBuildInfo mDeviceBuild;
     private IFolderBuildInfo mFolderBuild;
 
diff --git a/src/com/android/tradefed/build/FolderBuildInfo.java b/src/com/android/tradefed/build/FolderBuildInfo.java
index 689421c..03d88cc 100644
--- a/src/com/android/tradefed/build/FolderBuildInfo.java
+++ b/src/com/android/tradefed/build/FolderBuildInfo.java
@@ -26,6 +26,7 @@
  */
 public class FolderBuildInfo extends BuildInfo implements IFolderBuildInfo {
 
+    private static final long serialVersionUID = BuildSerializedVersion.VERSION;
     private File mRootDir;
 
     /**
diff --git a/src/com/android/tradefed/build/KernelBuildInfo.java b/src/com/android/tradefed/build/KernelBuildInfo.java
index 44b6bdd..bc8c868 100644
--- a/src/com/android/tradefed/build/KernelBuildInfo.java
+++ b/src/com/android/tradefed/build/KernelBuildInfo.java
@@ -22,6 +22,7 @@
  * A {@link IBuildInfo} that represents a kernel build.
  */
 public class KernelBuildInfo extends BuildInfo implements IKernelBuildInfo {
+    private static final long serialVersionUID = BuildSerializedVersion.VERSION;
     private final static String KERNEL_FILE = "kernel";
 
     private String mSha1 = null;
diff --git a/src/com/android/tradefed/build/KernelDeviceBuildInfo.java b/src/com/android/tradefed/build/KernelDeviceBuildInfo.java
index 4031efd..2c72401 100644
--- a/src/com/android/tradefed/build/KernelDeviceBuildInfo.java
+++ b/src/com/android/tradefed/build/KernelDeviceBuildInfo.java
@@ -22,6 +22,7 @@
  */
 public class KernelDeviceBuildInfo extends BuildInfo implements IDeviceBuildInfo,
         IKernelBuildInfo {
+    private static final long serialVersionUID = BuildSerializedVersion.VERSION;
     private IDeviceBuildInfo mDeviceBuild = new DeviceBuildInfo();
     private IKernelBuildInfo mKernelBuild = new KernelBuildInfo();
 
diff --git a/src/com/android/tradefed/build/OtaDeviceBuildInfo.java b/src/com/android/tradefed/build/OtaDeviceBuildInfo.java
index a065c6f..c0ae59b 100644
--- a/src/com/android/tradefed/build/OtaDeviceBuildInfo.java
+++ b/src/com/android/tradefed/build/OtaDeviceBuildInfo.java
@@ -36,6 +36,7 @@
  */
 public class OtaDeviceBuildInfo implements IDeviceBuildInfo {
 
+    private static final long serialVersionUID = BuildSerializedVersion.VERSION;
     protected IDeviceBuildInfo mOtaBuild;
     protected IDeviceBuildInfo mBaselineBuild;
     protected boolean mReportTargetBuild = false;
diff --git a/src/com/android/tradefed/build/OtaToolsDeviceBuildInfo.java b/src/com/android/tradefed/build/OtaToolsDeviceBuildInfo.java
index a72f527..319e43b 100644
--- a/src/com/android/tradefed/build/OtaToolsDeviceBuildInfo.java
+++ b/src/com/android/tradefed/build/OtaToolsDeviceBuildInfo.java
@@ -24,6 +24,7 @@
  * An {@link OtaDeviceBuildInfo} that also contains an otatools directory.
  */
 public class OtaToolsDeviceBuildInfo extends OtaDeviceBuildInfo {
+    private static final long serialVersionUID = BuildSerializedVersion.VERSION;
     private File mOtaToolsDir;
 
     /**
diff --git a/src/com/android/tradefed/build/OtatoolsBuildInfo.java b/src/com/android/tradefed/build/OtatoolsBuildInfo.java
index 3a3f7ad..9f9d3f1 100644
--- a/src/com/android/tradefed/build/OtatoolsBuildInfo.java
+++ b/src/com/android/tradefed/build/OtatoolsBuildInfo.java
@@ -22,6 +22,7 @@
  * An {@link IBuildInfo} that contains otatools artifacts.
  */
 public class OtatoolsBuildInfo extends BuildInfo {
+    private static final long serialVersionUID = BuildSerializedVersion.VERSION;
     private static final String SECURITY_DIR_NAME = "otatools_security";
     private static final String BIN_DIR_NAME = "otatools_bin";
     private static final String FRAMEWORK_DIR_NAME = "otatools_framework";
@@ -82,4 +83,4 @@
     public File getReleasetoolsDir() {
         return getFile(RELEASETOOLS_DIR_NAME);
     }
-}
\ No newline at end of file
+}
diff --git a/src/com/android/tradefed/build/SdkBuildInfo.java b/src/com/android/tradefed/build/SdkBuildInfo.java
index a495135..6c3efe4 100644
--- a/src/com/android/tradefed/build/SdkBuildInfo.java
+++ b/src/com/android/tradefed/build/SdkBuildInfo.java
@@ -30,6 +30,7 @@
  */
 public class SdkBuildInfo extends BuildInfo implements ISdkBuildInfo {
 
+    private static final long serialVersionUID = BuildSerializedVersion.VERSION;
     private File mTestDir = null;
     private File mSdkDir = null;
     private boolean mDeleteSdkDirParent;
diff --git a/src/com/android/tradefed/build/SdkFolderBuildInfo.java b/src/com/android/tradefed/build/SdkFolderBuildInfo.java
index 6f8bdd4..0424f19 100644
--- a/src/com/android/tradefed/build/SdkFolderBuildInfo.java
+++ b/src/com/android/tradefed/build/SdkFolderBuildInfo.java
@@ -25,6 +25,7 @@
  */
 public class SdkFolderBuildInfo extends BuildInfo implements ISdkBuildInfo, IFolderBuildInfo {
 
+    private static final long serialVersionUID = BuildSerializedVersion.VERSION;
     private ISdkBuildInfo mSdkBuild;
     private IFolderBuildInfo mFolderBuild;
 
diff --git a/src/com/android/tradefed/build/VersionedFile.java b/src/com/android/tradefed/build/VersionedFile.java
index 2c71b07..133f767 100644
--- a/src/com/android/tradefed/build/VersionedFile.java
+++ b/src/com/android/tradefed/build/VersionedFile.java
@@ -20,6 +20,7 @@
 
 /** Data structure representing a file that has an associated version. */
 public class VersionedFile implements Serializable {
+    private static final long serialVersionUID = BuildSerializedVersion.VERSION;
     private final File mFile;
     private final String mVersion;
 
diff --git a/src/com/android/tradefed/command/CommandRunner.java b/src/com/android/tradefed/command/CommandRunner.java
index 89e8834..6ebeace 100644
--- a/src/com/android/tradefed/command/CommandRunner.java
+++ b/src/com/android/tradefed/command/CommandRunner.java
@@ -18,6 +18,7 @@
 
 import com.android.tradefed.config.ConfigurationException;
 import com.android.tradefed.config.GlobalConfiguration;
+import com.android.tradefed.device.NoDeviceException;
 
 import com.google.common.annotations.VisibleForTesting;
 
@@ -33,6 +34,8 @@
     private ICommandScheduler mScheduler;
     private ExitCode mErrorCode = ExitCode.NO_ERROR;
 
+    private static final long CHECK_DEVICE_TIMEOUT = 15000;
+
     public CommandRunner() {}
 
     public ExitCode getErrorCode() {
@@ -59,6 +62,12 @@
         e.printStackTrace();
     }
 
+    /** Returns the timeout after which to check for the command. */
+    @VisibleForTesting
+    long getCheckDeviceTimeout() {
+        return CHECK_DEVICE_TIMEOUT;
+    }
+
     /**
      * The main method to run the command.
      *
@@ -77,6 +86,15 @@
             mScheduler.shutdownOnEmpty();
         }
         try {
+            mScheduler.join(getCheckDeviceTimeout());
+            // After 15 seconds we check if the command was executed.
+            if (mScheduler.getReadyCommandCount() > 0) {
+                printStackTrace(new NoDeviceException("No device was allocated for the command."));
+                mErrorCode = ExitCode.NO_DEVICE_ALLOCATED;
+                mScheduler.removeAllCommands();
+                mScheduler.shutdown();
+                return;
+            }
             mScheduler.join();
             // If no error code has been raised yet, we checked the invocation error code.
             if (ExitCode.NO_ERROR.equals(mErrorCode)) {
@@ -109,7 +127,8 @@
         DEVICE_UNRESPONSIVE(3),
         DEVICE_UNAVAILABLE(4),
         FATAL_HOST_ERROR(5),
-        THROWABLE_EXCEPTION(6);
+        THROWABLE_EXCEPTION(6),
+        NO_DEVICE_ALLOCATED(7);
 
         private final int mCodeValue;
 
diff --git a/src/com/android/tradefed/command/CommandScheduler.java b/src/com/android/tradefed/command/CommandScheduler.java
index 2fa6118..ed4eaee 100644
--- a/src/com/android/tradefed/command/CommandScheduler.java
+++ b/src/com/android/tradefed/command/CommandScheduler.java
@@ -16,8 +16,6 @@
 
 package com.android.tradefed.command;
 
-import com.google.common.annotations.VisibleForTesting;
-
 import com.android.ddmlib.DdmPreferences;
 import com.android.ddmlib.IDevice;
 import com.android.ddmlib.Log;
@@ -30,7 +28,6 @@
 import com.android.tradefed.command.remote.RemoteClient;
 import com.android.tradefed.command.remote.RemoteException;
 import com.android.tradefed.command.remote.RemoteManager;
-import com.android.tradefed.config.ConfigurationDef;
 import com.android.tradefed.config.ConfigurationException;
 import com.android.tradefed.config.ConfigurationFactory;
 import com.android.tradefed.config.GlobalConfiguration;
@@ -75,6 +72,8 @@
 import com.android.tradefed.util.keystore.IKeyStoreFactory;
 import com.android.tradefed.util.keystore.KeyStoreException;
 
+import com.google.common.annotations.VisibleForTesting;
+
 import org.json.JSONException;
 
 import java.io.File;
@@ -448,17 +447,6 @@
             mDeviceManager = deviceManager;
         }
 
-        @Deprecated
-        @Override
-        public void invocationComplete(ITestDevice device, FreeDeviceState deviceState) {
-            IInvocationContext context = new InvocationContext();
-            // Fake a single device context for compatibility
-            context.addAllocatedDevice(ConfigurationDef.DEFAULT_DEVICE_NAME, device);
-            Map<ITestDevice, FreeDeviceState> state = new HashMap<>();
-            state.put(device, deviceState);
-            invocationComplete(context, state);
-        }
-
         @Override
         public void invocationComplete(IInvocationContext context,
                 Map<ITestDevice, FreeDeviceState> devicesStates) {
@@ -2138,4 +2126,9 @@
         mLastInvocationExitCode = code;
         mLastInvocationThrowable = throwable;
     }
+
+    @Override
+    public synchronized int getReadyCommandCount() {
+        return mReadyCommands.size();
+    }
 }
diff --git a/src/com/android/tradefed/command/ICommandScheduler.java b/src/com/android/tradefed/command/ICommandScheduler.java
index 5ae5c9a..b24c359 100644
--- a/src/com/android/tradefed/command/ICommandScheduler.java
+++ b/src/com/android/tradefed/command/ICommandScheduler.java
@@ -44,17 +44,6 @@
          * Callback when entire invocation has completed, including all
          * {@link ITestInvocationListener#invocationEnded(long)} events.
          *
-         * @param device
-         * @param deviceState
-         * @deprecated use {@link #invocationComplete(IInvocationContext, Map)}.
-         */
-        @Deprecated
-        public void invocationComplete(ITestDevice device, FreeDeviceState deviceState);
-
-        /**
-         * Callback when entire invocation has completed, including all
-         * {@link ITestInvocationListener#invocationEnded(long)} events.
-         *
          * @param metadata
          * @param devicesStates
          */
@@ -202,8 +191,15 @@
     public void join() throws InterruptedException;
 
     /**
-     * Waits for scheduler to start running, including waiting for handover from old TF to
-     * complete if applicable.
+     * Waits for scheduler to complete or timeout after the duration specified in milliseconds.
+     *
+     * @see Thread#join(long)
+     */
+    public void join(long millis) throws InterruptedException;
+
+    /**
+     * Waits for scheduler to start running, including waiting for handover from old TF to complete
+     * if applicable.
      */
     public void await() throws InterruptedException;
 
@@ -291,4 +287,7 @@
      * and a stack trace that can be returned.
      */
     public void setLastInvocationExitCode(ExitCode code, Throwable stack);
+
+    /** Returns the number of Commands in ready state in the queue. */
+    public int getReadyCommandCount();
 }
diff --git a/src/com/android/tradefed/command/remote/ExecCommandTracker.java b/src/com/android/tradefed/command/remote/ExecCommandTracker.java
index 2b3709f..28d7ebc 100644
--- a/src/com/android/tradefed/command/remote/ExecCommandTracker.java
+++ b/src/com/android/tradefed/command/remote/ExecCommandTracker.java
@@ -16,11 +16,9 @@
 package com.android.tradefed.command.remote;
 
 import com.android.tradefed.command.ICommandScheduler.IScheduledInvocationListener;
-import com.android.tradefed.config.Configuration;
 import com.android.tradefed.device.FreeDeviceState;
 import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.invoker.IInvocationContext;
-import com.android.tradefed.invoker.InvocationContext;
 
 import com.google.common.collect.ImmutableMap;
 
@@ -45,21 +43,6 @@
         mErrorDetails = outputStream.toString();
     }
 
-    /**
-     * {@inheritDoc}
-     * @deprecated use {@link #invocationComplete(IInvocationContext, Map)} instead.
-     */
-    @Deprecated
-    @Override
-    public void invocationComplete(ITestDevice device, FreeDeviceState deviceState) {
-        IInvocationContext stubMeta = new InvocationContext();
-        stubMeta.addAllocatedDevice(Configuration.DEVICE_NAME, device);
-        // Stub metadata for compatibility
-        Map<ITestDevice, FreeDeviceState> state = new HashMap<>();
-        state.put(device, deviceState);
-        invocationComplete(stubMeta, state);
-    }
-
     @Override
     public void invocationComplete(IInvocationContext metadata,
             Map<ITestDevice, FreeDeviceState> devicesStates) {
diff --git a/src/com/android/tradefed/config/Configuration.java b/src/com/android/tradefed/config/Configuration.java
index 1d76b43..26e2166 100644
--- a/src/com/android/tradefed/config/Configuration.java
+++ b/src/com/android/tradefed/config/Configuration.java
@@ -1217,6 +1217,8 @@
             ConfigurationUtil.dumpClassToXml(serializer, TEST_TYPE_NAME, test);
         }
 
+        ConfigurationUtil.dumpClassToXml(
+                serializer, CONFIGURATION_DESCRIPTION_TYPE_NAME, getConfigurationDescription());
         ConfigurationUtil.dumpClassToXml(serializer, LOGGER_TYPE_NAME, getLogOutput());
         ConfigurationUtil.dumpClassToXml(serializer, LOG_SAVER_TYPE_NAME, getLogSaver());
         for (ITestInvocationListener listener : getTestInvocationListeners()) {
diff --git a/src/com/android/tradefed/config/ConfigurationDescriptor.java b/src/com/android/tradefed/config/ConfigurationDescriptor.java
index 94aac9e..1cf1345 100644
--- a/src/com/android/tradefed/config/ConfigurationDescriptor.java
+++ b/src/com/android/tradefed/config/ConfigurationDescriptor.java
@@ -15,10 +15,13 @@
  */
 package com.android.tradefed.config;
 
+import com.android.tradefed.build.BuildSerializedVersion;
+import com.android.tradefed.testtype.IAbi;
 import com.android.tradefed.util.MultiMap;
 
 import com.google.common.annotations.VisibleForTesting;
 
+import java.io.Serializable;
 import java.util.ArrayList;
 import java.util.List;
 
@@ -28,7 +31,8 @@
  * xml.
  */
 @OptionClass(alias = "config-descriptor")
-public class ConfigurationDescriptor {
+public class ConfigurationDescriptor implements Serializable {
+    private static final long serialVersionUID = BuildSerializedVersion.VERSION;
 
     @Option(name = "test-suite-tag", description = "A membership tag to suite. Can be repeated.")
     private List<String> mSuiteTags = new ArrayList<>();
@@ -46,6 +50,18 @@
     )
     private boolean mNotShardable = false;
 
+    @Option(
+        name = "not-strict-shardable",
+        description =
+                "A metadata to allows a suite configuration to specify that it cannot be "
+                        + "sharded in a strict context (independent shards). If a config is already "
+                        + "not-shardable, it will be not-strict-shardable."
+    )
+    private boolean mNotStrictShardable = false;
+
+    /** Optional Abi information the configuration will be run against. */
+    private IAbi mAbi = null;
+
     /** Returns the list of suite tags the test is part of. */
     public List<String> getSuiteTags() {
         return mSuiteTags;
@@ -75,4 +91,19 @@
     public boolean isNotShardable() {
         return mNotShardable;
     }
+
+    /** Returns if the configuration is strict shardable or not as part of a suite */
+    public boolean isNotStrictShardable() {
+        return mNotStrictShardable;
+    }
+
+    /** Sets the abi the configuration is going to run against. */
+    public void setAbi(IAbi abi) {
+        mAbi = abi;
+    }
+
+    /** Returns the abi the configuration is running against if known, null otherwise. */
+    public IAbi getAbi() {
+        return mAbi;
+    }
 }
diff --git a/src/com/android/tradefed/config/ConfigurationFactory.java b/src/com/android/tradefed/config/ConfigurationFactory.java
index 236eb39..628702b 100644
--- a/src/com/android/tradefed/config/ConfigurationFactory.java
+++ b/src/com/android/tradefed/config/ConfigurationFactory.java
@@ -193,8 +193,8 @@
      *     on the value of environment variables.
      */
     @VisibleForTesting
-    List<File> getTestCasesDirs() {
-        return SystemUtil.getTestCasesDirs();
+    List<File> getExternalTestCasesDirs() {
+        return SystemUtil.getExternalTestCasesDirs();
     }
 
     /**
@@ -216,7 +216,7 @@
     @VisibleForTesting
     File getTestCaseConfigPath(String name) {
         String[] possibleConfigFileNames = {name + ".xml", name + ".config"};
-        for (File testCasesDir : getTestCasesDirs()) {
+        for (File testCasesDir : getExternalTestCasesDirs()) {
             for (String configFileName : possibleConfigFileNames) {
                 File config = FileUtil.findFile(testCasesDir, configFileName);
                 if (config != null) {
@@ -623,7 +623,7 @@
      */
     @VisibleForTesting
     Set<String> getConfigNamesFromTestCases(String subPath) {
-        return ConfigurationUtil.getConfigNamesFromDirs(subPath, getTestCasesDirs());
+        return ConfigurationUtil.getConfigNamesFromDirs(subPath, getExternalTestCasesDirs());
     }
 
     /**
diff --git a/src/com/android/tradefed/device/DeviceManager.java b/src/com/android/tradefed/device/DeviceManager.java
index 06ae4c8..d31fee2 100644
--- a/src/com/android/tradefed/device/DeviceManager.java
+++ b/src/com/android/tradefed/device/DeviceManager.java
@@ -59,10 +59,14 @@
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.regex.Pattern;
 
 @OptionClass(alias = "dmgr", global_namespace = false)
 public class DeviceManager implements IDeviceManager {
 
+    /** Display string for unknown properties */
+    public static final String UNKNOWN_DISPLAY_STRING = "unknown";
+
     /** max wait time in ms for fastboot devices command to complete */
     private static final long FASTBOOT_CMD_TIMEOUT = 1 * 60 * 1000;
     /** time to wait in ms between fastboot devices requests */
@@ -85,6 +89,17 @@
     private static final String EMULATOR_SERIAL_PREFIX = "emulator";
     private static final String TCP_DEVICE_SERIAL_PREFIX = "tcp-device";
 
+    /**
+     * Pattern for a device listed by 'adb devices':
+     *
+     * <p>List of devices attached
+     *
+     * <p>serial1 device
+     *
+     * <p>serial2 device
+     */
+    private static final String DEVICE_LIST_PATTERN = "(.*)(\n)(%s)(\\s+)(device)(.*?)";
+
     protected DeviceMonitorMultiplexer mDvcMon = new DeviceMonitorMultiplexer();
 
     private boolean mIsInitialized = false;
@@ -528,17 +543,28 @@
     /**
      * Helper method to convert from a {@link com.android.tradefed.device.FreeDeviceState} to a
      * {@link com.android.tradefed.device.DeviceEvent}
+     *
      * @param managedDevice
      */
-    static DeviceEvent getEventFromFree(IManagedTestDevice managedDevice, FreeDeviceState deviceState) {
+    private DeviceEvent getEventFromFree(
+            IManagedTestDevice managedDevice, FreeDeviceState deviceState) {
         switch (deviceState) {
             case UNRESPONSIVE:
                 return DeviceEvent.FREE_UNRESPONSIVE;
             case AVAILABLE:
                 return DeviceEvent.FREE_AVAILABLE;
             case UNAVAILABLE:
-                if (managedDevice.getDeviceState() == TestDeviceState.NOT_AVAILABLE) {
-                    return DeviceEvent.FREE_UNKNOWN;
+                // We double check if device is still showing in adb or not to confirm the
+                // connection is gone.
+                if (TestDeviceState.NOT_AVAILABLE.equals(managedDevice.getDeviceState())) {
+                    String devices = executeGlobalAdbCommand("devices");
+                    Pattern p =
+                            Pattern.compile(
+                                    String.format(
+                                            DEVICE_LIST_PATTERN, managedDevice.getSerialNumber()));
+                    if (devices == null || !p.matcher(devices).find()) {
+                        return DeviceEvent.FREE_UNKNOWN;
+                    }
                 }
                 return DeviceEvent.FREE_UNAVAILABLE;
             case IGNORE:
@@ -883,7 +909,8 @@
                             d.getDeviceClass(),
                             getDisplay(d.getMacAddress()),
                             getDisplay(d.getSimState()),
-                            getDisplay(d.getSimOperator())));
+                            getDisplay(d.getSimOperator()),
+                            idevice));
         }
         return serialStates;
     }
@@ -956,7 +983,7 @@
      * Return the displayable string for given object
      */
     private String getDisplay(Object o) {
-        return o == null ? "unknown" : o.toString();
+        return o == null ? UNKNOWN_DISPLAY_STRING : o.toString();
     }
 
     /**
diff --git a/src/com/android/tradefed/device/LargeOutputReceiver.java b/src/com/android/tradefed/device/LargeOutputReceiver.java
index 8af95af..d402db1 100644
--- a/src/com/android/tradefed/device/LargeOutputReceiver.java
+++ b/src/com/android/tradefed/device/LargeOutputReceiver.java
@@ -82,7 +82,7 @@
     public synchronized InputStreamSource getData() {
         if (mOutStream != null) {
             try {
-                return new SnapshotInputStreamSource(mOutStream.getData());
+                return new SnapshotInputStreamSource("LargeOutputReceiver", mOutStream.getData());
             } catch (IOException e) {
                 CLog.e("failed to get %s data for %s.", mDescriptor, mSerialNumber);
                 CLog.e(e);
diff --git a/src/com/android/tradefed/device/NativeDevice.java b/src/com/android/tradefed/device/NativeDevice.java
index 978a6a6..b50a22a 100644
--- a/src/com/android/tradefed/device/NativeDevice.java
+++ b/src/com/android/tradefed/device/NativeDevice.java
@@ -131,10 +131,6 @@
     private static final int ENCRYPTION_INPLACE_TIMEOUT_MIN = 2 * 60;
     /** Encrypting with wipe can take up to 20 minutes. */
     private static final long ENCRYPTION_WIPE_TIMEOUT_MIN = 20;
-    /** Beginning of the string returned by vdc for "vdc cryptfs enablecrypto". */
-    private static final String ENCRYPTION_SUPPORTED_CODE = "500";
-    /** Message in the string returned by vdc for "vdc cryptfs enablecrypto". */
-    private static final String ENCRYPTION_SUPPORTED_USAGE = "Usage: ";
 
     /** The time in ms to wait before starting logcat for a device */
     private int mLogStartDelay = 5*1000;
@@ -883,11 +879,19 @@
                             status = true;
                         } catch (SyncException e) {
                             CLog.w(
-                                    "Failed to push %s to %s on device %s. Message %s",
+                                    "Failed to push %s to %s on device %s. Message: '%s'. "
+                                            + "Error code: %s",
                                     localFile.getAbsolutePath(),
                                     remoteFilePath,
                                     getSerialNumber(),
-                                    e.getMessage());
+                                    e.getMessage(),
+                                    e.getErrorCode());
+                            // TODO: check if ddmlib can report a better error
+                            if (SyncError.TRANSFER_PROTOCOL_ERROR.equals(e.getErrorCode())) {
+                                if (e.getMessage().contains("Permission denied")) {
+                                    return false;
+                                }
+                            }
                             throw e;
                         } finally {
                             if (syncService != null) {
@@ -2626,6 +2630,9 @@
         // if its necessary or not
         if (isAdbRoot()) {
             CLog.i("adb is already running as root on %s", getSerialNumber());
+            // Still check for online, in some case we could see the root, but device could be
+            // very early in its cycle.
+            waitForDeviceOnline();
             return true;
         }
         // Don't enable root if user requested no root
@@ -2943,8 +2950,10 @@
         }
         enableAdbRoot();
         String output = executeShellCommand("vdc cryptfs enablecrypto").trim();
-        mIsEncryptionSupported = (output != null && output.startsWith(ENCRYPTION_SUPPORTED_CODE) &&
-                output.contains(ENCRYPTION_SUPPORTED_USAGE));
+
+        mIsEncryptionSupported =
+                (output != null
+                        && Pattern.matches("(500)(\\s+)(\\d+)(\\s+)(Usage)(.*)(:)(.*)", output));
         return mIsEncryptionSupported;
     }
 
@@ -3179,7 +3188,8 @@
                         getSerialNumber());
             } else {
                 try {
-                    return new SnapshotInputStreamSource(mEmulatorOutput.getData());
+                    return new SnapshotInputStreamSource(
+                            "getEmulatorOutput", mEmulatorOutput.getData());
                 } catch (IOException e) {
                     CLog.e("Failed to get %s data.", getSerialNumber());
                     CLog.e(e);
diff --git a/src/com/android/tradefed/device/StubDevice.java b/src/com/android/tradefed/device/StubDevice.java
index dba528d..01cdbdb 100644
--- a/src/com/android/tradefed/device/StubDevice.java
+++ b/src/com/android/tradefed/device/StubDevice.java
@@ -449,6 +449,19 @@
         throw new IOException("stub");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public void executeShellCommand(
+            String command,
+            IShellOutputReceiver receiver,
+            long maxTimeout,
+            long maxTimeToOutputResponse,
+            TimeUnit maxTimeUnits)
+            throws TimeoutException, AdbCommandRejectedException, ShellCommandUnresponsiveException,
+                    IOException {
+        throw new IOException("stub");
+    }
+
     /**
      * {@inheritDoc}
      */
diff --git a/src/com/android/tradefed/device/TestDevice.java b/src/com/android/tradefed/device/TestDevice.java
index 066ef12..dd9beca 100644
--- a/src/com/android/tradefed/device/TestDevice.java
+++ b/src/com/android/tradefed/device/TestDevice.java
@@ -1145,7 +1145,7 @@
      */
     @Override
     IWifiHelper createWifiHelper() throws DeviceNotAvailableException {
-        return new WifiHelper(this);
+        return new WifiHelper(this, mOptions.getWifiUtilAPKPath());
     }
 
     /** {@inheritDoc} */
diff --git a/src/com/android/tradefed/device/TestDeviceOptions.java b/src/com/android/tradefed/device/TestDeviceOptions.java
index 3562763..8675b71 100644
--- a/src/com/android/tradefed/device/TestDeviceOptions.java
+++ b/src/com/android/tradefed/device/TestDeviceOptions.java
@@ -92,6 +92,9 @@
                     + " a binary exponential back-offs when retrying.")
     private boolean mWifiExpoRetryEnabled = true;
 
+    @Option(name = "wifiutil-apk-path", description = "path to the wifiutil APK file")
+    private String mWifiUtilAPKPath = null;
+
     @Option(name = "post-boot-command",
             description = "shell command to run after reboots during invocation")
     private List<String> mPostBootCommands = new ArrayList<String>();
@@ -326,4 +329,9 @@
     public boolean isWifiExpoRetryEnabled() {
         return mWifiExpoRetryEnabled;
     }
+
+    /** @return the wifiutil apk path */
+    public String getWifiUtilAPKPath() {
+        return mWifiUtilAPKPath;
+    }
 }
diff --git a/src/com/android/tradefed/device/WifiHelper.java b/src/com/android/tradefed/device/WifiHelper.java
index 2f12a0a..6a4e7da 100644
--- a/src/com/android/tradefed/device/WifiHelper.java
+++ b/src/com/android/tradefed/device/WifiHelper.java
@@ -21,6 +21,8 @@
 import com.android.tradefed.util.IRunUtil;
 import com.android.tradefed.util.RunUtil;
 
+import com.google.common.annotations.VisibleForTesting;
+
 import org.json.JSONException;
 import org.json.JSONObject;
 
@@ -61,10 +63,17 @@
     private static final long DEFAULT_WIFI_STATE_TIMEOUT = 30*1000;
 
     private final ITestDevice mDevice;
+    private File mWifiUtilApkFile;
 
     public WifiHelper(ITestDevice device) throws DeviceNotAvailableException {
         mDevice = device;
-        ensureDeviceSetup();
+        ensureDeviceSetup(null);
+    }
+
+    public WifiHelper(ITestDevice device, String wifiUtilApkPath)
+            throws DeviceNotAvailableException {
+        mDevice = device;
+        ensureDeviceSetup(wifiUtilApkPath);
     }
 
     /**
@@ -76,7 +85,7 @@
         return RunUtil.getDefault();
     }
 
-    void ensureDeviceSetup() throws DeviceNotAvailableException {
+    void ensureDeviceSetup(String wifiUtilApkPath) throws DeviceNotAvailableException {
         final String inst = mDevice.executeShellCommand(CHECK_PACKAGE_CMD);
         if (inst != null) {
             Matcher matcher = PACKAGE_VERSION_PAT.matcher(inst);
@@ -92,11 +101,10 @@
         }
 
         // Attempt to install utility
-        File apkTempFile = null;
         try {
-            apkTempFile = extractWifiUtilApk();
+            setupWifiUtilApkFile(wifiUtilApkPath);
 
-            final String error = mDevice.installPackage(apkTempFile, true);
+            final String error = mDevice.installPackage(mWifiUtilApkFile, true);
             if (error == null) {
                 // Installed successfully; good to go.
                 return;
@@ -108,10 +116,31 @@
             throw new RuntimeException(String.format(
                     "Failed to unpack WifiUtil utility: %s", e.getMessage()));
         } finally {
-            FileUtil.deleteFile(apkTempFile);
+            // Delete the tmp file only if the APK is copied from classpath
+            if (wifiUtilApkPath == null) {
+                FileUtil.deleteFile(mWifiUtilApkFile);
+            }
         }
     }
 
+    private void setupWifiUtilApkFile(String wifiUtilApkPath) throws IOException {
+        if (wifiUtilApkPath != null) {
+            mWifiUtilApkFile = new File(wifiUtilApkPath);
+        } else {
+            mWifiUtilApkFile = extractWifiUtilApk();
+        }
+    }
+
+    /**
+     * Get the {@link File} object of the APK file.
+     *
+     * <p>Exposed for unit testing.
+     */
+    @VisibleForTesting
+    File getWifiUtilApkFile() {
+        return mWifiUtilApkFile;
+    }
+
     /**
      * Helper method to extract the wifi util apk from the classpath
      */
diff --git a/src/com/android/tradefed/invoker/IInvocationContext.java b/src/com/android/tradefed/invoker/IInvocationContext.java
index 6d0bb37..474f326 100644
--- a/src/com/android/tradefed/invoker/IInvocationContext.java
+++ b/src/com/android/tradefed/invoker/IInvocationContext.java
@@ -23,15 +23,16 @@
 import com.android.tradefed.util.MultiMap;
 import com.android.tradefed.util.UniqueMultiMap;
 
+import java.io.Serializable;
 import java.util.List;
 import java.util.Map;
 
 /**
- * Holds information about the Invocation for the tests to access if needed.
- * Tests should not modify the context contained here so only getters will be available, except for
- * the context attributes for reporting purpose.
+ * Holds information about the Invocation for the tests to access if needed. Tests should not modify
+ * the context contained here so only getters will be available, except for the context attributes
+ * for reporting purpose.
  */
-public interface IInvocationContext {
+public interface IInvocationContext extends Serializable {
 
     /**
      * Return the number of devices allocated for the invocation.
@@ -121,7 +122,7 @@
     /** Add several invocation attributes at once through a {@link UniqueMultiMap}. */
     public void addInvocationAttributes(UniqueMultiMap<String, String> attributesMap);
 
-    /** Returns the map of invocation attributes. */
+    /** Returns a copy of the map containing all the invocation attributes. */
     public MultiMap<String, String> getAttributes();
 
     /** Sets the descriptor associated with the test configuration that launched the invocation */
diff --git a/src/com/android/tradefed/invoker/InvocationContext.java b/src/com/android/tradefed/invoker/InvocationContext.java
index 68040e1..f8a4be4 100644
--- a/src/com/android/tradefed/invoker/InvocationContext.java
+++ b/src/com/android/tradefed/invoker/InvocationContext.java
@@ -16,6 +16,7 @@
 package com.android.tradefed.invoker;
 
 import com.android.tradefed.build.BuildInfo;
+import com.android.tradefed.build.BuildSerializedVersion;
 import com.android.tradefed.build.IBuildInfo;
 import com.android.tradefed.config.ConfigurationDescriptor;
 import com.android.tradefed.device.ITestDevice;
@@ -25,20 +26,25 @@
 import com.android.tradefed.util.MultiMap;
 import com.android.tradefed.util.UniqueMultiMap;
 
+import java.io.IOException;
+import java.io.ObjectInputStream;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
 
 /**
  * Generic implementation of a {@link IInvocationContext}.
  */
 public class InvocationContext implements IInvocationContext {
+    private static final long serialVersionUID = BuildSerializedVersion.VERSION;
 
-    private Map<ITestDevice, IBuildInfo> mAllocatedDeviceAndBuildMap;
-    /** Map of the configuration device name and the actual {@link ITestDevice} **/
-    private Map<String, ITestDevice> mNameAndDeviceMap;
+    // Transient field are not serialized
+    private transient Map<ITestDevice, IBuildInfo> mAllocatedDeviceAndBuildMap;
+    /** Map of the configuration device name and the actual {@link ITestDevice} * */
+    private transient Map<String, ITestDevice> mNameAndDeviceMap;
     private Map<String, IBuildInfo> mNameAndBuildinfoMap;
     private final UniqueMultiMap<String, String> mInvocationAttributes =
             new UniqueMultiMap<String, String>();
@@ -49,6 +55,8 @@
     /** module invocation context (when running as part of a {@link ITestSuite} */
     private IInvocationContext mModuleContext;
 
+    private boolean mLocked;
+
     /**
      * Creates a {@link BuildInfo} using default attribute values.
      */
@@ -73,6 +81,10 @@
     @Override
     public void addAllocatedDevice(String devicename, ITestDevice testDevice) {
         mNameAndDeviceMap.put(devicename, testDevice);
+        // back fill the information if possible
+        if (mNameAndBuildinfoMap.get(devicename) != null) {
+            mAllocatedDeviceAndBuildMap.put(testDevice, mNameAndBuildinfoMap.get(devicename));
+        }
     }
 
     /**
@@ -81,6 +93,13 @@
     @Override
     public void addAllocatedDevice(Map<String, ITestDevice> deviceWithName) {
         mNameAndDeviceMap.putAll(deviceWithName);
+        // back fill the information if possible
+        for (Entry<String, ITestDevice> entry : deviceWithName.entrySet()) {
+            if (mNameAndBuildinfoMap.get(entry.getKey()) != null) {
+                mAllocatedDeviceAndBuildMap.put(
+                        entry.getValue(), mNameAndBuildinfoMap.get(entry.getKey()));
+            }
+        }
     }
 
     /**
@@ -167,19 +186,30 @@
      */
     @Override
     public void addInvocationAttribute(String attributeName, String attributeValue) {
+        if (mLocked) {
+            throw new IllegalStateException(
+                    "Attempting to add invocation attribute during a test.");
+        }
         mInvocationAttributes.put(attributeName, attributeValue);
     }
 
     /** {@inheritDoc} */
     @Override
     public void addInvocationAttributes(UniqueMultiMap<String, String> attributesMap) {
+        if (mLocked) {
+            throw new IllegalStateException(
+                    "Attempting to add invocation attribute during a test.");
+        }
         mInvocationAttributes.putAll(attributesMap);
     }
 
     /** {@inheritDoc} */
     @Override
     public MultiMap<String, String> getAttributes() {
-        return mInvocationAttributes;
+        // Return a copy of the map to avoid unwanted modifications.
+        UniqueMultiMap<String, String> copy = new UniqueMultiMap<>();
+        copy.putAll(mInvocationAttributes);
+        return copy;
     }
 
     /**
@@ -257,4 +287,18 @@
     public IInvocationContext getModuleInvocationContext() {
         return mModuleContext;
     }
+
+    /** Lock the context to prevent more invocation attributes to be added. */
+    public void lockAttributes() {
+        mLocked = true;
+    }
+
+    /** Special java method that allows for custom deserialization. */
+    private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException {
+        // our "pseudo-constructor"
+        in.defaultReadObject();
+        // now we are a "live" object again, so let's init the transient field
+        mAllocatedDeviceAndBuildMap = new HashMap<ITestDevice, IBuildInfo>();
+        mNameAndDeviceMap = new LinkedHashMap<String, ITestDevice>();
+    }
 }
diff --git a/src/com/android/tradefed/invoker/ShardListener.java b/src/com/android/tradefed/invoker/ShardListener.java
index 2cd2c69..ea0f1fb 100644
--- a/src/com/android/tradefed/invoker/ShardListener.java
+++ b/src/com/android/tradefed/invoker/ShardListener.java
@@ -25,7 +25,9 @@
 import com.android.tradefed.result.ITestInvocationListener;
 import com.android.tradefed.result.InputStreamSource;
 import com.android.tradefed.result.LogDataType;
+import com.android.tradefed.util.TimeUtil;
 
+import java.util.Collection;
 import java.util.Map;
 
 /**
@@ -51,10 +53,8 @@
 
     /**
      * {@inheritDoc}
-     * @deprecated use {@link #invocationStarted(IInvocationContext)} instead.
      */
     @Override
-    @Deprecated
     public void invocationStarted(IInvocationContext context) {
         super.invocationStarted(context);
         synchronized (mMasterListener) {
@@ -112,6 +112,7 @@
     public void invocationEnded(long elapsedTime) {
         super.invocationEnded(elapsedTime);
         synchronized (mMasterListener) {
+            logShardContent(getRunResults());
             for (TestRunResult runResult : getRunResults()) {
                 mMasterListener.testRunStarted(runResult.getName(), runResult.getNumTests());
                 forwardTestResults(runResult.getTestResults());
@@ -150,4 +151,18 @@
             }
         }
     }
+
+    /** Log the content of the shard for easier debugging. */
+    private void logShardContent(Collection<TestRunResult> listResults) {
+        CLog.d("=================================================");
+        CLog.d(
+                "========== Shard Primary Device %s ==========",
+                getInvocationContext().getDevices().get(0).getSerialNumber());
+        for (TestRunResult runRes : listResults) {
+            CLog.d(
+                    "\tRan '%s' in %s",
+                    runRes.getName(), TimeUtil.formatElapsedTime(runRes.getElapsedTime()));
+        }
+        CLog.d("=================================================");
+    }
 }
diff --git a/src/com/android/tradefed/invoker/TestInvocation.java b/src/com/android/tradefed/invoker/TestInvocation.java
index 2d42caa..605f701 100644
--- a/src/com/android/tradefed/invoker/TestInvocation.java
+++ b/src/com/android/tradefed/invoker/TestInvocation.java
@@ -20,6 +20,7 @@
 import com.android.tradefed.build.BuildRetrievalError;
 import com.android.tradefed.build.IBuildInfo;
 import com.android.tradefed.build.IBuildProvider;
+import com.android.tradefed.build.IDeviceBuildInfo;
 import com.android.tradefed.build.IDeviceBuildProvider;
 import com.android.tradefed.command.CommandRunner.ExitCode;
 import com.android.tradefed.config.GlobalConfiguration;
@@ -59,13 +60,16 @@
 import com.android.tradefed.testtype.IRemoteTest;
 import com.android.tradefed.testtype.IResumableTest;
 import com.android.tradefed.testtype.IRetriableTest;
+import com.android.tradefed.util.FileUtil;
 import com.android.tradefed.util.IRunUtil;
 import com.android.tradefed.util.RunInterruptedException;
 import com.android.tradefed.util.RunUtil;
 import com.android.tradefed.util.StreamUtil;
+import com.android.tradefed.util.SystemUtil;
 
 import com.google.common.annotations.VisibleForTesting;
 
+import java.io.File;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -87,7 +91,7 @@
 public class TestInvocation implements ITestInvocation {
 
     /**
-     * Format of the key in {@link IInvocationContext} to log the battery level for each step of the
+     * Format of the key in {@link IBuildInfo} to log the battery level for each step of the
      * invocation. (Setup, test, tear down).
      */
     private static final String BATTERY_ATTRIBUTE_FORMAT_KEY = "%s-battery-%s";
@@ -207,6 +211,30 @@
             CLog.w("Using the test-tag from the build_provider. Consider updating your config to"
                     + " have no alias/namespace in front of test-tag.");
         }
+
+        // Load environment tests dir.
+        if (info instanceof IDeviceBuildInfo) {
+            File testsDir = ((IDeviceBuildInfo) info).getTestsDir();
+            if (testsDir != null && testsDir.exists()) {
+                for (File externalTestDir : getExternalTestCasesDirs()) {
+                    try {
+                        File subDir = new File(testsDir, externalTestDir.getName());
+                        FileUtil.recursiveSimlink(externalTestDir, subDir);
+                    } catch (IOException e) {
+                        CLog.e(
+                                "Failed to load external test dir %s. Ignoring it.",
+                                externalTestDir);
+                        CLog.e(e);
+                    }
+                }
+            }
+        }
+    }
+
+    /** Returns the list of external directories to Tradefed coming from the environment. */
+    @VisibleForTesting
+    List<File> getExternalTestCasesDirs() {
+        return SystemUtil.getExternalTestCasesDirs();
     }
 
     /**
@@ -295,6 +323,8 @@
         ITestDevice badDevice = null;
 
         startInvocation(config, context, listener);
+        // Ensure that no unexpected attributes are added afterward
+        ((InvocationContext) context).lockAttributes();
         try {
             logDeviceBatteryLevel(context, "initial");
             prepareAndRun(config, context, listener);
@@ -539,7 +569,9 @@
                 }
             }
             // Extra tear down step for the device
-            device.postInvocationTearDown();
+            if (!config.getCommandOptions().shouldSkipPreDeviceSetup()) {
+                device.postInvocationTearDown();
+            }
         }
 
         if (throwable != null) {
@@ -786,10 +818,13 @@
             try {
                 Integer batteryLevel = device.getBattery(500, TimeUnit.MILLISECONDS).get();
                 CLog.v("%s - %s - %d%%", BATT_TAG, event, batteryLevel);
-                context.addInvocationAttribute(
-                        String.format(
-                                BATTERY_ATTRIBUTE_FORMAT_KEY, testDevice.getSerialNumber(), event),
-                        batteryLevel.toString());
+                context.getBuildInfo(testDevice)
+                        .addBuildAttribute(
+                                String.format(
+                                        BATTERY_ATTRIBUTE_FORMAT_KEY,
+                                        testDevice.getSerialNumber(),
+                                        event),
+                                batteryLevel.toString());
                 continue;
             } catch (InterruptedException | ExecutionException e) {
                 // fall through
diff --git a/src/com/android/tradefed/invoker/shard/StrictShardHelper.java b/src/com/android/tradefed/invoker/shard/StrictShardHelper.java
index b0d179b..55e2041 100644
--- a/src/com/android/tradefed/invoker/shard/StrictShardHelper.java
+++ b/src/com/android/tradefed/invoker/shard/StrictShardHelper.java
@@ -22,9 +22,12 @@
 import com.android.tradefed.testtype.IBuildReceiver;
 import com.android.tradefed.testtype.IDeviceTest;
 import com.android.tradefed.testtype.IRemoteTest;
+import com.android.tradefed.testtype.IRuntimeHintProvider;
 import com.android.tradefed.testtype.IShardableTest;
 import com.android.tradefed.testtype.IStrictShardableTest;
 import com.android.tradefed.testtype.suite.ITestSuite;
+import com.android.tradefed.testtype.suite.ModuleMerger;
+import com.android.tradefed.util.TimeUtil;
 
 import java.util.ArrayList;
 import java.util.Collection;
@@ -53,7 +56,17 @@
             updateConfigIfSharded(config, shardCount, shardIndex);
         } else {
             List<IRemoteTest> listAllTests = getAllTests(config, shardCount, context);
-            config.setTests(splitTests(listAllTests, shardCount, shardIndex));
+            // We cannot shuffle to get better average results
+            normalizeDistribution(listAllTests, shardCount);
+            List<IRemoteTest> splitList;
+            if (shardCount == 1) {
+                // not sharded
+                splitList = listAllTests;
+            } else {
+                splitList = splitTests(listAllTests, shardCount, shardIndex);
+            }
+            aggregateSuiteModules(splitList);
+            config.setTests(splitList);
         }
         return false;
     }
@@ -135,19 +148,93 @@
      */
     private List<IRemoteTest> splitTests(
             List<IRemoteTest> fullList, int shardCount, int shardIndex) {
-        if (shardCount == 1) {
-            // Not sharded
-            return fullList;
+        List<List<IRemoteTest>> shards = new ArrayList<>();
+
+        // Generate all the shards
+        for (int i = 0; i < shardCount; i++) {
+            List<IRemoteTest> shardList;
+            if (i >= fullList.size()) {
+                // Return empty list when we don't have enough tests for all the shards.
+                shardList = new ArrayList<IRemoteTest>();
+                shards.add(shardList);
+                continue;
+            }
+            int numPerShard = (int) Math.ceil(fullList.size() / (float) shardCount);
+            if (i == shardCount - 1) {
+                // last shard take everything remaining.
+                shardList = fullList.subList(i * numPerShard, fullList.size());
+                shards.add(shardList);
+                continue;
+            }
+            shardList = fullList.subList(i * numPerShard, numPerShard + (i * numPerShard));
+            shards.add(shardList);
         }
-        if (shardIndex >= fullList.size()) {
-            // Return empty list when we don't have enough tests for all the shards.
-            return new ArrayList<IRemoteTest>();
+
+        // do last minute rebalancing
+        topBottom(shards);
+        return shards.get(shardIndex);
+    }
+
+    /**
+     * Move around predictably the tests in order to have a better uniformization of the tests in
+     * each shard.
+     */
+    private void normalizeDistribution(List<IRemoteTest> listAllTests, int shardCount) {
+        final int numRound = shardCount;
+        final int distance = shardCount + 1;
+        for (int i = 0; i < numRound; i++) {
+            for (int j = 0; j < listAllTests.size(); j = j + distance) {
+                // Push the test at the end
+                IRemoteTest push = listAllTests.remove(j);
+                listAllTests.add(push);
+            }
         }
-        int numPerShard = (int) Math.ceil(fullList.size() / (float) shardCount);
-        if (shardIndex == shardCount - 1) {
-            // last shard take everything remaining.
-            return fullList.subList(shardIndex * numPerShard, fullList.size());
+    }
+
+    /**
+     * Special handling for suite from {@link ITestSuite}. We aggregate the tests in the same shard
+     * in order to optimize target_preparation step.
+     *
+     * @param tests the {@link List} of {@link IRemoteTest} for that shard.
+     */
+    private void aggregateSuiteModules(List<IRemoteTest> tests) {
+        List<IRemoteTest> dupList = new ArrayList<>(tests);
+        for (int i = 0; i < dupList.size(); i++) {
+            if (dupList.get(i) instanceof ITestSuite) {
+                // We iterate the other tests to see if we can find another from the same module.
+                for (int j = i + 1; j < dupList.size(); j++) {
+                    // If the test was not already merged
+                    if (tests.contains(dupList.get(j))) {
+                        if (dupList.get(j) instanceof ITestSuite) {
+                            if (ModuleMerger.arePartOfSameSuite(
+                                    (ITestSuite) dupList.get(i), (ITestSuite) dupList.get(j))) {
+                                ModuleMerger.mergeSplittedITestSuite(
+                                        (ITestSuite) dupList.get(i), (ITestSuite) dupList.get(j));
+                                tests.remove(dupList.get(j));
+                            }
+                        }
+                    }
+                }
+            }
         }
-        return fullList.subList(shardIndex * numPerShard, numPerShard + (shardIndex * numPerShard));
+    }
+
+    private void topBottom(List<List<IRemoteTest>> allShards) {
+        // Generate approximate RuntimeHint for each shard
+        int index = 0;
+        CLog.e("============================");
+        for (List<IRemoteTest> shard : allShards) {
+            long aggTime = 0l;
+            CLog.e("++++++++++++++++++ SHARD %s +++++++++++++++", index);
+            for (IRemoteTest test : shard) {
+                if (test instanceof IRuntimeHintProvider) {
+                    aggTime += ((IRuntimeHintProvider) test).getRuntimeHint();
+                }
+            }
+            CLog.e("Shard %s approximate time: %s", index, TimeUtil.formatElapsedTime(aggTime));
+            index++;
+            CLog.e("+++++++++++++++++++++++++++++++++++++++++++");
+        }
+        CLog.e("============================");
     }
 }
diff --git a/src/com/android/tradefed/log/FileLogger.java b/src/com/android/tradefed/log/FileLogger.java
index 12d3e24..f9b13d9 100644
--- a/src/com/android/tradefed/log/FileLogger.java
+++ b/src/com/android/tradefed/log/FileLogger.java
@@ -197,7 +197,7 @@
             try {
                 // create a InputStream from log file
                 mLogStream.flush();
-                return new SnapshotInputStreamSource(mLogStream.getData());
+                return new SnapshotInputStreamSource("FileLogger", mLogStream.getData());
             } catch (IOException e) {
                 System.err.println("Failed to get log");
                 e.printStackTrace();
diff --git a/src/com/android/tradefed/log/LogReceiver.java b/src/com/android/tradefed/log/LogReceiver.java
index 832337a..658a1b9 100644
--- a/src/com/android/tradefed/log/LogReceiver.java
+++ b/src/com/android/tradefed/log/LogReceiver.java
@@ -1,3 +1,18 @@
+/*
+ * Copyright (C) 2016 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.log;
 
 import com.android.tradefed.device.BackgroundDeviceAction;
@@ -6,6 +21,7 @@
 import com.android.tradefed.result.ITestInvocationListener;
 import com.android.tradefed.result.InputStreamSource;
 import com.android.tradefed.result.LogDataType;
+import com.android.tradefed.util.StreamUtil;
 
 public class LogReceiver {
     private BackgroundDeviceAction mDeviceAction;
@@ -61,6 +77,11 @@
     }
 
     public void postLog(ITestInvocationListener listener) {
-        listener.testLog(getDescriptor(), LogDataType.TEXT, getData());
+        InputStreamSource stream = getData();
+        try {
+            listener.testLog(getDescriptor(), LogDataType.TEXT, getData());
+        } finally {
+            StreamUtil.cancel(stream);
+        }
     }
 }
\ No newline at end of file
diff --git a/src/com/android/tradefed/result/JUnit4ResultForwarder.java b/src/com/android/tradefed/result/JUnit4ResultForwarder.java
index d52cc35..07086e4 100644
--- a/src/com/android/tradefed/result/JUnit4ResultForwarder.java
+++ b/src/com/android/tradefed/result/JUnit4ResultForwarder.java
@@ -16,7 +16,10 @@
 package com.android.tradefed.result;
 
 import com.android.ddmlib.testrunner.TestIdentifier;
+import com.android.tradefed.testtype.DeviceJUnit4ClassRunner.LogAnnotation;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner.MetricAnnotation;
+import com.android.tradefed.testtype.MetricTestCase.LogHolder;
+import com.android.tradefed.util.StreamUtil;
 
 import org.junit.runner.Description;
 import org.junit.runner.notification.Failure;
@@ -72,6 +75,14 @@
                 if (a instanceof MetricAnnotation) {
                     metrics.putAll(((MetricAnnotation) a).mMetrics);
                 }
+                if (a instanceof LogAnnotation) {
+                    // Log all the logs found.
+                    for (LogHolder log : ((LogAnnotation) a).mLogs) {
+                        mListener.testLog(log.mDataName, log.mDataType, log.mDataStream);
+                        StreamUtil.cancel(log.mDataStream);
+                    }
+                    ((LogAnnotation) a).mLogs.clear();
+                }
             }
         }
         //description.
diff --git a/src/com/android/tradefed/result/JUnitToInvocationResultForwarder.java b/src/com/android/tradefed/result/JUnitToInvocationResultForwarder.java
index b2357d3..57d6f5c 100644
--- a/src/com/android/tradefed/result/JUnitToInvocationResultForwarder.java
+++ b/src/com/android/tradefed/result/JUnitToInvocationResultForwarder.java
@@ -92,6 +92,24 @@
         }
     }
 
+    /**
+     * Callback from JUnit3 forwarder in order to get the logs from a test.
+     *
+     * @param dataName a String descriptive name of the data. e.g. "device_logcat". Note dataName
+     *     may not be unique per invocation. ie implementers must be able to handle multiple calls
+     *     with same dataName
+     * @param dataType the LogDataType of the data
+     * @param dataStream the InputStreamSource of the data. Implementers should call
+     *     createInputStream to start reading the data, and ensure to close the resulting
+     *     InputStream when complete. Callers should ensure the source of the data remains present
+     *     and accessible until the testLog method completes.
+     */
+    public void testLog(String dataName, LogDataType dataType, InputStreamSource dataStream) {
+        for (ITestInvocationListener listener : mInvocationListeners) {
+            listener.testLog(dataName, dataType, dataStream);
+        }
+    }
+
     /** {@inheritDoc} */
     @Override
     public void startTest(Test test) {
diff --git a/src/com/android/tradefed/result/SnapshotInputStreamSource.java b/src/com/android/tradefed/result/SnapshotInputStreamSource.java
index 464d927..cc9f803 100644
--- a/src/com/android/tradefed/result/SnapshotInputStreamSource.java
+++ b/src/com/android/tradefed/result/SnapshotInputStreamSource.java
@@ -32,16 +32,14 @@
     private File mBackingFile;
     private boolean mIsCancelled = false;
 
-    /**
-     * Constructor for a file-backed {@link InputStreamSource}
-     */
-    public SnapshotInputStreamSource(InputStream stream) {
+    /** Constructor for a file-backed {@link InputStreamSource} */
+    public SnapshotInputStreamSource(String name, InputStream stream) {
         if (stream == null) {
             throw new NullPointerException();
         }
 
         try {
-            mBackingFile = createBackingFile(stream);
+            mBackingFile = createBackingFile(name, stream);
         } catch (IOException e) {
             // Log an error and invalidate ourself
             CLog.e("Received IOException while trying to wrap a stream");
@@ -52,11 +50,12 @@
 
     /**
      * Create the backing file and fill it with the contents of {@code stream}.
-     * <p />
-     * Exposed for unit testing
+     *
+     * <p>Exposed for unit testing
      */
-    File createBackingFile(InputStream stream) throws IOException {
-        File backingFile = FileUtil.createTempFile(this.getClass().getSimpleName() + "_", ".txt");
+    File createBackingFile(String name, InputStream stream) throws IOException {
+        File backingFile =
+                FileUtil.createTempFile(name + "_" + this.getClass().getSimpleName() + "_", ".txt");
         FileUtil.writeToFile(stream, backingFile);
         return backingFile;
     }
diff --git a/src/com/android/tradefed/result/suite/SuiteResultReporter.java b/src/com/android/tradefed/result/suite/SuiteResultReporter.java
index a236564..5b57a77 100644
--- a/src/com/android/tradefed/result/suite/SuiteResultReporter.java
+++ b/src/com/android/tradefed/result/suite/SuiteResultReporter.java
@@ -232,41 +232,50 @@
         }
     }
 
-    /**
-     * Print the collected times for Module preparation and tear Down. This only log to debug since
-     * it will be quite verbose for full run.
-     */
+    /** Print the collected times for Module preparation and tear Down. */
     private void printPreparationMetrics(Map<String, ModulePrepTimes> metrics) {
         if (metrics.isEmpty()) {
             return;
         }
-        CLog.d("============== Modules Preparation Times ==============");
+        CLog.logAndDisplay(
+                LogLevel.INFO, "============== Modules Preparation Times ==============");
         long totalPrep = 0l;
         long totalTear = 0l;
 
         for (String moduleName : metrics.keySet()) {
-            CLog.d("    %s => %s", moduleName, metrics.get(moduleName).toString());
+            CLog.logAndDisplay(
+                    LogLevel.INFO, "    %s => %s", moduleName, metrics.get(moduleName).toString());
             totalPrep += metrics.get(moduleName).mPrepTime;
             totalTear += metrics.get(moduleName).mTearDownTime;
         }
-        CLog.d(
+        CLog.logAndDisplay(
+                LogLevel.INFO,
                 "Total preparation time: %s  ||  Total tear down time: %s",
-                TimeUtil.formatElapsedTime(totalPrep), TimeUtil.formatElapsedTime(totalTear));
-        CLog.d("=======================================================");
+                TimeUtil.formatElapsedTime(totalPrep),
+                TimeUtil.formatElapsedTime(totalTear));
+        CLog.logAndDisplay(
+                LogLevel.INFO, "=======================================================");
     }
 
     private void printModuleCheckersMetric(List<TestRunResult> moduleCheckerResults) {
         if (moduleCheckerResults.isEmpty()) {
             return;
         }
-        CLog.d("============== Modules Checkers Times ==============");
+        CLog.logAndDisplay(LogLevel.INFO, "============== Modules Checkers Times ==============");
         long totalTime = 0l;
         for (TestRunResult t : moduleCheckerResults) {
-            CLog.d("    %s: %s", t.getName(), TimeUtil.formatElapsedTime(t.getElapsedTime()));
+            CLog.logAndDisplay(
+                    LogLevel.INFO,
+                    "    %s: %s",
+                    t.getName(),
+                    TimeUtil.formatElapsedTime(t.getElapsedTime()));
             totalTime += t.getElapsedTime();
         }
-        CLog.d("Total module checkers time: %s", TimeUtil.formatElapsedTime(totalTime));
-        CLog.d("====================================================");
+        CLog.logAndDisplay(
+                LogLevel.INFO,
+                "Total module checkers time: %s",
+                TimeUtil.formatElapsedTime(totalTime));
+        CLog.logAndDisplay(LogLevel.INFO, "====================================================");
     }
 
     public int getTotalModules() {
diff --git a/src/com/android/tradefed/targetprep/DeviceFlashPreparer.java b/src/com/android/tradefed/targetprep/DeviceFlashPreparer.java
index 96f4569..f691e70 100644
--- a/src/com/android/tradefed/targetprep/DeviceFlashPreparer.java
+++ b/src/com/android/tradefed/targetprep/DeviceFlashPreparer.java
@@ -247,6 +247,7 @@
         getRunUtil().allowInterrupt(false);
         try {
             IDeviceBuildInfo deviceBuild = (IDeviceBuildInfo)buildInfo;
+            checkDeviceProductType(device, deviceBuild);
             device.setRecoveryMode(RecoveryMode.ONLINE);
             IDeviceFlasher flasher = createFlasher(device);
             flasher.setWipeTimeout(mWipeTimeout);
@@ -303,6 +304,19 @@
     }
 
     /**
+     * Possible check before flashing to ensure the device is as expected compare to the build info.
+     *
+     * @param device the {@link ITestDevice} to flash.
+     * @param deviceBuild the {@link IDeviceBuildInfo} used to flash.
+     * @throws BuildError
+     * @throws DeviceNotAvailableException
+     */
+    protected void checkDeviceProductType(ITestDevice device, IDeviceBuildInfo deviceBuild)
+            throws BuildError, DeviceNotAvailableException {
+        // empty of purpose
+    }
+
+    /**
      * Verifies the expected build matches the actual build on device after flashing
      * @throws DeviceNotAvailableException
      */
diff --git a/src/com/android/tradefed/targetprep/InstallAllTestZipAppsSetup.java b/src/com/android/tradefed/targetprep/InstallAllTestZipAppsSetup.java
new file mode 100644
index 0000000..aa7894b
--- /dev/null
+++ b/src/com/android/tradefed/targetprep/InstallAllTestZipAppsSetup.java
@@ -0,0 +1,258 @@
+/*
+ * Copyright (C) 2017 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.targetprep;
+
+import com.android.tradefed.build.IBuildInfo;
+import com.android.tradefed.config.Option;
+import com.android.tradefed.config.OptionClass;
+import com.android.tradefed.device.DeviceNotAvailableException;
+import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.log.LogUtil.CLog;
+import com.android.tradefed.util.AaptParser;
+import com.android.tradefed.util.FileUtil;
+import com.android.tradefed.util.ZipUtil2;
+
+import org.apache.commons.compress.archivers.zip.ZipFile;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+
+/**
+ * A {@link ITargetPreparer} that installs all apps in a test zip. For individual test app install
+ * please look at {@link TestAppInstallSetup}.
+ */
+@OptionClass(alias = "all-tests-zip-installer")
+public class InstallAllTestZipAppsSetup implements ITargetCleaner {
+    @Option(
+        name = "install-arg",
+        description =
+                "Additional arguments to be passed to install command, "
+                        + "including leading dash, e.g. \"-d\""
+    )
+    private Collection<String> mInstallArgs = new ArrayList<>();
+
+    @Option(
+        name = "cleanup-apks",
+        description =
+                "Whether apks installed should be uninstalled after test. Note that the "
+                        + "preparer does not verify if the apks are successfully removed."
+    )
+    private boolean mCleanup = true;
+
+    @Option(
+        name = "stop-install-on-failure",
+        description =
+                "Whether to stop the preparer by throwing an exception or only log the "
+                        + "error on continue."
+    )
+    private boolean mStopInstallOnFailure = true;
+
+    @Option(name = "test-zip-name", description = "File name for test zip containing APKs.")
+    private String mTestZipName;
+
+    @Option(name = "disable", description = "Disable this target preparer.")
+    private boolean mDisable;
+
+    List<String> mPackagesInstalled = new ArrayList<>();
+
+    public void setTestZipName(String testZipName) {
+        mTestZipName = testZipName;
+    }
+
+    public void setStopInstallOnFailure(boolean stopInstallOnFailure) {
+        mStopInstallOnFailure = stopInstallOnFailure;
+    }
+
+    public void setDisable(boolean disable) {
+        mDisable = disable;
+    }
+
+    public void setCleanup(boolean cleanup) {
+        mCleanup = cleanup;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void setUp(ITestDevice device, IBuildInfo buildInfo)
+            throws TargetSetupError, DeviceNotAvailableException {
+        if (mDisable) {
+            CLog.d("InstallAllTestZipAppsSetup disabled, skipping setUp");
+            return;
+        }
+        File testsZip = getZipFile(device, buildInfo);
+        // Locate test dir where the test zip file was unzip to.
+        if (testsZip == null) {
+            throw new TargetSetupError(
+                    "Failed to find a valid test zip directory.", device.getDeviceDescriptor());
+        }
+        File testsDir;
+        try {
+            testsDir = extractZip(testsZip);
+        } catch (IOException e) {
+            throw new TargetSetupError(
+                    "Failed to extract test zip.", e, device.getDeviceDescriptor());
+        }
+
+        try {
+            installApksRecursively(testsDir, device);
+        } finally {
+            FileUtil.recursiveDelete(testsDir);
+        }
+    }
+
+    /**
+     * Returns the zip file.
+     *
+     * @param buildInfo {@link IBuildInfo} containing files.
+     * @return the {@link File} for the zip file.
+     */
+    File getZipFile(ITestDevice device, IBuildInfo buildInfo) throws TargetSetupError {
+        if (mTestZipName == null) {
+            throw new TargetSetupError("test-zip-name is null.", device.getDeviceDescriptor());
+        }
+        return buildInfo.getFile(mTestZipName);
+    }
+
+    /**
+     * Install all apks found in a given directory.
+     *
+     * @param directory {@link File} directory to install from.
+     * @param device {@link ITestDevice} to install all apks to.
+     * @throws TargetSetupError
+     * @throws DeviceNotAvailableException
+     */
+    void installApksRecursively(File directory, ITestDevice device)
+            throws TargetSetupError, DeviceNotAvailableException {
+        if (directory == null || !directory.isDirectory()) {
+            throw new TargetSetupError("Invalid test directory!", device.getDeviceDescriptor());
+        }
+        CLog.d("Installing all apks found in dir %s ...", directory.getAbsolutePath());
+        File[] files = directory.listFiles();
+        for (File f : files) {
+            if (f.isDirectory()) {
+                installApksRecursively(f, device);
+            } else if (FileUtil.getExtension(f.getAbsolutePath()).toLowerCase().equals(".apk")) {
+                installApk(f, device);
+            } else {
+                CLog.d("Skipping %s because it is not an apk", f.getAbsolutePath());
+            }
+        }
+    }
+
+    /**
+     * Extract the given zip file to a local dir.
+     *
+     * <p>Exposed so unit tests can mock
+     *
+     * @param testsZip
+     * @return the {@link File} referencing the zip output.
+     * @throws IOException
+     */
+    File extractZip(File testsZip) throws IOException {
+        File testsDir = null;
+        try (ZipFile zip = new ZipFile(testsZip)) {
+            testsDir = FileUtil.createTempDir("tests-zip_");
+            ZipUtil2.extractZip(zip, testsDir);
+        } catch (IOException e) {
+            FileUtil.recursiveDelete(testsDir);
+            throw e;
+        }
+        return testsDir;
+    }
+
+    /**
+     * Installs a single app to the device.
+     *
+     * @param appFile {@link File} of the apk to install.
+     * @param device {@link ITestDevice} to install the apk to.
+     * @throws TargetSetupError
+     * @throws DeviceNotAvailableException
+     */
+    void installApk(File appFile, ITestDevice device)
+            throws TargetSetupError, DeviceNotAvailableException {
+        CLog.d("Installing apk from %s ...", appFile.getAbsolutePath());
+        String result = device.installPackage(appFile, true, mInstallArgs.toArray(new String[] {}));
+        if (result == null) {
+            // only consider cleanup if install was successful
+            if (mCleanup) {
+                addApkToInstalledList(appFile, device);
+            }
+        } else if (mStopInstallOnFailure) {
+            // if flag is true, we stop the sequence for an exception.
+            throw new TargetSetupError(
+                    String.format(
+                            "Failed to install %s on %s. Reason: '%s'",
+                            appFile, device.getSerialNumber(), result),
+                    device.getDeviceDescriptor());
+        } else {
+            CLog.e(
+                    "Failed to install %s on %s. Reason: '%s'",
+                    appFile, device.getSerialNumber(), result);
+        }
+    }
+
+    /**
+     * Adds an app to the list of apps to uninstall in tearDown() if necessary.
+     *
+     * @param appFile {@link File} of apk.
+     * @param device {@link ITestDevice} apk was installed on.
+     * @throws TargetSetupError
+     */
+    void addApkToInstalledList(File appFile, ITestDevice device) throws TargetSetupError {
+        String packageName = getAppPackageName(appFile);
+        if (packageName == null) {
+            throw new TargetSetupError(
+                    "apk installed but AaptParser failed", device.getDeviceDescriptor());
+        }
+        mPackagesInstalled.add(packageName);
+    }
+
+    /**
+     * Returns the package name for a an apk. Returns null if any errors.
+     *
+     * @param appFile {@link File} of apk.
+     * @return Package name of appFile, null if errors.
+     */
+    String getAppPackageName(File appFile) {
+        AaptParser parser = AaptParser.parse(appFile);
+        if (parser == null) {
+            return null;
+        }
+        return parser.getPackageName();
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void tearDown(ITestDevice device, IBuildInfo buildInfo, Throwable e)
+            throws DeviceNotAvailableException {
+        if (mDisable) {
+            CLog.d("InstallAllTestZipAppsSetup disabled, skipping tearDown");
+            return;
+        }
+        if (mCleanup && !(e instanceof DeviceNotAvailableException)) {
+            for (String packageName : mPackagesInstalled) {
+                String msg = device.uninstallPackage(packageName);
+                if (msg != null) {
+                    CLog.w(String.format("error uninstalling package '%s': %s", packageName, msg));
+                }
+            }
+        }
+    }
+}
diff --git a/src/com/android/tradefed/targetprep/PushFilePreparer.java b/src/com/android/tradefed/targetprep/PushFilePreparer.java
index 75c28a4..7247a79 100644
--- a/src/com/android/tradefed/targetprep/PushFilePreparer.java
+++ b/src/com/android/tradefed/targetprep/PushFilePreparer.java
@@ -25,13 +25,11 @@
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.util.FileUtil;
-import com.android.tradefed.util.SystemUtil;
 
 import java.io.File;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
-import java.util.List;
 
 /**
  * A {@link ITargetPreparer} that attempts to push any number of files from any host path to any
@@ -45,8 +43,6 @@
     private static final String MEDIA_SCAN_INTENT =
             "am broadcast -a android.intent.action.MEDIA_MOUNTED -d file://%s "
                     + "--receiver-include-background";
-    private static final String HOST_TESTCASES = "host/testcases";
-    private static final String TARGET_TESTCASES = "target/testcases";
 
     @Option(name="push", description=
             "A push-spec, formatted as '/path/to/srcfile.txt->/path/to/destfile.txt' or " +
@@ -92,18 +88,6 @@
     }
 
     /**
-     * Get a list of {@link File} of the test cases directories
-     *
-     * <p>The wrapper function is for unit test to mock the system calls.
-     *
-     * @return a list of {@link File} of directories of the test cases folder of build output, based
-     *     on the value of environment variables.
-     */
-    List<File> getTestCasesDirs() {
-        return SystemUtil.getTestCasesDirs();
-    }
-
-    /**
      * Resolve relative file path via {@link IBuildInfo} and test cases directories.
      *
      * @param buildInfo the build artifact information
@@ -118,26 +102,10 @@
                 return src;
             }
         }
-        List<File> testCasesDirs = getTestCasesDirs();
 
-        // Search for source file in tests directory if buildInfo is IDeviceBuildInfo.
         if (buildInfo instanceof IDeviceBuildInfo) {
-            IDeviceBuildInfo deviceBuildInfo = (IDeviceBuildInfo) buildInfo;
-            File testsDir = deviceBuildInfo.getTestsDir();
-            // Add all possible paths to the testcases directory list.
-            if (testsDir != null) {
-                testCasesDirs.addAll(
-                        Arrays.asList(
-                                testsDir,
-                                FileUtil.getFileForPath(testsDir, HOST_TESTCASES),
-                                FileUtil.getFileForPath(testsDir, TARGET_TESTCASES)));
-            }
-        }
-        for (File dir : testCasesDirs) {
-            src = FileUtil.getFileForPath(dir, fileName);
-            if (src != null && src.exists()) {
-                return src;
-            }
+            File testsDir = ((IDeviceBuildInfo) buildInfo).getTestsDir();
+            return FileUtil.findFile(testsDir, fileName);
         }
         return null;
     }
diff --git a/src/com/android/tradefed/targetprep/TestAppInstallSetup.java b/src/com/android/tradefed/targetprep/TestAppInstallSetup.java
index eab1b5d..ce94055 100644
--- a/src/com/android/tradefed/targetprep/TestAppInstallSetup.java
+++ b/src/com/android/tradefed/targetprep/TestAppInstallSetup.java
@@ -29,7 +29,6 @@
 import com.android.tradefed.util.AaptParser;
 import com.android.tradefed.util.AbiFormatter;
 import com.android.tradefed.util.BuildTestsZipUtils;
-import com.android.tradefed.util.SystemUtil;
 
 import java.io.File;
 import java.io.IOException;
@@ -131,18 +130,6 @@
         }
     }
 
-    /**
-     * Get a list of {@link File} of the test cases directories
-     *
-     * <p>The wrapper function is for unit test to mock the system calls.
-     *
-     * @return a list of {@link File} of directories of the test cases folder of build output, based
-     *     on the value of environment variables.
-     */
-    List<File> getTestCasesDirs() {
-        return SystemUtil.getTestCasesDirs();
-    }
-
     /** {@inheritDoc} */
     @Override
     public void setUp(ITestDevice device, IBuildInfo buildInfo)
@@ -155,11 +142,6 @@
             mPackagesInstalled = new ArrayList<>();
         }
 
-        // Force to look for apk files in build ouput's test cases directory.
-        for (File testCasesDir : getTestCasesDirs()) {
-            setAltDir(testCasesDir);
-        }
-
         for (String testAppName : mTestFileNames) {
             if (testAppName == null || testAppName.trim().isEmpty()) {
                 continue;
diff --git a/src/com/android/tradefed/targetprep/TestFilePushSetup.java b/src/com/android/tradefed/targetprep/TestFilePushSetup.java
index 7763b78..4e19756 100644
--- a/src/com/android/tradefed/targetprep/TestFilePushSetup.java
+++ b/src/com/android/tradefed/targetprep/TestFilePushSetup.java
@@ -156,6 +156,10 @@
                             "Could not find test file %s directory in extracted tests.zip",
                             fileName), device.getDeviceDescriptor());
                 } else {
+                    CLog.w(String.format(
+                            "Could not find test file %s directory in extracted tests.zip, but" +
+                            "will continue test setup as throw-if-not-found is set to false",
+                            fileName));
                     continue;
                 }
             }
@@ -170,7 +174,7 @@
             device.executeShellCommand(String.format("chown system.system %s", remoteFileName));
             filePushed++;
         }
-        if (filePushed == 0) {
+        if (filePushed == 0 && mThrowIfNoFile) {
             throw new TargetSetupError("No file is pushed from tests.zip",
                     device.getDeviceDescriptor());
         }
diff --git a/src/com/android/tradefed/testtype/Abi.java b/src/com/android/tradefed/testtype/Abi.java
index 2b0c5bb..1a24df5 100644
--- a/src/com/android/tradefed/testtype/Abi.java
+++ b/src/com/android/tradefed/testtype/Abi.java
@@ -44,4 +44,8 @@
         return mBitness;
     }
 
+    @Override
+    public String toString() {
+        return "{" + mName + ", bitness=" + mBitness + "}";
+    }
 }
\ No newline at end of file
diff --git a/src/com/android/tradefed/testtype/AndroidJUnitTest.java b/src/com/android/tradefed/testtype/AndroidJUnitTest.java
index 2af7fe9..d44f57c 100644
--- a/src/com/android/tradefed/testtype/AndroidJUnitTest.java
+++ b/src/com/android/tradefed/testtype/AndroidJUnitTest.java
@@ -32,6 +32,7 @@
 import java.io.File;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 import java.util.Set;
 
@@ -104,6 +105,12 @@
             description="The device directory path to which the test filtering files are pushed")
     private String mTestFilterDir = "/data/local/tmp/ajur";
 
+    @Option(
+        name = "ajur-max-shard",
+        description = "The maximum number of shard we want to allow the test to shard into"
+    )
+    private Integer mMaxShard = null;
+
     private String mDeviceIncludeFile = null;
     private String mDeviceExcludeFile = null;
     private int mTotalShards = 0;
@@ -217,13 +224,13 @@
         // if mIncludeTestFile is set, perform filtering with this file
         if (mIncludeTestFile != null) {
             mDeviceIncludeFile = mTestFilterDir.replaceAll("/$", "") + "/" + INCLUDE_FILE;
-            pushTestFile(mIncludeTestFile, mDeviceIncludeFile);
+            pushTestFile(mIncludeTestFile, mDeviceIncludeFile, listener);
         }
 
         // if mExcludeTestFile is set, perform filtering with this file
         if (mExcludeTestFile != null) {
             mDeviceExcludeFile = mTestFilterDir.replaceAll("/$", "") + "/" + EXCLUDE_FILE;
-            pushTestFile(mExcludeTestFile, mDeviceExcludeFile);
+            pushTestFile(mExcludeTestFile, mDeviceExcludeFile, listener);
         }
         if (mTotalShards > 0 && !isShardable() && mShardIndex != 0) {
             // If not shardable, only first shard can run.
@@ -305,21 +312,36 @@
         }
     }
 
-    /*
+    /**
+     * Push the testFile to the requested destination. This should only be called for a non-null
+     * testFile
+     *
      * @param testFile file to be pushed from the host to the device.
      * @param destination the path on the device to which testFile is pushed
-     * This should only be called for a non-null testFile
+     * @param listener {@link ITestInvocationListener} to report failures.
      */
-    private void pushTestFile(File testFile, String destination) throws DeviceNotAvailableException {
+    private void pushTestFile(File testFile, String destination, ITestInvocationListener listener)
+            throws DeviceNotAvailableException {
         if (!testFile.canRead() || !testFile.isFile()) {
-            throw new IllegalArgumentException(
-                    String.format("Cannot read test file %s", testFile.getAbsolutePath()));
+            String message = String.format("Cannot read test file %s", testFile.getAbsolutePath());
+            reportEarlyFailure(listener, message);
+            throw new IllegalArgumentException(message);
         }
         ITestDevice device = getDevice();
-        if (!device.pushFile(testFile, destination)) {
-            throw new RuntimeException(String.format("Failed to push file %s to %s for %s "
-                    + "in pushTestFile", testFile.getAbsolutePath(), destination,
-                    device.getSerialNumber()));
+        try {
+            if (!device.pushFile(testFile, destination)) {
+                String message =
+                        String.format(
+                                "Failed to push file %s to %s for %s in pushTestFile",
+                                testFile.getAbsolutePath(), destination, device.getSerialNumber());
+                reportEarlyFailure(listener, message);
+                throw new RuntimeException(message);
+            }
+            // in case the folder was created as 'root' we make is usable.
+            device.executeShellCommand(String.format("chown -R shell:shell %s", mTestFilterDir));
+        } catch (DeviceNotAvailableException e) {
+            reportEarlyFailure(listener, e.getMessage());
+            throw e;
         }
     }
 
@@ -328,6 +350,12 @@
         device.executeShellCommand(String.format("rm %s", deviceTestFile));
     }
 
+    private void reportEarlyFailure(ITestInvocationListener listener, String errorMessage) {
+        listener.testRunStarted("AndroidJUnitTest_setupError", 0);
+        listener.testRunFailed(errorMessage);
+        listener.testRunEnded(0, Collections.emptyMap());
+    }
+
     /**
      * Return if a string is the name of a Class or a Method.
      */
@@ -360,6 +388,9 @@
         if (!isShardable()) {
             return null;
         }
+        if (mMaxShard != null) {
+            shardCount = Math.min(shardCount, mMaxShard);
+        }
         if (!mIsSharded && shardCount > 1) {
             mIsSharded = true;
             Collection<IRemoteTest> shards = new ArrayList<>(shardCount);
diff --git a/src/com/android/tradefed/testtype/DeviceJUnit4ClassRunner.java b/src/com/android/tradefed/testtype/DeviceJUnit4ClassRunner.java
index 2f1c7f3..68d32b3 100644
--- a/src/com/android/tradefed/testtype/DeviceJUnit4ClassRunner.java
+++ b/src/com/android/tradefed/testtype/DeviceJUnit4ClassRunner.java
@@ -16,7 +16,11 @@
 package com.android.tradefed.testtype;
 
 import com.android.tradefed.build.IBuildInfo;
+import com.android.tradefed.config.Option;
 import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.result.InputStreamSource;
+import com.android.tradefed.result.LogDataType;
+import com.android.tradefed.testtype.MetricTestCase.LogHolder;
 
 import org.junit.rules.ExternalResource;
 import org.junit.rules.TestRule;
@@ -26,19 +30,24 @@
 import org.junit.runners.model.Statement;
 
 import java.lang.annotation.Annotation;
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 /**
- * JUnit4 test runner that also accommodate {@link IDeviceTest}.
- * Should be specify above JUnit4 Test with the RunWith annotation.
+ * JUnit4 test runner that also accommodate {@link IDeviceTest}. Should be specify above JUnit4 Test
+ * with the RunWith annotation.
  */
-public class DeviceJUnit4ClassRunner extends BlockJUnit4ClassRunner implements IDeviceTest,
-        IBuildReceiver, IAbiReceiver {
+public class DeviceJUnit4ClassRunner extends BlockJUnit4ClassRunner
+        implements IDeviceTest, IBuildReceiver, IAbiReceiver, ISetOptionReceiver {
     private ITestDevice mDevice;
     private IBuildInfo mBuildInfo;
     private IAbi mAbi;
 
+    @Option(name = HostTest.SET_OPTION_NAME, description = HostTest.SET_OPTION_DESC)
+    private List<String> mKeyValueOptions = new ArrayList<>();
+
     public DeviceJUnit4ClassRunner(Class<?> klass) throws InitializationError {
         super(klass);
     }
@@ -65,6 +74,8 @@
         if (testObj instanceof IAbiReceiver) {
             ((IAbiReceiver) testObj).setAbi(mAbi);
         }
+        // Set options of test object
+        HostTest.setOptionToLoadedObject(testObj, mKeyValueOptions);
         return testObj;
     }
 
@@ -166,4 +177,64 @@
             return null;
         }
     }
+
+    /**
+     * Implementation of {@link ExternalResource} and {@link TestRule}. This rule allows to log logs
+     * during a test case (inside @Test). It guarantees that the log list is cleaned between tests,
+     * so the same rule object can be re-used.
+     *
+     * <pre>Example:
+     * &#064;Rule
+     * public TestLogData logs = new TestLogData();
+     *
+     * &#064;Test
+     * public void testFoo() {
+     *     logs.addTestLog("logcat", LogDataType.LOGCAT, new FileInputStreamSource(logcatFile));
+     * }
+     *
+     * &#064;Test
+     * public void testFoo2() {
+     *     logs.addTestLog("logcat2", LogDataType.LOGCAT, new FileInputStreamSource(logcatFile2));
+     * }
+     * </pre>
+     */
+    public static class TestLogData extends ExternalResource {
+        private Description mDescription;
+        private List<LogHolder> mLogs = new ArrayList<>();
+
+        @Override
+        public Statement apply(Statement base, Description description) {
+            mDescription = description;
+            return super.apply(base, description);
+        }
+
+        public final void addTestLog(
+                String dataName, LogDataType dataType, InputStreamSource dataStream) {
+            mLogs.add(new LogHolder(dataName, dataType, dataStream));
+        }
+
+        @Override
+        protected void after() {
+            // we inject a Description with an annotation carrying metrics.
+            // We have to go around, since Description cannot be extended and RunNotifier
+            // does not give us a lot of flexibility to find our metrics back.
+            mDescription.addChild(
+                    Description.createTestDescription("LOGS", "LOGS", new LogAnnotation(mLogs)));
+        }
+    }
+
+    /** Fake annotation meant to carry logs to the reporters. */
+    public static class LogAnnotation implements Annotation {
+
+        public List<LogHolder> mLogs = new ArrayList<>();
+
+        public LogAnnotation(List<LogHolder> logs) {
+            mLogs.addAll(logs);
+        }
+
+        @Override
+        public Class<? extends Annotation> annotationType() {
+            return null;
+        }
+    }
 }
diff --git a/src/com/android/tradefed/testtype/DeviceTestResult.java b/src/com/android/tradefed/testtype/DeviceTestResult.java
index 48840ac..7f57210 100644
--- a/src/com/android/tradefed/testtype/DeviceTestResult.java
+++ b/src/com/android/tradefed/testtype/DeviceTestResult.java
@@ -17,6 +17,8 @@
 
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.result.JUnitToInvocationResultForwarder;
+import com.android.tradefed.testtype.MetricTestCase.LogHolder;
+import com.android.tradefed.util.StreamUtil;
 
 import junit.framework.AssertionFailedError;
 import junit.framework.Protectable;
@@ -100,10 +102,24 @@
     public void endTest(Test test) {
         Map<String, String> metrics = new HashMap<>();
         if (test instanceof MetricTestCase) {
-            metrics.putAll(((MetricTestCase) test).mMetrics);
+            MetricTestCase metricTest = (MetricTestCase) test;
+            metrics.putAll(metricTest.mMetrics);
             // reset the metric for next test.
-            ((MetricTestCase) test).mMetrics = new HashMap<String, String>();
+            metricTest.mMetrics = new HashMap<String, String>();
+
+            // testLog the log files
+            for (TestListener each : cloneListeners()) {
+                for (LogHolder log : metricTest.mLogs) {
+                    if (each instanceof JUnitToInvocationResultForwarder) {
+                        ((JUnitToInvocationResultForwarder) each)
+                                .testLog(log.mDataName, log.mDataType, log.mDataStream);
+                    }
+                    StreamUtil.cancel(log.mDataStream);
+                }
+            }
+            metricTest.mLogs.clear();
         }
+
         for (TestListener each : cloneListeners()) {
             // when possible pass the metrics collected from the tests to our reporters.
             if (!metrics.isEmpty() && each instanceof JUnitToInvocationResultForwarder) {
diff --git a/src/com/android/tradefed/testtype/GTest.java b/src/com/android/tradefed/testtype/GTest.java
index 06ae015..dcff95b 100644
--- a/src/com/android/tradefed/testtype/GTest.java
+++ b/src/com/android/tradefed/testtype/GTest.java
@@ -105,6 +105,13 @@
             description = "adb shell command(s) to run before GTest.")
     private List<String> mBeforeTestCmd = new ArrayList<>();
 
+
+    @Option(
+        name = "reboot-before-test",
+        description = "Reboot the device before the test suite starts."
+    )
+    private boolean mRebootBeforeTest = false;
+
     @Option(name = "after-test-cmd",
             description = "adb shell command(s) to run after GTest.")
     private List<String> mAfterTestCmd = new ArrayList<>();
@@ -534,6 +541,12 @@
             for (String cmd : mBeforeTestCmd) {
                 testDevice.executeShellCommand(cmd);
             }
+
+            if (mRebootBeforeTest) {
+                CLog.d("Rebooting device before test starts as requested.");
+                testDevice.reboot();
+            }
+
             String cmd = getGTestCmdLine(fullPath, flags);
             // ensure that command is not too long for adb
             if (cmd.length() < GTEST_CMD_CHAR_LIMIT) {
diff --git a/src/com/android/tradefed/testtype/HostTest.java b/src/com/android/tradefed/testtype/HostTest.java
index 249332a..ae2cf57 100644
--- a/src/com/android/tradefed/testtype/HostTest.java
+++ b/src/com/android/tradefed/testtype/HostTest.java
@@ -48,9 +48,11 @@
 import java.lang.reflect.AnnotatedElement;
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
+import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.Deque;
 import java.util.Enumeration;
 import java.util.HashSet;
 import java.util.LinkedHashSet;
@@ -75,6 +77,7 @@
                 IStrictShardableTest,
                 IRuntimeHintProvider {
 
+
     @Option(name = "class", description = "The JUnit test classes to run, in the format "
             + "<package>.<class>. eg. \"com.android.foo.Bar\". This field can be repeated.",
             importance = Importance.IF_UNSET)
@@ -85,10 +88,15 @@
             importance = Importance.IF_UNSET)
     private String mMethodName;
 
-    @Option(name = "set-option", description = "Options to be passed down to the class "
-            + "under test, key and value should be separated by colon \":\"; for example, if class "
-            + "under test supports \"--iteration 1\" from a command line, it should be passed in as"
-            + " \"--set-option iteration:1\"; escaping of \":\" is currently not supported")
+    public static final String SET_OPTION_NAME = "set-option";
+    public static final String SET_OPTION_DESC =
+            "Options to be passed down to the class under test, key and value should be "
+                    + "separated by colon \":\"; for example, if class under test supports "
+                    + "\"--iteration 1\" from a command line, it should be passed in as"
+                    + " \"--set-option iteration:1\" or \"--set-option iteration:key=value\" for "
+                    + "passing options to map; escaping of \":\" \"=\" is currently not supported";
+
+    @Option(name = SET_OPTION_NAME, description = SET_OPTION_DESC)
     private List<String> mKeyValueOptions = new ArrayList<>();
 
     @Option(name = "include-annotation",
@@ -113,12 +121,23 @@
     )
     private long mRuntimeHint = 60000; // 1 minute
 
+    enum ShardUnit {
+        CLASS, METHOD;
+    }
+
+    @Option(name = "shard-unit",
+            description = "Shard by class or method")
+    private ShardUnit mShardUnit = ShardUnit.CLASS;
+
     private ITestDevice mDevice;
     private IBuildInfo mBuildInfo;
     private IAbi mAbi;
     private TestFilterHelper mFilterHelper;
     private boolean mSkipTestClassCheck = false;
 
+    private List<Object> mTestMethods;
+    private int mNumTestCases = -1;
+
     private static final String EXCLUDE_NO_TEST_FAILURE = "org.junit.runner.manipulation.Filter";
     private static final String TEST_FULL_NAME_FORMAT = "%s#%s";
 
@@ -179,6 +198,13 @@
     }
 
     /**
+     * @return true if shard-unit is method; false otherwise
+     */
+    private boolean shardUnitIsMethod() {
+        return ShardUnit.METHOD.equals(mShardUnit);
+    }
+
+    /**
      * {@inheritDoc}
      */
     @Override
@@ -214,6 +240,11 @@
      * Return the number of test cases across all classes part of the tests
      */
     public int countTestCases() {
+        if (mTestMethods != null) {
+            return mTestMethods.size();
+        } else if (mNumTestCases >= 0) {
+            return mNumTestCases;
+        }
         // Ensure filters are set in the helper
         mFilterHelper.addAllIncludeAnnotation(mIncludeAnnotations);
         mFilterHelper.addAllExcludeAnnotation(mExcludeAnnotations);
@@ -253,7 +284,7 @@
                 count++;
             }
         }
-        return count;
+        return mNumTestCases = count;
     }
 
     /**
@@ -329,6 +360,17 @@
         if (testObj instanceof IAbiReceiver) {
             ((IAbiReceiver)testObj).setAbi(mAbi);
         }
+        // managed runner should have the same set-option to pass option too.
+        if (testObj instanceof ISetOptionReceiver) {
+            try {
+                OptionSetter setter = new OptionSetter(testObj);
+                for (String item : mKeyValueOptions) {
+                    setter.setOptionValue(SET_OPTION_NAME, item);
+                }
+            } catch (ConfigurationException e) {
+                throw new RuntimeException(e);
+            }
+        }
     }
 
     /**
@@ -349,49 +391,23 @@
         if (mMethodName != null && classes.size() > 1) {
             throw new IllegalArgumentException("Method name given with multiple test classes");
         }
-        for (Class<?> classObj : classes) {
+        if (mTestMethods != null) {
+            runTestCases(listener);
+        } else {
+            runTestClasses(listener);
+        }
+    }
+
+    private void runTestClasses(ITestInvocationListener listener) throws DeviceNotAvailableException {
+        for (Class<?> classObj : getClasses()) {
             if (IRemoteTest.class.isAssignableFrom(classObj)) {
                 IRemoteTest test = (IRemoteTest) loadObject(classObj);
                 applyFilters(classObj, test);
-                if (mCollectTestsOnly) {
-                    // Collect only mode is propagated to the test.
-                    if (test instanceof ITestCollector) {
-                        ((ITestCollector) test).setCollectTestsOnly(true);
-                    } else {
-                        throw new IllegalArgumentException(
-                                String.format(
-                                        "%s does not implement ITestCollector", test.getClass()));
-                    }
-                }
-                test.run(listener);
+                runRemoteTest(listener, test);
             } else if (Test.class.isAssignableFrom(classObj)) {
-                if (mCollectTestsOnly) {
-                    // Collect only mode, fake the junit test execution.
-                    TestSuite junitTest = collectTests(collectClasses(classObj));
-                    listener.testRunStarted(classObj.getName(), junitTest.countTestCases());
-                    Map<String, String> empty = Collections.emptyMap();
-                    for (int i = 0; i < junitTest.countTestCases(); i++) {
-                        Test t = junitTest.testAt(i);
-                        // Test does not have a getName method.
-                        // using the toString format instead: <testName>(className)
-                        String testName = t.toString().split("\\(")[0];
-                        TestIdentifier testId =
-                                new TestIdentifier(t.getClass().getName(), testName);
-                        listener.testStarted(testId);
-                        listener.testEnded(testId, empty);
-                    }
-                    Map<String, String> emptyMap = Collections.emptyMap();
-                    listener.testRunEnded(0, emptyMap);
-                } else {
-                    JUnitRunUtil.runTest(listener, collectTests(collectClasses(classObj)),
-                            classObj.getName());
-                }
+                TestSuite junitTest = collectTests(collectClasses(classObj));
+                runJUnit3Tests(listener, junitTest, classObj.getName());
             } else if (hasJUnit4Annotation(classObj)) {
-                // Running in a full JUnit4 manner, no downgrade to JUnit3 {@link Test}
-                JUnitCore runnerCore = new JUnitCore();
-                JUnit4ResultForwarder list = new JUnit4ResultForwarder(listener);
-                runnerCore.addListener(list);
-                Request req = Request.aClass(classObj);
                 // Include the method name filtering
                 Set<String> includes = mFilterHelper.getIncludeFilters();
                 if (mMethodName != null) {
@@ -399,36 +415,11 @@
                             mMethodName));
                 }
 
+                // Running in a full JUnit4 manner, no downgrade to JUnit3 {@link Test}
+                Request req = Request.aClass(classObj);
                 req = req.filterWith(new JUnit4TestFilter(mFilterHelper));
-                // If no tests are remaining after filtering, it returns an Error Runner.
                 Runner checkRunner = req.getRunner();
-                if (!(checkRunner instanceof ErrorReportingRunner)) {
-                    long startTime = System.currentTimeMillis();
-                    listener.testRunStarted(classObj.getName(), checkRunner.testCount());
-                    if (mCollectTestsOnly) {
-                        fakeDescriptionExecution(checkRunner.getDescription(), list);
-                    } else {
-                        setTestObjectInformation(checkRunner);
-                        runnerCore.run(checkRunner);
-                    }
-                    listener.testRunEnded(System.currentTimeMillis() - startTime,
-                            Collections.emptyMap());
-                } else {
-                    // Special case where filtering leaves no tests to run, we report no failure
-                    // in this case.
-                    if (EXCLUDE_NO_TEST_FAILURE.equals(
-                            checkRunner.getDescription().getClassName())) {
-                        listener.testRunStarted(classObj.getName(), 0);
-                        listener.testRunEnded(0, Collections.emptyMap());
-                    } else {
-                        // Run the Error runner to get the failures from test classes.
-                        listener.testRunStarted(classObj.getName(), checkRunner.testCount());
-                        RunNotifier failureNotifier = new RunNotifier();
-                        failureNotifier.addListener(list);
-                        checkRunner.run(failureNotifier);
-                        listener.testRunEnded(0, Collections.emptyMap());
-                    }
-                }
+                runJUnit4Tests(listener, checkRunner, classObj.getName());
             } else {
                 throw new IllegalArgumentException(
                         String.format("%s is not a supported test", classObj.getName()));
@@ -436,6 +427,103 @@
         }
     }
 
+    private void runTestCases(ITestInvocationListener listener) throws DeviceNotAvailableException {
+        for (Object obj : getTestMethods()) {
+            if (IRemoteTest.class.isInstance(obj)) {
+                IRemoteTest test = (IRemoteTest) obj;
+                runRemoteTest(listener, test);
+            } else if (TestSuite.class.isInstance(obj)) {
+                TestSuite junitTest = (TestSuite) obj;
+                runJUnit3Tests(listener, junitTest, junitTest.getName());
+            } else if (Description.class.isInstance(obj)) {
+                // Running in a full JUnit4 manner, no downgrade to JUnit3 {@link Test}
+                Description desc = (Description) obj;
+                Request req = Request.aClass(desc.getTestClass());
+                Runner checkRunner = req.filterWith(desc).getRunner();
+                runJUnit4Tests(listener, checkRunner, desc.getClassName());
+            } else {
+                throw new IllegalArgumentException(
+                        String.format("%s is not a supported test", obj));
+            }
+        }
+    }
+
+    private void runRemoteTest(ITestInvocationListener listener, IRemoteTest test)
+            throws DeviceNotAvailableException {
+        if (mCollectTestsOnly) {
+            // Collect only mode is propagated to the test.
+            if (test instanceof ITestCollector) {
+                ((ITestCollector) test).setCollectTestsOnly(true);
+            } else {
+                throw new IllegalArgumentException(
+                        String.format(
+                                "%s does not implement ITestCollector", test.getClass()));
+            }
+        }
+        test.run(listener);
+    }
+
+    private void runJUnit3Tests(
+            ITestInvocationListener listener, TestSuite junitTest, String className)
+            throws DeviceNotAvailableException {
+        if (mCollectTestsOnly) {
+            // Collect only mode, fake the junit test execution.
+            listener.testRunStarted(className, junitTest.countTestCases());
+            Map<String, String> empty = Collections.emptyMap();
+            for (int i = 0; i < junitTest.countTestCases(); i++) {
+                Test t = junitTest.testAt(i);
+                // Test does not have a getName method.
+                // using the toString format instead: <testName>(className)
+                String testName = t.toString().split("\\(")[0];
+                TestIdentifier testId =
+                        new TestIdentifier(t.getClass().getName(), testName);
+                listener.testStarted(testId);
+                listener.testEnded(testId, empty);
+            }
+            Map<String, String> emptyMap = Collections.emptyMap();
+            listener.testRunEnded(0, emptyMap);
+        } else {
+            JUnitRunUtil.runTest(listener, junitTest, className);
+        }
+    }
+
+
+    private void runJUnit4Tests(
+            ITestInvocationListener listener, Runner checkRunner, String className) {
+        JUnitCore runnerCore = new JUnitCore();
+        JUnit4ResultForwarder list = new JUnit4ResultForwarder(listener);
+        runnerCore.addListener(list);
+
+        // If no tests are remaining after filtering, it returns an Error Runner.
+        if (!(checkRunner instanceof ErrorReportingRunner)) {
+            long startTime = System.currentTimeMillis();
+            listener.testRunStarted(className, checkRunner.testCount());
+            if (mCollectTestsOnly) {
+                fakeDescriptionExecution(checkRunner.getDescription(), list);
+            } else {
+                setTestObjectInformation(checkRunner);
+                runnerCore.run(checkRunner);
+            }
+            listener.testRunEnded(System.currentTimeMillis() - startTime,
+                    Collections.emptyMap());
+        } else {
+            // Special case where filtering leaves no tests to run, we report no failure
+            // in this case.
+            if (EXCLUDE_NO_TEST_FAILURE.equals(
+                    checkRunner.getDescription().getClassName())) {
+                listener.testRunStarted(className, 0);
+                listener.testRunEnded(0, Collections.emptyMap());
+            } else {
+                // Run the Error runner to get the failures from test classes.
+                listener.testRunStarted(className, checkRunner.testCount());
+                RunNotifier failureNotifier = new RunNotifier();
+                failureNotifier.addListener(list);
+                checkRunner.run(failureNotifier);
+                listener.testRunEnded(0, Collections.emptyMap());
+            }
+        }
+    }
+
     /**
      * Helper to fake the execution of JUnit4 Tests, using the {@link Description}
      */
@@ -513,6 +601,65 @@
         return suite;
     }
 
+    private List<Object> getTestMethods() throws IllegalArgumentException  {
+        if (mTestMethods != null) {
+            return mTestMethods;
+        }
+        mTestMethods = new ArrayList<>();
+        mFilterHelper.addAllIncludeAnnotation(mIncludeAnnotations);
+        mFilterHelper.addAllExcludeAnnotation(mExcludeAnnotations);
+        List<Class<?>> classes = getClasses();
+        for (Class<?> classObj : classes) {
+            if (Test.class.isAssignableFrom(classObj)) {
+                TestSuite suite = collectTests(collectClasses(classObj));
+                for (int i = 0; i < suite.testCount(); i++) {
+                    TestSuite singletonSuite = new TestSuite();
+                    singletonSuite.setName(classObj.getName());
+                    Test testObj = suite.testAt(i);
+                    singletonSuite.addTest(testObj);
+                    if (IRemoteTest.class.isInstance(testObj)) {
+                        setTestObjectInformation(testObj);
+                    }
+                    mTestMethods.add(singletonSuite);
+                }
+            } else if (IRemoteTest.class.isAssignableFrom(classObj)) {
+                // a pure IRemoteTest is considered a test method itself
+                IRemoteTest test = (IRemoteTest) loadObject(classObj);
+                applyFilters(classObj, test);
+                mTestMethods.add(test);
+            } else if (hasJUnit4Annotation(classObj)) {
+                // Running in a full JUnit4 manner, no downgrade to JUnit3 {@link Test}
+                Request req = Request.aClass(classObj);
+                // Include the method name filtering
+                Set<String> includes = mFilterHelper.getIncludeFilters();
+                if (mMethodName != null) {
+                    includes.add(String.format(TEST_FULL_NAME_FORMAT, classObj.getName(),
+                            mMethodName));
+                }
+
+                req = req.filterWith(new JUnit4TestFilter(mFilterHelper));
+                Runner checkRunner = req.getRunner();
+                Deque<Description> descriptions = new ArrayDeque<>();
+                descriptions.push(checkRunner.getDescription());
+                while (!descriptions.isEmpty()) {
+                    Description desc = descriptions.pop();
+                    if (desc.isTest()) {
+                        mTestMethods.add(desc);
+                    }
+                    List<Description> children = desc.getChildren();
+                    Collections.reverse(children);
+                    for (Description child : children) {
+                        descriptions.push(child);
+                    }
+                }
+            } else {
+                throw new IllegalArgumentException(
+                        String.format("%s is not a supported test", classObj.getName()));
+            }
+        }
+        return mTestMethods;
+    }
+
     protected List<Class<?>> getClasses() throws IllegalArgumentException  {
         List<Class<?>> classes = new ArrayList<>();
         for (String className : mClasses) {
@@ -549,24 +696,7 @@
         try {
             Object testObj = classObj.newInstance();
             // set options
-            if (!mKeyValueOptions.isEmpty()) {
-                try {
-                    OptionSetter setter = new OptionSetter(testObj);
-                    for (String item : mKeyValueOptions) {
-                        String[] fields = item.split(":");
-                        if (fields.length == 2) {
-                            setter.setOptionValue(fields[0], fields[1]);
-                        } else if (fields.length == 3) {
-                            setter.setOptionValue(fields[0], fields[1], fields[2]);
-                        } else {
-                            throw new RuntimeException(
-                                String.format("invalid option spec \"%s\"", item));
-                        }
-                    }
-                } catch (ConfigurationException ce) {
-                    throw new RuntimeException("error passing options down to test class", ce);
-                }
-            }
+            setOptionToLoadedObject(testObj, mKeyValueOptions);
             // Set the test information if needed.
             if (setInfo) {
                 setTestObjectInformation(testObj);
@@ -582,6 +712,43 @@
     }
 
     /**
+     * Helper for Device Runners to use to set options the same way as HostTest, from set-option.
+     *
+     * @param testObj the object that will receive the options.
+     * @param keyValueOptions the list of options formatted as HostTest set-option requires.
+     */
+    public static void setOptionToLoadedObject(Object testObj, List<String> keyValueOptions) {
+        if (!keyValueOptions.isEmpty()) {
+            try {
+                OptionSetter setter = new OptionSetter(testObj);
+                for (String item : keyValueOptions) {
+                    String[] fields = item.split(":");
+                    if (fields.length == 2) {
+                        if (fields[1].contains("=")) {
+                            String[] values = fields[1].split("=");
+                            if (values.length != 2) {
+                                throw new RuntimeException(
+                                        String.format(
+                                                "set-option provided '%s' format is invalid. Only one "
+                                                        + "'=' is allowed",
+                                                item));
+                            }
+                            setter.setOptionValue(fields[0], values[0], values[1]);
+                        } else {
+                            setter.setOptionValue(fields[0], fields[1]);
+                        }
+                    } else {
+                        throw new RuntimeException(
+                                String.format("invalid option spec \"%s\"", item));
+                    }
+                }
+            } catch (ConfigurationException ce) {
+                throw new RuntimeException("error passing options down to test class", ce);
+            }
+        }
+    }
+
+    /**
      * Check if an elements that has annotation pass the filter. Exposed for unit testing.
      * @param annotatedElement
      * @return false if the test should not run.
@@ -641,10 +808,13 @@
     }
 
     /**
-     * We split by --class, and if each individual IRemoteTest is splitable we split them too.
+     * We split by individual by either test class or method.
      */
     @Override
-    public Collection<IRemoteTest> split() {
+    public Collection<IRemoteTest> split(int shardCount) {
+        if (shardCount < 1) {
+            throw new IllegalArgumentException("Must have at least 1 shard");
+        }
         List<IRemoteTest> listTests = new ArrayList<>();
         List<Class<?>> classes = getClasses();
         if (classes.isEmpty()) {
@@ -653,22 +823,62 @@
         if (mMethodName != null && classes.size() > 1) {
             throw new IllegalArgumentException("Method name given with multiple test classes");
         }
-        if (classes.size() == 1) {
-            // Cannot shard if only no class or one class specified
-            // TODO: Consider doing class sharding too if its a suite.
+        List<? extends Object> testObjects;
+        if (shardUnitIsMethod()) {
+            testObjects = getTestMethods();
+        } else {
+            testObjects = classes;
+            // ignore shardCount when shard unit is class;
+            // simply shard by the number of classes
+            shardCount = testObjects.size();
+        }
+        if (testObjects.size() == 1) {
             return null;
         }
-        for (Class<?> classObj : classes) {
-            HostTest test = createHostTest(classObj);
-            test.mRuntimeHint = mRuntimeHint / classes.size();
-            // Carry over non-annotation filters to shards.
-            test.addAllExcludeFilters(mFilterHelper.getExcludeFilters());
-            test.addAllIncludeFilters(mFilterHelper.getIncludeFilters());
-            listTests.add(test);
+        int i = 0;
+        int numTotalTestCases = countTestCases();
+        for (Object testObj : testObjects) {
+            Class<?> classObj = Class.class.isInstance(testObj) ? (Class<?>)testObj : null;
+            HostTest test;
+            if (i >= listTests.size()) {
+                test = createHostTest(classObj);
+                test.mRuntimeHint = 0;
+                // Carry over non-annotation filters to shards.
+                test.addAllExcludeFilters(mFilterHelper.getExcludeFilters());
+                test.addAllIncludeFilters(mFilterHelper.getIncludeFilters());
+                listTests.add(test);
+            }
+            test = (HostTest) listTests.get(i);
+            Collection<? extends Object> subTests;
+            if (classObj != null) {
+                test.addClassName(classObj.getName());
+                subTests = test.mClasses;
+            } else {
+                test.addTestMethod(testObj);
+                subTests = test.mTestMethods;
+            }
+            test.mRuntimeHint = mRuntimeHint * subTests.size() / numTotalTestCases;
+            i = (i + 1) % shardCount;
         }
+
         return listTests;
     }
 
+    private void addTestMethod(Object testObject) {
+        if (mTestMethods == null) {
+            mTestMethods = new ArrayList<>();
+            mClasses.clear();
+        }
+        mTestMethods.add(testObject);
+        if (IRemoteTest.class.isInstance(testObject)) {
+            addClassName(testObject.getClass().getName());
+        } else if (TestSuite.class.isInstance(testObject)) {
+            addClassName(((TestSuite)testObject).getName());
+        } else if (Description.class.isInstance(testObject)) {
+            addClassName(((Description)testObject).getTestClass().getName());
+        }
+    }
+
     /**
      * Add a class to be ran by HostTest.
      */
@@ -698,7 +908,6 @@
 
     @Override
     public IRemoteTest getTestShard(int shardCount, int shardIndex) {
-        IRemoteTest test = null;
         List<Class<?>> classes = getClasses();
         if (classes.isEmpty()) {
             throw new IllegalArgumentException("Missing Test class name");
@@ -706,29 +915,16 @@
         if (mMethodName != null && classes.size() > 1) {
             throw new IllegalArgumentException("Method name given with multiple test classes");
         }
-        int numTotalTestCases = countTestCases();
-        int i = 0;
-        for (Class<?> classObj : classes) {
-            if (i % shardCount == shardIndex) {
-                if (test == null) {
-                    test = createHostTest(classObj);
-                } else {
-                    ((HostTest) test).addClassName(classObj.getName());
-                }
-                // Carry over non-annotation filters to shards.
-                ((HostTest) test).addAllExcludeFilters(mFilterHelper.getExcludeFilters());
-                ((HostTest) test).addAllIncludeFilters(mFilterHelper.getIncludeFilters());
-            }
-            i++;
-        }
+        HostTest test = createTestShard(shardCount, shardIndex);
         // In case we don't have enough classes to shard, we return a Stub.
         if (test == null) {
             test = createHostTest(null);
-            ((HostTest) test).mSkipTestClassCheck = true;
-            ((HostTest) test).mClasses.clear();
-            ((HostTest) test).mRuntimeHint = 0l;
+            test.mSkipTestClassCheck = true;
+            test.mClasses.clear();
+            test.mRuntimeHint = 0l;
         } else {
-            int newCount = ((HostTest) test).countTestCases();
+            int newCount = test.countTestCases();
+            int numTotalTestCases = countTestCases();
             // In case of counting inconsistency we raise the issue. Should not happen if we are
             // counting properly. Here as a security.
             if (newCount > numTotalTestCases) {
@@ -738,11 +934,35 @@
             // update the runtime hint on pro-rate of number of tests.
             if (newCount == 0) {
                 // In case there is not tests left.
-                ((HostTest) test).mRuntimeHint = 0l;
+                test.mRuntimeHint = 0l;
             } else {
-                ((HostTest) test).mRuntimeHint = (mRuntimeHint * newCount) / numTotalTestCases;
+                test.mRuntimeHint = (mRuntimeHint * newCount) / numTotalTestCases;
             }
         }
         return test;
     }
+
+    private HostTest createTestShard(int shardCount, int shardIndex) {
+        int i = 0;
+        HostTest test = null;
+        List<? extends Object> tests = shardUnitIsMethod() ? getTestMethods() : getClasses();
+        for (Object testObj : tests) {
+            Class<?> classObj = Class.class.isInstance(testObj) ? (Class<?>)testObj : null;
+            if (i % shardCount == shardIndex) {
+                if (test == null) {
+                    test = createHostTest(classObj);
+                }
+                if (classObj != null) {
+                    test.addClassName(classObj.getName());
+                } else {
+                    test.addTestMethod(testObj);
+                }
+                // Carry over non-annotation filters to shards.
+                test.addAllExcludeFilters(mFilterHelper.getExcludeFilters());
+                test.addAllIncludeFilters(mFilterHelper.getIncludeFilters());
+            }
+            i++;
+        }
+        return test;
+    }
 }
diff --git a/src/com/android/tradefed/testtype/ISetOptionReceiver.java b/src/com/android/tradefed/testtype/ISetOptionReceiver.java
new file mode 100644
index 0000000..1ca04d3
--- /dev/null
+++ b/src/com/android/tradefed/testtype/ISetOptionReceiver.java
@@ -0,0 +1,24 @@
+/*
+ * Copyright (C) 2017 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.tradefed.config.Option;
+
+/**
+ * Implementation of this interface should have an {@link Option} with a "set-option" name linked to
+ * {@link HostTest#SET_OPTION_NAME}.
+ */
+public interface ISetOptionReceiver {}
diff --git a/src/com/android/tradefed/testtype/InstrumentationTest.java b/src/com/android/tradefed/testtype/InstrumentationTest.java
index d99c1ef..206679e 100644
--- a/src/com/android/tradefed/testtype/InstrumentationTest.java
+++ b/src/com/android/tradefed/testtype/InstrumentationTest.java
@@ -118,6 +118,15 @@
                     + "the next test. For no timeout, set to 0.")
     private int mTestTimeout = 5 * 60 * 1000;  // default to 5 minutes
 
+    @Option(
+        name = "max-timeout",
+        description =
+                "Sets the max timeout for the instrumentation to terminate. "
+                        + "For no timeout, set to 0.",
+        isTimeVal = true
+    )
+    private long mMaxTimeout = 0l;
+
     @Option(name = "size",
             description="Restrict test to a specific test size.")
     private String mTestSize = null;
@@ -440,6 +449,11 @@
         return mTestTimeout;
     }
 
+    /** Returns the max timeout set for the instrumentation. */
+    public long getMaxTimeout() {
+        return mMaxTimeout;
+    }
+
     /**
      * Set the optional file to install that contains the tests.
      *
@@ -692,6 +706,7 @@
                     mTestTimeout, mShellTimeout));
         }
         runner.setMaxTimeToOutputResponse(mShellTimeout, TimeUnit.MILLISECONDS);
+        runner.setMaxTimeout(mMaxTimeout, TimeUnit.MILLISECONDS);
         addInstrumentationArg(TEST_TIMEOUT_INST_ARGS_KEY, Long.toString(mTestTimeout));
     }
 
diff --git a/src/com/android/tradefed/testtype/MetricTestCase.java b/src/com/android/tradefed/testtype/MetricTestCase.java
index 92a9ade..00a321b 100644
--- a/src/com/android/tradefed/testtype/MetricTestCase.java
+++ b/src/com/android/tradefed/testtype/MetricTestCase.java
@@ -15,9 +15,15 @@
  */
 package com.android.tradefed.testtype;
 
+import com.android.tradefed.result.InputStreamSource;
+import com.android.tradefed.result.LogDataType;
+import com.android.tradefed.result.SnapshotInputStreamSource;
+
 import junit.framework.TestCase;
 
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 /**
@@ -28,6 +34,7 @@
 public class MetricTestCase extends TestCase {
 
     public Map<String, String> mMetrics = new HashMap<>();
+    public List<LogHolder> mLogs = new ArrayList<>();
 
     public MetricTestCase() {
         super();
@@ -47,4 +54,36 @@
     public final void addTestMetric(String key, String value) {
         mMetrics.put(key, value);
     }
+
+    /**
+     * Callback from JUnit3 forwarder in order to get the logs from a test.
+     *
+     * @param dataName a String descriptive name of the data. e.g. "device_logcat". Note dataName
+     *     may not be unique per invocation. ie implementers must be able to handle multiple calls
+     *     with same dataName
+     * @param dataType the LogDataType of the data
+     * @param dataStream the InputStreamSource of the data. Implementers should call
+     *     createInputStream to start reading the data, and ensure to close the resulting
+     *     InputStream when complete. Callers should ensure the source of the data remains present
+     *     and accessible until the testLog method completes.
+     */
+    public final void addTestLog(
+            String dataName, LogDataType dataType, InputStreamSource dataStream) {
+        mLogs.add(new LogHolder(dataName, dataType, dataStream));
+    }
+
+    /** Structure to hold a log file to be reported. */
+    public static class LogHolder {
+        public final String mDataName;
+        public final LogDataType mDataType;
+        public final InputStreamSource mDataStream;
+
+        public LogHolder(String dataName, LogDataType dataType, InputStreamSource dataStream) {
+            mDataName = dataName;
+            mDataType = dataType;
+            // We hold a copy because the caller will most likely cancel the stream after.
+            mDataStream =
+                    new SnapshotInputStreamSource("LogHolder", dataStream.createInputStream());
+        }
+    }
 }
diff --git a/src/com/android/tradefed/testtype/NoisyDryRunTest.java b/src/com/android/tradefed/testtype/NoisyDryRunTest.java
index 8681c20..7783e87 100644
--- a/src/com/android/tradefed/testtype/NoisyDryRunTest.java
+++ b/src/com/android/tradefed/testtype/NoisyDryRunTest.java
@@ -28,6 +28,7 @@
 import com.android.tradefed.log.LogUtil.CLog;
 import com.android.tradefed.result.ITestInvocationListener;
 import com.android.tradefed.util.QuotationAwareTokenizer;
+import com.android.tradefed.util.RunUtil;
 import com.android.tradefed.util.StreamUtil;
 import com.android.tradefed.util.keystore.StubKeyStoreClient;
 
@@ -41,9 +42,16 @@
  */
 public class NoisyDryRunTest implements IRemoteTest {
 
+    private static final long SLEEP_INTERVAL_MILLI_SEC = 5 * 1000;
+
     @Option(name = "cmdfile", description = "The cmdfile to run noisy dry run on.")
     private String mCmdfile = null;
 
+    @Option(name = "timeout",
+            description = "The timeout to wait cmd file be ready.",
+            isTimeVal = true)
+    private long mTimeoutMilliSec = 0;
+
     @Override
     public void run(ITestInvocationListener listener) throws DeviceNotAvailableException {
         List<CommandLine> commands = testCommandFile(listener, mCmdfile);
@@ -59,7 +67,9 @@
         listener.testStarted(parseFileTest);
         CommandFileParser parser = new CommandFileParser();
         try {
-            return parser.parseFile(new File(filename));
+            File file = new File(filename);
+            checkFileWithTimeout(file);
+            return parser.parseFile(file);
         } catch (IOException | ConfigurationException e) {
             listener.testFailed(parseFileTest, StreamUtil.getStackTrace(e));
             return null;
@@ -69,6 +79,35 @@
         }
     }
 
+    /**
+     * If the file doesn't exist, we want to wait a while and check.
+     *
+     * @param file
+     * @throws IOException
+     */
+    @VisibleForTesting
+    void checkFileWithTimeout(File file) throws IOException {
+        long timeout = currentTimeMillis() + mTimeoutMilliSec;
+        while (!file.exists() && currentTimeMillis() < timeout) {
+            CLog.w("%s doesn't exist, wait and recheck.", file.getAbsoluteFile());
+            sleep();
+        }
+        if (!file.exists()) {
+            throw new IOException(
+                    String.format("%s doesn't exist.", file.getAbsoluteFile()));
+        }
+    }
+
+    @VisibleForTesting
+    long currentTimeMillis() {
+        return System.currentTimeMillis();
+    }
+
+    @VisibleForTesting
+    void sleep() {
+        RunUtil.getDefault().sleep(SLEEP_INTERVAL_MILLI_SEC);
+    }
+
     private void testCommandLines(ITestInvocationListener listener, List<CommandLine> commands) {
         listener.testRunStarted(NoisyDryRunTest.class.getCanonicalName() + "_parseCommands",
                 commands.size());
@@ -93,9 +132,4 @@
         }
         listener.testRunEnded(0, new HashMap<String, String>());
     }
-
-    @VisibleForTesting
-    void setCmdFile(String cmdfile) {
-        mCmdfile = cmdfile;
-    }
 }
diff --git a/src/com/android/tradefed/testtype/SubprocessTfLauncher.java b/src/com/android/tradefed/testtype/SubprocessTfLauncher.java
index 51bad42..0c6245e 100644
--- a/src/com/android/tradefed/testtype/SubprocessTfLauncher.java
+++ b/src/com/android/tradefed/testtype/SubprocessTfLauncher.java
@@ -99,8 +99,8 @@
     protected List<String> mCmdArgs = null;
     // The absolute path to the build's root directory.
     protected String mRootDir = null;
+    protected IConfiguration mConfig;
     private IInvocationContext mContext;
-    private IConfiguration mConfig;
 
     @Override
     public void setInvocationContext(IInvocationContext invocationContext) {
@@ -168,9 +168,15 @@
             // If the global configuration is not set in option, create a filtered global
             // configuration for subprocess to use.
             try {
+                String[] configs =
+                        new String[] {
+                            GlobalConfiguration.DEVICE_MANAGER_TYPE_NAME,
+                            GlobalConfiguration.KEY_STORE_TYPE_NAME
+                        };
                 File filteredGlobalConfig =
                         FileUtil.createTempFile("filtered_global_config", ".config");
-                GlobalConfiguration.getInstance().cloneConfigWithFilter(filteredGlobalConfig, null);
+                GlobalConfiguration.getInstance()
+                        .cloneConfigWithFilter(filteredGlobalConfig, configs);
                 mFilteredGlobalConfig = filteredGlobalConfig.getAbsolutePath();
                 mGlobalConfig = mFilteredGlobalConfig;
             } catch (IOException e) {
diff --git a/src/com/android/tradefed/testtype/TfTestLauncher.java b/src/com/android/tradefed/testtype/TfTestLauncher.java
index 87aae1f..561f10c 100644
--- a/src/com/android/tradefed/testtype/TfTestLauncher.java
+++ b/src/com/android/tradefed/testtype/TfTestLauncher.java
@@ -29,6 +29,7 @@
 import com.android.tradefed.util.HprofAllocSiteParser;
 import com.android.tradefed.util.RunUtil;
 import com.android.tradefed.util.StreamUtil;
+import com.android.tradefed.util.TarUtil;
 
 import com.google.common.annotations.VisibleForTesting;
 
@@ -120,7 +121,7 @@
                 // cutoff the min value we look at.
                 String hprofAgent =
                         String.format(
-                                "-agentlib:hprof=heap=sites,cutoff=0.01,depth=12,verbose=n,file=%s",
+                                "-agentlib:hprof=heap=sites,cutoff=0.01,depth=16,verbose=n,file=%s",
                                 mHprofFile.getAbsolutePath());
                 args.add(hprofAgent);
             }
@@ -371,11 +372,17 @@
             return;
         }
         InputStreamSource memory = null;
+        File tmpGzip = null;
         try {
-            memory = new FileInputStreamSource(hprofFile);
-            listener.testLog("hprof", LogDataType.TEXT, memory);
+            tmpGzip = TarUtil.gzip(hprofFile);
+            memory = new FileInputStreamSource(tmpGzip);
+            listener.testLog("hprof", LogDataType.GZIP, memory);
+        } catch (IOException e) {
+            CLog.e(e);
+            return;
         } finally {
             StreamUtil.cancel(memory);
+            FileUtil.deleteFile(tmpGzip);
         }
         HprofAllocSiteParser parser = new HprofAllocSiteParser();
         try {
diff --git a/src/com/android/tradefed/testtype/UiAutomatorRunner.java b/src/com/android/tradefed/testtype/UiAutomatorRunner.java
index c311890..969f930 100644
--- a/src/com/android/tradefed/testtype/UiAutomatorRunner.java
+++ b/src/com/android/tradefed/testtype/UiAutomatorRunner.java
@@ -60,6 +60,7 @@
     private String[] mJarPaths;
     private String mPackageName;
     // default to no timeout
+    private long mMaxTimeout = 0l;
     private long mMaxTimeToOutputResponse = 0;
     private IDevice mRemoteDevice;
     private String mRunName;
@@ -306,8 +307,8 @@
         mParser = new InstrumentationResultParser(runName, listeners);
 
         try {
-            mRemoteDevice.executeShellCommand(cmdLine,
-                    mParser, mMaxTimeToOutputResponse, TimeUnit.MILLISECONDS);
+            mRemoteDevice.executeShellCommand(
+                    cmdLine, mParser, mMaxTimeout, mMaxTimeToOutputResponse, TimeUnit.MILLISECONDS);
         } catch (IOException e) {
             CLog.w(String.format("IOException %1$s when running tests %2$s on %3$s",
                     e.toString(), getPackageName(), mRemoteDevice.getSerialNumber()));
@@ -349,4 +350,9 @@
     public void setEnforceTimeStamp(boolean arg0) {
         // ignore, UiAutomator runner does not need this.
     }
+
+    @Override
+    public void setMaxTimeout(long maxTimeout, TimeUnit unit) {
+        mMaxTimeout = unit.toMillis(maxTimeout);
+    }
 }
diff --git a/src/com/android/tradefed/testtype/VersionedTfLauncher.java b/src/com/android/tradefed/testtype/VersionedTfLauncher.java
index 08a398d..96307a9 100644
--- a/src/com/android/tradefed/testtype/VersionedTfLauncher.java
+++ b/src/com/android/tradefed/testtype/VersionedTfLauncher.java
@@ -21,6 +21,7 @@
 import com.android.tradefed.config.OptionCopier;
 import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.device.NullDevice;
+import com.android.tradefed.device.StubDevice;
 import com.android.tradefed.util.StringEscapeUtils;
 
 import java.io.File;
@@ -82,7 +83,7 @@
             ITestDevice device = mDeviceInfos.entrySet().iterator().next().getKey();
             if (device.getIDevice() instanceof NullDevice) {
                 mCmdArgs.add("--null-device");
-            } else {
+            } else if (!(device.getIDevice() instanceof StubDevice)) {
                 String serial = device.getSerialNumber();
                 mCmdArgs.add("--serial");
                 mCmdArgs.add(serial);
diff --git a/src/com/android/tradefed/testtype/suite/ITestSuite.java b/src/com/android/tradefed/testtype/suite/ITestSuite.java
index 6fc8a15..37a0e60 100644
--- a/src/com/android/tradefed/testtype/suite/ITestSuite.java
+++ b/src/com/android/tradefed/testtype/suite/ITestSuite.java
@@ -23,6 +23,7 @@
 import com.android.tradefed.config.OptionCopier;
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.invoker.IInvocationContext;
 import com.android.tradefed.log.LogUtil.CLog;
 import com.android.tradefed.result.ITestInvocationListener;
 import com.android.tradefed.result.InputStreamSource;
@@ -31,9 +32,12 @@
 import com.android.tradefed.suite.checker.ISystemStatusCheckerReceiver;
 import com.android.tradefed.testtype.IBuildReceiver;
 import com.android.tradefed.testtype.IDeviceTest;
+import com.android.tradefed.testtype.IInvocationContextReceiver;
 import com.android.tradefed.testtype.IRemoteTest;
+import com.android.tradefed.testtype.IRuntimeHintProvider;
 import com.android.tradefed.testtype.IShardableTest;
 import com.android.tradefed.testtype.ITestCollector;
+import com.android.tradefed.util.TimeUtil;
 
 import java.util.ArrayList;
 import java.util.Collection;
@@ -52,7 +56,9 @@
                 IBuildReceiver,
                 ISystemStatusCheckerReceiver,
                 IShardableTest,
-                ITestCollector {
+                ITestCollector,
+                IInvocationContextReceiver,
+                IRuntimeHintProvider {
 
     public static final String MODULE_CHECKER_PRE = "PreModuleChecker";
     public static final String MODULE_CHECKER_POST = "PostModuleChecker";
@@ -109,6 +115,7 @@
     private ITestDevice mDevice;
     private IBuildInfo mBuildInfo;
     private List<ISystemStatusChecker> mSystemStatusCheckers;
+    private IInvocationContext mContext;
 
     // Sharding attributes
     private boolean mIsSharded = false;
@@ -157,8 +164,12 @@
                                         "Configuration %s cannot be run in a suite.",
                                         config.getValue().getName())));
             }
-            ModuleDefinition module = new ModuleDefinition(config.getKey(),
-                    config.getValue().getTests(), config.getValue().getTargetPreparers());
+            ModuleDefinition module =
+                    new ModuleDefinition(
+                            config.getKey(),
+                            config.getValue().getTests(),
+                            config.getValue().getTargetPreparers(),
+                            config.getValue().getConfigurationDescription());
             module.setDevice(mDevice);
             module.setBuild(mBuildInfo);
             runModules.add(module);
@@ -208,7 +219,15 @@
                 if (module.hasTests()) {
                     continue;
                 }
-                runSingleModule(module, listener, failureListener);
+
+                try {
+                    mContext.setModuleInvocationContext(module.getModuleInvocationContext());
+                    runSingleModule(module, listener, failureListener);
+                } finally {
+                    // clear out module invocation context since we are now done with module
+                    // execution
+                    mContext.setModuleInvocationContext(null);
+                }
             }
         } catch (DeviceNotAvailableException e) {
             CLog.e(
@@ -456,4 +475,31 @@
     public void setShouldMakeDynamicModule(boolean dynamicModule) {
         mShouldMakeDynamicModule = dynamicModule;
     }
+
+    /** {@inheritDoc} */
+    @Override
+    public void setInvocationContext(IInvocationContext invocationContext) {
+        mContext = invocationContext;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public long getRuntimeHint() {
+        if (mDirectModule != null) {
+            CLog.e(
+                    "    %s: %s",
+                    mDirectModule.getId(),
+                    TimeUtil.formatElapsedTime(mDirectModule.getRuntimeHint()));
+            return mDirectModule.getRuntimeHint();
+        }
+        return 0l;
+    }
+
+    /**
+     * Returns the {@link ModuleDefinition} to be executed directly, or null if none yet (when the
+     * ITestSuite has not been sharded yet).
+     */
+    public ModuleDefinition getDirectModule() {
+        return mDirectModule;
+    }
 }
diff --git a/src/com/android/tradefed/testtype/suite/ModuleDefinition.java b/src/com/android/tradefed/testtype/suite/ModuleDefinition.java
index ef193b7..8a0c5b1 100644
--- a/src/com/android/tradefed/testtype/suite/ModuleDefinition.java
+++ b/src/com/android/tradefed/testtype/suite/ModuleDefinition.java
@@ -20,10 +20,12 @@
 import com.android.ddmlib.testrunner.TestResult;
 import com.android.ddmlib.testrunner.TestRunResult;
 import com.android.tradefed.build.IBuildInfo;
+import com.android.tradefed.config.ConfigurationDescriptor;
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.device.DeviceUnresponsiveException;
 import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.invoker.IInvocationContext;
+import com.android.tradefed.invoker.InvocationContext;
 import com.android.tradefed.log.ILogRegistry.EventType;
 import com.android.tradefed.log.ITestLogger;
 import com.android.tradefed.log.LogRegistry;
@@ -39,6 +41,7 @@
 import com.android.tradefed.testtype.IBuildReceiver;
 import com.android.tradefed.testtype.IDeviceTest;
 import com.android.tradefed.testtype.IRemoteTest;
+import com.android.tradefed.testtype.IRuntimeHintProvider;
 import com.android.tradefed.testtype.ITestCollector;
 import com.android.tradefed.util.StreamUtil;
 
@@ -62,6 +65,8 @@
     public static final String MODULE_NAME = "module-name";
     public static final String MODULE_ABI = "module-abi";
 
+    private final IInvocationContext mModuleInvocationContext;
+
     private final String mId;
     private Collection<IRemoteTest> mTests = null;
     private List<ITargetPreparer> mPreparers = new ArrayList<>();
@@ -90,11 +95,25 @@
      * @param name unique name of the test configuration.
      * @param tests list of {@link IRemoteTest} that needs to run.
      * @param preparers list of {@link ITargetPreparer} to be used to setup the device.
+     * @param configDescriptor the {@link ConfigurationDescriptor} of the underlying module config.
      */
     public ModuleDefinition(
-            String name, Collection<IRemoteTest> tests, List<ITargetPreparer> preparers) {
+            String name,
+            Collection<IRemoteTest> tests,
+            List<ITargetPreparer> preparers,
+            ConfigurationDescriptor configDescriptor) {
         mId = name;
         mTests = tests;
+
+        mModuleInvocationContext = new InvocationContext();
+        mModuleInvocationContext.setConfigurationDescriptor(configDescriptor);
+        mModuleInvocationContext.addInvocationAttribute(MODULE_NAME, mId);
+        // If available in the suite, add the abi name
+        if (configDescriptor.getAbi() != null) {
+            mModuleInvocationContext.addInvocationAttribute(
+                    MODULE_ABI, configDescriptor.getAbi().getName());
+        }
+
         for (ITargetPreparer preparer : preparers) {
             mPreparers.add(preparer);
             if (preparer instanceof ITargetCleaner) {
@@ -121,6 +140,23 @@
     }
 
     /**
+     * Add some {@link IRemoteTest} to be executed as part of the module. Used when merging two
+     * modules.
+     */
+    void addTests(List<IRemoteTest> test) {
+        synchronized (mTests) {
+            mTests.addAll(test);
+        }
+    }
+
+    /** Returns the current number of {@link IRemoteTest} waiting to be executed. */
+    public int numTests() {
+        synchronized (mTests) {
+            return mTests.size();
+        }
+    }
+
+    /**
      * Return True if the Module still has {@link IRemoteTest} to run in its pool. False otherwise.
      */
     protected boolean hasTests() {
@@ -244,6 +280,7 @@
                     CLog.e("Module '%s' - test '%s' threw exception:", getId(), test.getClass());
                     CLog.e(re);
                     CLog.e("Proceeding to the next test.");
+                    reportFailure(new ResultForwarder(currentTestListener), re.getMessage());
                 } catch (DeviceUnresponsiveException due) {
                     // being able to catch a DeviceUnresponsiveException here implies that
                     // recovery was successful, and test execution should proceed to next
@@ -253,6 +290,7 @@
                                     + "successful, proceeding with next module. Stack trace:");
                     CLog.w(due);
                     CLog.w("Proceeding to the next test.");
+                    reportFailure(new ResultForwarder(currentTestListener), due.getMessage());
                 } catch (DeviceNotAvailableException dnae) {
                     // We do special logging of some information in Context of the module for easier
                     // debugging.
@@ -295,6 +333,10 @@
         }
     }
 
+    private void reportFailure(ITestInvocationListener listener, String errorMessage) {
+        listener.testRunFailed(errorMessage);
+    }
+
     /** Helper to log the device events. */
     private void logDeviceEvent(EventType event, String serial, Throwable t, String moduleId) {
         Map<String, String> args = new HashMap<>();
@@ -419,6 +461,19 @@
         return getId();
     }
 
+    /** Returns the approximate time to run all the tests in the module. */
+    public long getRuntimeHint() {
+        long hint = 0l;
+        for (IRemoteTest test : mTests) {
+            if (test instanceof IRuntimeHintProvider) {
+                hint += ((IRuntimeHintProvider) test).getRuntimeHint();
+            } else {
+                hint += 60000;
+            }
+        }
+        return hint;
+    }
+
     /** Returns the list of {@link ITargetPreparer} defined for this module. */
     @VisibleForTesting
     List<ITargetPreparer> getTargetPreparers() {
@@ -430,4 +485,9 @@
     List<IRemoteTest> getTests() {
         return new ArrayList<>(mTests);
     }
+
+    /** Returns the {@link IInvocationContext} associated with the module. */
+    public IInvocationContext getModuleInvocationContext() {
+        return mModuleInvocationContext;
+    }
 }
diff --git a/src/com/android/tradefed/testtype/suite/ModuleMerger.java b/src/com/android/tradefed/testtype/suite/ModuleMerger.java
new file mode 100644
index 0000000..9410a2e
--- /dev/null
+++ b/src/com/android/tradefed/testtype/suite/ModuleMerger.java
@@ -0,0 +1,64 @@
+/*
+ * Copyright (C) 2017 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.suite;
+
+/**
+ * Helper class for operation related to merging {@link ITestSuite} and {@link ModuleDefinition}
+ * after a split.
+ */
+public class ModuleMerger {
+
+    private static void mergeModules(ModuleDefinition module1, ModuleDefinition module2) {
+        if (!module1.getId().equals(module2.getId())) {
+            throw new IllegalArgumentException(
+                    String.format(
+                            "Modules must have the same id to be mergeable: received %s and "
+                                    + "%s",
+                            module1.getId(), module2.getId()));
+        }
+        module1.addTests(module2.getTests());
+    }
+
+    /**
+     * Merge the modules from one suite to another.
+     *
+     * @param suite1 the suite that will receive the module from the other.
+     * @param suite2 the suite that will give the module.
+     */
+    public static void mergeSplittedITestSuite(ITestSuite suite1, ITestSuite suite2) {
+        if (suite1.getDirectModule() == null) {
+            throw new IllegalArgumentException("suite was not a splitted suite.");
+        }
+        if (suite2.getDirectModule() == null) {
+            throw new IllegalArgumentException("suite was not a splitted suite.");
+        }
+        mergeModules(suite1.getDirectModule(), suite2.getDirectModule());
+    }
+
+    /** Returns true if the two suites are part of the same original split. False otherwise. */
+    public static boolean arePartOfSameSuite(ITestSuite suite1, ITestSuite suite2) {
+        if (suite1.getDirectModule() == null) {
+            return false;
+        }
+        if (suite2.getDirectModule() == null) {
+            return false;
+        }
+        if (!suite1.getDirectModule().getId().equals(suite2.getDirectModule().getId())) {
+            return false;
+        }
+        return true;
+    }
+}
diff --git a/src/com/android/tradefed/testtype/suite/ModuleSplitter.java b/src/com/android/tradefed/testtype/suite/ModuleSplitter.java
index eb04cff..83d6c19 100644
--- a/src/com/android/tradefed/testtype/suite/ModuleSplitter.java
+++ b/src/com/android/tradefed/testtype/suite/ModuleSplitter.java
@@ -88,12 +88,17 @@
             boolean dynamicModule) {
         // If this particular configuration module is declared as 'not shardable' we take it whole
         // but still split the individual IRemoteTest in a pool.
-        if (config.getConfigurationDescription().isNotShardable()) {
+        if (config.getConfigurationDescription().isNotShardable()
+                || (!dynamicModule
+                        && config.getConfigurationDescription().isNotStrictShardable())) {
             for (int i = 0; i < config.getTests().size(); i++) {
                 if (dynamicModule) {
                     ModuleDefinition module =
                             new ModuleDefinition(
-                                    moduleName, config.getTests(), clonePreparers(config));
+                                    moduleName,
+                                    config.getTests(),
+                                    clonePreparers(config),
+                                    config.getConfigurationDescription());
                     currentList.add(module);
                 } else {
                     addModuleToListFromSingleTest(
@@ -114,7 +119,10 @@
                         for (int i = 0; i < shardCount; i++) {
                             ModuleDefinition module =
                                     new ModuleDefinition(
-                                            moduleName, shardedTests, clonePreparers(config));
+                                            moduleName,
+                                            shardedTests,
+                                            clonePreparers(config),
+                                            config.getConfigurationDescription());
                             currentList.add(module);
                         }
                     } else {
@@ -144,7 +152,11 @@
         List<IRemoteTest> testList = new ArrayList<>();
         testList.add(test);
         ModuleDefinition module =
-                new ModuleDefinition(moduleName, testList, clonePreparers(config));
+                new ModuleDefinition(
+                        moduleName,
+                        testList,
+                        clonePreparers(config),
+                        config.getConfigurationDescription());
         currentList.add(module);
     }
 
diff --git a/src/com/android/tradefed/util/BluetoothUtils.java b/src/com/android/tradefed/util/BluetoothUtils.java
index 7728b1e..cfc8b4b 100644
--- a/src/com/android/tradefed/util/BluetoothUtils.java
+++ b/src/com/android/tradefed/util/BluetoothUtils.java
@@ -55,6 +55,7 @@
     public static final String BTSNOOP_CMD = "setprop persist.bluetooth.btsnoopenable ";
     public static final String BTSNOOP_ENABLE_CMD = BTSNOOP_CMD + "true";
     public static final String BTSNOOP_DISABLE_CMD = BTSNOOP_CMD + "false";
+    public static final String GOLD_BTSNOOP_LOG_PATH = "/data/misc/bluetooth/logs/btsnoop_hci.log";
     public static final String O_BUILD = "O";
 
     /**
@@ -192,6 +193,8 @@
             throws DeviceNotAvailableException {
         if (isGoldAndAbove(device)) {
             device.executeShellCommand(BTSNOOP_ENABLE_CMD);
+            disable(device);
+            enable(device);
             return true;
         }
         return enableBtsnoopLogging(device, null);
@@ -222,6 +225,8 @@
             throws DeviceNotAvailableException {
         if (isGoldAndAbove(device)) {
             device.executeShellCommand(BTSNOOP_DISABLE_CMD);
+            disable(device);
+            enable(device);
             return true;
         }
         return disableBtsnoopLogging(device, null);
@@ -305,6 +310,9 @@
      */
     public static String getBtSnoopLogFilePath(ITestDevice device)
             throws DeviceNotAvailableException {
+        if (isGoldAndAbove(device)) {
+            return GOLD_BTSNOOP_LOG_PATH;
+        }
         String snoopfileSetting =
                 device.executeShellCommand(
                         String.format("cat %s | grep BtSnoopFileName", BT_STACK_CONF));
diff --git a/src/com/android/tradefed/util/BuildTestsZipUtils.java b/src/com/android/tradefed/util/BuildTestsZipUtils.java
index 310a9b9..8858b1d 100644
--- a/src/com/android/tradefed/util/BuildTestsZipUtils.java
+++ b/src/com/android/tradefed/util/BuildTestsZipUtils.java
@@ -61,8 +61,6 @@
                 dirs.add(FileUtil.getFileForPath(dir, "DATA", "priv-app", apkBase));
                 // Files in out dir will be in data/app/apk_name
                 dirs.add(FileUtil.getFileForPath(dir, "data", "app", apkBase));
-                // Files in testcases directory will be in //apkBase
-                dirs.add(FileUtil.getFileForPath(dir, apkBase));
             }
         }
         // reverse the order so ones provided via command line last can be searched first
@@ -70,12 +68,14 @@
 
         List<File> expandedTestDirs = new ArrayList<>();
         if (buildInfo != null && buildInfo instanceof IDeviceBuildInfo) {
-            File testsDir = ((IDeviceBuildInfo)buildInfo).getTestsDir();
+            File testsDir = ((IDeviceBuildInfo) buildInfo).getTestsDir();
             if (testsDir != null && testsDir.exists()) {
                 expandedTestDirs.add(FileUtil.getFileForPath(testsDir, "DATA", "app"));
                 expandedTestDirs.add(FileUtil.getFileForPath(testsDir, "DATA", "app", apkBase));
-                expandedTestDirs.add(FileUtil.getFileForPath(
-                    testsDir, "DATA", "priv-app", apkBase));
+                expandedTestDirs.add(
+                        FileUtil.getFileForPath(testsDir, "DATA", "priv-app", apkBase));
+                // Files in testcases directory will be in base build info tests dir.
+                expandedTestDirs.add(FileUtil.findFile(testsDir, apkBase));
             }
         }
         if (altDirBehavior == null) {
@@ -97,8 +97,9 @@
         }
 
         for (File dir : dirs) {
-            File testAppFile = new File(dir, apkFileName);
-            if (testAppFile.exists()) {
+            // Recursively search each folder
+            File testAppFile = FileUtil.findFile(dir, apkFileName);
+            if (testAppFile != null && testAppFile.exists()) {
                 return testAppFile;
             }
         }
diff --git a/src/com/android/tradefed/util/FileUtil.java b/src/com/android/tradefed/util/FileUtil.java
index b537a1b..2f87111 100644
--- a/src/com/android/tradefed/util/FileUtil.java
+++ b/src/com/android/tradefed/util/FileUtil.java
@@ -234,12 +234,14 @@
      * Internal helper to determine if 'chmod' is available on the system OS.
      */
     protected static boolean chmodExists() {
-        CommandResult result = RunUtil.getDefault().runTimedCmd(10 * 1000, sChmod);
+        // Silence the scary process exception when chmod is missing, we will log instead.
+        CommandResult result = RunUtil.getDefault().runTimedCmdSilently(10 * 1000, sChmod);
         // We expect a status fail because 'chmod' requires arguments.
         if (CommandStatus.FAILED.equals(result.getStatus()) &&
                 result.getStderr().contains("chmod: missing operand")) {
             return true;
         }
+        CLog.w("Chmod is not supported by this OS.");
         return false;
     }
 
@@ -392,15 +394,10 @@
      * @throws IOException if failed to hardlink file
      */
     public static void hardlinkFile(File origFile, File destFile) throws IOException {
-        if (!origFile.exists()) {
-            throw new IOException(String.format("Cannot hardlink %s. File does not exist",
-                    origFile.getAbsolutePath()));
-        }
         // `ln src dest` will create a hardlink (note: not `ln -s src dest`, which creates symlink)
         // note that this will fail across filesystem boundaries
         // FIXME: should probably just fall back to normal copy if this fails
-        CommandResult result = RunUtil.getDefault().runTimedCmd(10 * 1000, "ln",
-                origFile.getAbsolutePath(), destFile.getAbsolutePath());
+        CommandResult result = linkFile(origFile, destFile, false);
         if (!result.getStatus().equals(CommandStatus.SUCCESS)) {
             throw new IOException(String.format(
                     "Failed to hardlink %s to %s.  Across filesystem boundary?",
@@ -409,6 +406,40 @@
     }
 
     /**
+     * A helper method that simlinks a file to another file
+     *
+     * @param origFile the original file
+     * @param destFile the destination file
+     * @throws IOException if failed to simlink file
+     */
+    public static void simlinkFile(File origFile, File destFile) throws IOException {
+        CommandResult res = linkFile(origFile, destFile, true);
+        if (!CommandStatus.SUCCESS.equals(res.getStatus())) {
+            throw new IOException("Error trying to simlink: " + res.getStderr());
+        }
+    }
+
+    private static CommandResult linkFile(File origFile, File destFile, boolean simlink)
+            throws IOException {
+        if (!origFile.exists()) {
+            String link = simlink ? "simlink" : "hardlink";
+            throw new IOException(
+                    String.format(
+                            "Cannot %s %s. File does not exist", link, origFile.getAbsolutePath()));
+        }
+        List<String> cmd = new ArrayList<>();
+        cmd.add("ln");
+        if (simlink) {
+            cmd.add("-s");
+        }
+        cmd.add(origFile.getAbsolutePath());
+        cmd.add(destFile.getAbsolutePath());
+        CommandResult result =
+                RunUtil.getDefault().runTimedCmd(10 * 1000, cmd.toArray(new String[0]));
+        return result;
+    }
+
+    /**
      * Recursively hardlink folder contents.
      * <p/>
      * Only supports copying of files and directories - symlinks are not copied. If the destination
@@ -434,6 +465,31 @@
     }
 
     /**
+     * Recursively simlink folder contents.
+     *
+     * <p>Only supports copying of files and directories - symlinks are not copied. If the
+     * destination directory does not exist, it will be created.
+     *
+     * @param sourceDir the folder that contains the files to copy
+     * @param destDir the destination folder
+     * @throws IOException
+     */
+    public static void recursiveSimlink(File sourceDir, File destDir) throws IOException {
+        if (!destDir.isDirectory() && !destDir.mkdir()) {
+            throw new IOException(
+                    String.format("Could not create directory %s", destDir.getAbsolutePath()));
+        }
+        for (File childFile : sourceDir.listFiles()) {
+            File destChild = new File(destDir, childFile.getName());
+            if (childFile.isDirectory()) {
+                recursiveSimlink(childFile, destChild);
+            } else if (childFile.isFile()) {
+                simlinkFile(childFile, destChild);
+            }
+        }
+    }
+
+    /**
      * A helper method that copies a file's contents to a local file
      *
      * @param origFile the original file to be copied
diff --git a/src/com/android/tradefed/util/MultiMap.java b/src/com/android/tradefed/util/MultiMap.java
index 1123e6d..8ed39e4 100644
--- a/src/com/android/tradefed/util/MultiMap.java
+++ b/src/com/android/tradefed/util/MultiMap.java
@@ -15,6 +15,8 @@
  */
 package com.android.tradefed.util;
 
+import com.android.tradefed.build.BuildSerializedVersion;
+
 import com.google.common.base.Objects;
 
 import java.io.Serializable;
@@ -27,6 +29,7 @@
 /** A {@link Map} that supports multiple values per key. */
 public class MultiMap<K, V> implements Serializable {
 
+    private static final long serialVersionUID = BuildSerializedVersion.VERSION;
     private final Map<K, List<V>> mInternalMap;
 
     public MultiMap() {
diff --git a/src/com/android/tradefed/util/SystemUtil.java b/src/com/android/tradefed/util/SystemUtil.java
index 3b55662..b1e9284 100644
--- a/src/com/android/tradefed/util/SystemUtil.java
+++ b/src/com/android/tradefed/util/SystemUtil.java
@@ -16,6 +16,8 @@
 
 package com.android.tradefed.util;
 
+import com.android.tradefed.build.IBuildInfo;
+import com.android.tradefed.build.IDeviceBuildInfo;
 import com.android.tradefed.log.LogUtil.CLog;
 
 import com.google.common.annotations.VisibleForTesting;
@@ -42,6 +44,9 @@
 
     static final String ENV_ANDROID_PRODUCT_OUT = "ANDROID_PRODUCT_OUT";
 
+    private static final String HOST_TESTCASES = "host/testcases";
+    private static final String TARGET_TESTCASES = "target/testcases";
+
     /**
      * Get the value of an environment variable.
      *
@@ -55,13 +60,8 @@
         return System.getenv(name);
     }
 
-    /**
-     * Get a list of {@link File} of the test cases directories
-     *
-     * @return a list of {@link File} of directories of the test cases folder of build output, based
-     *     on the value of environment variables.
-     */
-    public static List<File> getTestCasesDirs() {
+    /** Get a list of {@link File} pointing to tests directories external to Tradefed. */
+    public static List<File> getExternalTestCasesDirs() {
         List<File> testCasesDirs = new ArrayList<File>();
         // TODO(b/36782030): Add ENV_ANDROID_HOST_OUT_TESTCASES back to the list.
         Set<String> testCasesDirNames =
@@ -84,6 +84,38 @@
     }
 
     /**
+     * Get a list of {@link File} of the test cases directories
+     *
+     * @param buildInfo the build artifact information. Set it to null if build info is not
+     *     available or there is no need to get test cases directories from build info.
+     * @return a list of {@link File} of directories of the test cases folder of build output, based
+     *     on the value of environment variables and the given build info.
+     */
+    public static List<File> getTestCasesDirs(IBuildInfo buildInfo) {
+        List<File> testCasesDirs = new ArrayList<File>();
+        testCasesDirs.addAll(getExternalTestCasesDirs());
+
+        // TODO: Remove this logic after Versioned TF V2 is implemented, in which staging build
+        // artifact will be done by the parent process, and the test cases dirs will be set by
+        // environment variables.
+        // Add tests dir from build info.
+        if (buildInfo instanceof IDeviceBuildInfo) {
+            IDeviceBuildInfo deviceBuildInfo = (IDeviceBuildInfo) buildInfo;
+            File testsDir = deviceBuildInfo.getTestsDir();
+            // Add all possible paths to the testcases directory list.
+            if (testsDir != null) {
+                testCasesDirs.addAll(
+                        Arrays.asList(
+                                testsDir,
+                                FileUtil.getFileForPath(testsDir, HOST_TESTCASES),
+                                FileUtil.getFileForPath(testsDir, TARGET_TESTCASES)));
+            }
+        }
+
+        return testCasesDirs;
+    }
+
+    /**
      * Gets the product specific output dir from an Android build tree. Typically this location
      * contains images for various device partitions, bootloader, radio and so on.
      *
diff --git a/src/com/android/tradefed/util/TarUtil.java b/src/com/android/tradefed/util/TarUtil.java
index 62b3ffd..d9dfb2a 100644
--- a/src/com/android/tradefed/util/TarUtil.java
+++ b/src/com/android/tradefed/util/TarUtil.java
@@ -37,6 +37,7 @@
 import java.util.LinkedList;
 import java.util.List;
 import java.util.zip.GZIPInputStream;
+import java.util.zip.GZIPOutputStream;
 
 /**
  * Utility to manipulate a tar file. It wraps the commons-compress in order to provide tar support.
@@ -126,6 +127,32 @@
     }
 
     /**
+     * Utility function to gzip (.gz) a file. the .gz extension will be added to base file name.
+     *
+     * @param inputFile the {@link File} to be gzipped.
+     * @return the gzipped file.
+     * @throws IOException
+     */
+    public static File gzip(final File inputFile) throws IOException {
+        File outputFile = FileUtil.createTempFile(inputFile.getName(), ".gz");
+        GZIPOutputStream out = null;
+        FileInputStream in = null;
+        try {
+            out = new GZIPOutputStream(new FileOutputStream(outputFile));
+            in = new FileInputStream(inputFile);
+            IOUtils.copy(in, out);
+        } catch (IOException e) {
+            // delete the tmp file if we failed to gzip.
+            FileUtil.deleteFile(outputFile);
+            throw e;
+        } finally {
+            StreamUtil.close(in);
+            StreamUtil.close(out);
+        }
+        return outputFile;
+    }
+
+    /**
      * Helper to extract and log to the reporters a tar gz file and its content
      *
      * @param listener the {@link ITestLogger} where to log the files.
diff --git a/src/com/android/tradefed/util/UniqueMultiMap.java b/src/com/android/tradefed/util/UniqueMultiMap.java
index a9d190d..f53b9d2 100644
--- a/src/com/android/tradefed/util/UniqueMultiMap.java
+++ b/src/com/android/tradefed/util/UniqueMultiMap.java
@@ -15,6 +15,8 @@
  */
 package com.android.tradefed.util;
 
+import com.android.tradefed.build.BuildSerializedVersion;
+
 import java.util.Collection;
 
 /**
@@ -25,6 +27,8 @@
  */
 public class UniqueMultiMap<K, V> extends MultiMap<K, V> {
 
+    private static final long serialVersionUID = BuildSerializedVersion.VERSION;
+
     @Override
     public V put(K key, V value) {
         Collection<V> values = get(key);
diff --git a/tests/res/testdata/SmallRawImage.raw b/tests/res/testdata/SmallRawImage.raw
new file mode 100644
index 0000000..6cacd5f
--- /dev/null
+++ b/tests/res/testdata/SmallRawImage.raw
Binary files differ
diff --git a/tests/src/com/android/tradefed/FuncTests.java b/tests/src/com/android/tradefed/FuncTests.java
index 9df3455..e6054d5 100644
--- a/tests/src/com/android/tradefed/FuncTests.java
+++ b/tests/src/com/android/tradefed/FuncTests.java
@@ -21,45 +21,40 @@
 import com.android.tradefed.device.TestDeviceFuncTest;
 import com.android.tradefed.targetprep.AppSetupFuncTest;
 import com.android.tradefed.targetprep.DeviceSetupFuncTest;
-import com.android.tradefed.testtype.DeviceTestSuite;
+import com.android.tradefed.testtype.DeviceSuite;
 import com.android.tradefed.testtype.InstrumentationTestFuncTest;
 import com.android.tradefed.util.FileUtilFuncTest;
 import com.android.tradefed.util.RunUtilFuncTest;
 import com.android.tradefed.util.net.HttpHelperFuncTest;
 
-import junit.framework.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Suite.SuiteClasses;
 
 /**
  * A test suite for all Trade Federation functional tests.
- * <p/>
- * This suite requires a device.
+ *
+ * <p>This suite requires a device.
  */
-public class FuncTests extends DeviceTestSuite {
-
-    public FuncTests() {
-        super();
-        // build
-        this.addTestSuite(FileDownloadCacheFuncTest.class);
-        // command
-        this.addTestSuite(CommandSchedulerFuncTest.class);
-        // command.remote
-        this.addTestSuite(RemoteManagerFuncTest.class);
-        // device
-        this.addTestSuite(TestDeviceFuncTest.class);
-        // targetprep
-        this.addTestSuite(AppSetupFuncTest.class);
-        this.addTestSuite(DeviceSetupFuncTest.class);
-        // testtype
-        this.addTestSuite(InstrumentationTestFuncTest.class);
-        // util
-        this.addTestSuite(FileUtilFuncTest.class);
-        // TODO: temporarily remove from suite until we figure out how to install gtest data
-        //this.addTestSuite(GTestFuncTest.class);
-        this.addTestSuite(HttpHelperFuncTest.class);
-        this.addTestSuite(RunUtilFuncTest.class);
-    }
-
-    public static Test suite() {
-        return new FuncTests();
-    }
-}
+@RunWith(DeviceSuite.class)
+@SuiteClasses({
+    // build
+    FileDownloadCacheFuncTest.class,
+    // command
+    CommandSchedulerFuncTest.class,
+    // command.remote
+    RemoteManagerFuncTest.class,
+    // device
+    TestDeviceFuncTest.class,
+    // targetprep
+    AppSetupFuncTest.class,
+    DeviceSetupFuncTest.class,
+    // testtype
+    InstrumentationTestFuncTest.class,
+    // util
+    FileUtilFuncTest.class,
+    // TODO: temporarily remove from suite until we figure out how to install gtest data
+    //this.addTestSuite(GTestFuncTest.class);
+    HttpHelperFuncTest.class,
+    RunUtilFuncTest.class,
+})
+public class FuncTests {}
diff --git a/tests/src/com/android/tradefed/UnitTests.java b/tests/src/com/android/tradefed/UnitTests.java
index 62ff479..240e4c7 100644
--- a/tests/src/com/android/tradefed/UnitTests.java
+++ b/tests/src/com/android/tradefed/UnitTests.java
@@ -15,6 +15,7 @@
  */
 package com.android.tradefed;
 
+import com.android.tradefed.build.BootstrapBuildProviderTest;
 import com.android.tradefed.build.BuildInfoTest;
 import com.android.tradefed.build.DeviceBuildDescriptorTest;
 import com.android.tradefed.build.DeviceBuildInfoTest;
@@ -105,6 +106,7 @@
 import com.android.tradefed.targetprep.DeviceSetupTest;
 import com.android.tradefed.targetprep.FastbootDeviceFlasherTest;
 import com.android.tradefed.targetprep.FlashingResourcesParserTest;
+import com.android.tradefed.targetprep.InstallAllTestZipAppsSetupTest;
 import com.android.tradefed.targetprep.InstallApkSetupTest;
 import com.android.tradefed.targetprep.InstrumentationPreparerTest;
 import com.android.tradefed.targetprep.KernelFlashPreparerTest;
@@ -152,6 +154,7 @@
 import com.android.tradefed.testtype.suite.ITestSuiteTest;
 import com.android.tradefed.testtype.suite.ModuleDefinitionTest;
 import com.android.tradefed.testtype.suite.ModuleListenerTest;
+import com.android.tradefed.testtype.suite.ModuleMergerTest;
 import com.android.tradefed.testtype.suite.ModuleSplitterTest;
 import com.android.tradefed.testtype.suite.TestFailureListenerTest;
 import com.android.tradefed.testtype.suite.TfSuiteRunnerTest;
@@ -227,6 +230,7 @@
 @SuiteClasses({
 
     // build
+    BootstrapBuildProviderTest.class,
     BuildInfoTest.class,
     DeviceBuildInfoTest.class,
     DeviceBuildDescriptorTest.class,
@@ -336,6 +340,7 @@
     DeviceSetupTest.class,
     FastbootDeviceFlasherTest.class,
     FlashingResourcesParserTest.class,
+    InstallAllTestZipAppsSetupTest.class,
     InstallApkSetupTest.class,
     InstrumentationPreparerTest.class,
     KernelFlashPreparerTest.class,
@@ -394,6 +399,7 @@
     ITestSuiteTest.class,
     ModuleDefinitionTest.class,
     ModuleListenerTest.class,
+    ModuleMergerTest.class,
     ModuleSplitterTest.class,
     TestFailureListenerTest.class,
     TfSuiteRunnerTest.class,
diff --git a/tests/src/com/android/tradefed/build/BootstrapBuildProviderTest.java b/tests/src/com/android/tradefed/build/BootstrapBuildProviderTest.java
new file mode 100644
index 0000000..a533655
--- /dev/null
+++ b/tests/src/com/android/tradefed/build/BootstrapBuildProviderTest.java
@@ -0,0 +1,60 @@
+/*
+ * Copyright (C) 2017 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.build;
+
+import static org.junit.Assert.*;
+
+import com.android.tradefed.device.ITestDevice;
+
+import org.easymock.EasyMock;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Unit tests for {@link BootstrapBuildProvider}. */
+@RunWith(JUnit4.class)
+public class BootstrapBuildProviderTest {
+    private BootstrapBuildProvider mProvider;
+    private ITestDevice mMockDevice;
+
+    @Before
+    public void setUp() {
+        mProvider = new BootstrapBuildProvider();
+        mMockDevice = EasyMock.createMock(ITestDevice.class);
+    }
+
+    @Test
+    public void testGetBuild() throws Exception {
+        EasyMock.expect(mMockDevice.getBuildId()).andReturn("5");
+        EasyMock.expect(mMockDevice.waitForDeviceShell(EasyMock.anyLong())).andReturn(true);
+        EasyMock.expect(mMockDevice.getProperty(EasyMock.anyObject())).andStubReturn("property");
+        EasyMock.expect(mMockDevice.getProductVariant()).andStubReturn("variant");
+        EasyMock.expect(mMockDevice.getBuildFlavor()).andStubReturn("flavor");
+        EasyMock.expect(mMockDevice.getBuildAlias()).andStubReturn("alias");
+        EasyMock.replay(mMockDevice);
+        IBuildInfo res = mProvider.getBuild(mMockDevice);
+        assertNotNull(res);
+        try {
+            assertTrue(res instanceof IDeviceBuildInfo);
+            // Ensure tests dir is never null
+            assertTrue(((IDeviceBuildInfo) res).getTestsDir() != null);
+            EasyMock.verify(mMockDevice);
+        } finally {
+            mProvider.cleanUp(res);
+        }
+    }
+}
diff --git a/tests/src/com/android/tradefed/build/SdkBuildInfoTest.java b/tests/src/com/android/tradefed/build/SdkBuildInfoTest.java
index 64ca732..ecde698 100644
--- a/tests/src/com/android/tradefed/build/SdkBuildInfoTest.java
+++ b/tests/src/com/android/tradefed/build/SdkBuildInfoTest.java
@@ -36,12 +36,15 @@
     @Override
     protected void setUp() throws Exception {
         mMockRunUtil = EasyMock.createMock(IRunUtil.class);
-        mSdkBuild = new SdkBuildInfo() {
-            @Override
-            IRunUtil getRunUtil() {
-                return mMockRunUtil;
-            }
-        };
+        mSdkBuild =
+                new SdkBuildInfo() {
+                    private static final long serialVersionUID = BuildSerializedVersion.VERSION;
+
+                    @Override
+                    IRunUtil getRunUtil() {
+                        return mMockRunUtil;
+                    }
+                };
         mSdkBuild.setSdkDir(new File("tmp"));
     }
     /**
diff --git a/tests/src/com/android/tradefed/command/CommandRunnerTest.java b/tests/src/com/android/tradefed/command/CommandRunnerTest.java
index fb906af..c73410e 100644
--- a/tests/src/com/android/tradefed/command/CommandRunnerTest.java
+++ b/tests/src/com/android/tradefed/command/CommandRunnerTest.java
@@ -18,10 +18,13 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.doReturn;
 
 import com.android.tradefed.command.CommandRunner.ExitCode;
 import com.android.tradefed.config.ConfigurationException;
 import com.android.tradefed.config.GlobalConfiguration;
+import com.android.tradefed.device.MockDeviceManager;
 import com.android.tradefed.util.FileUtil;
 
 import org.junit.After;
@@ -29,6 +32,7 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
 
 import java.io.ByteArrayOutputStream;
 import java.io.File;
@@ -210,4 +214,43 @@
                 "Stack does not contain expected message: " + mStackTraceOutput,
                 mStackTraceOutput.contains(FAKE_CONFIG));
     }
+
+    /** Test that if the device is not allocated after a timeout, we throw a NoDeviceException. */
+    @Test
+    public void testRun_noDevice() throws Exception {
+        CommandScheduler mockScheduler = Mockito.spy(CommandScheduler.class);
+        CommandRunner mRunner =
+                new TestableCommandRunner() {
+                    @Override
+                    long getCheckDeviceTimeout() {
+                        return 200l;
+                    }
+
+                    @Override
+                    ICommandScheduler getCommandScheduler() {
+                        return mockScheduler;
+                    }
+                };
+        String[] args = {
+            mConfig.getAbsolutePath(),
+            "-s",
+            "impossibleSerialThatWillNotBeFound",
+            "--log-file-path",
+            mLogDir.getAbsolutePath()
+        };
+        doNothing().when(mockScheduler).initDeviceManager();
+        doReturn(new MockDeviceManager(1)).when(mockScheduler).getDeviceManager();
+        doNothing().when(mockScheduler).shutdownOnEmpty();
+        doNothing().when(mockScheduler).initLogging();
+        doNothing().when(mockScheduler).cleanUp();
+        mRunner.run(args);
+        Mockito.verify(mockScheduler).shutdownOnEmpty();
+        mockScheduler.join(5000);
+        assertEquals(ExitCode.NO_DEVICE_ALLOCATED, mRunner.getErrorCode());
+        assertTrue(
+                String.format("%s does not contains the expected output", mStackTraceOutput),
+                mStackTraceOutput.contains(
+                        "com.android.tradefed.device.NoDeviceException: No device was allocated "
+                                + "for the command."));
+    }
 }
diff --git a/tests/src/com/android/tradefed/config/ConfigurationFactoryTest.java b/tests/src/com/android/tradefed/config/ConfigurationFactoryTest.java
index 54ee993..3a94edd 100644
--- a/tests/src/com/android/tradefed/config/ConfigurationFactoryTest.java
+++ b/tests/src/com/android/tradefed/config/ConfigurationFactoryTest.java
@@ -1371,7 +1371,7 @@
         File tmpDir = externalConfig.getParentFile();
 
         ConfigurationFactory spyFactory = Mockito.spy(mFactory);
-        Mockito.doReturn(Arrays.asList(tmpDir)).when(spyFactory).getTestCasesDirs();
+        Mockito.doReturn(Arrays.asList(tmpDir)).when(spyFactory).getExternalTestCasesDirs();
 
         try {
             File config = spyFactory.getTestCaseConfigPath(configName);
@@ -1389,10 +1389,10 @@
         File tmpDir = FileUtil.createTempDir("config-check-var");
         try {
             ConfigurationFactory spyFactory = Mockito.spy(mFactory);
-            Mockito.doReturn(Arrays.asList(tmpDir)).when(spyFactory).getTestCasesDirs();
+            Mockito.doReturn(Arrays.asList(tmpDir)).when(spyFactory).getExternalTestCasesDirs();
             File config = spyFactory.getTestCaseConfigPath("non-exist");
             assertNull(config);
-            Mockito.verify(spyFactory, Mockito.times(1)).getTestCasesDirs();
+            Mockito.verify(spyFactory, Mockito.times(1)).getExternalTestCasesDirs();
         } finally {
             FileUtil.recursiveDelete(tmpDir);
         }
@@ -1409,7 +1409,7 @@
             File subDir = FileUtil.createTempDir("subdir", tmpDir);
             FileUtil.createTempFile("testconfig2", ".xml", subDir);
             ConfigurationFactory spyFactory = Mockito.spy(mFactory);
-            Mockito.doReturn(Arrays.asList(tmpDir)).when(spyFactory).getTestCasesDirs();
+            Mockito.doReturn(Arrays.asList(tmpDir)).when(spyFactory).getExternalTestCasesDirs();
             // looking at full path we get both configs
             Set<String> res = spyFactory.getConfigNamesFromTestCases(null);
             assertEquals(2, res.size());
diff --git a/tests/src/com/android/tradefed/config/ConfigurationTest.java b/tests/src/com/android/tradefed/config/ConfigurationTest.java
index d16d5d4..4a108c8 100644
--- a/tests/src/com/android/tradefed/config/ConfigurationTest.java
+++ b/tests/src/com/android/tradefed/config/ConfigurationTest.java
@@ -26,6 +26,7 @@
 import com.android.tradefed.device.IDeviceSelection;
 import com.android.tradefed.invoker.InvocationContext;
 import com.android.tradefed.log.ILeveledLogOutput;
+import com.android.tradefed.log.LogUtil.CLog;
 import com.android.tradefed.result.ITestInvocationListener;
 import com.android.tradefed.result.TextResultReporter;
 import com.android.tradefed.targetprep.ITargetPreparer;
@@ -637,6 +638,7 @@
             String content = FileUtil.readStringFromFile(test);
             assertTrue(content.length() > 100);
             assertTrue(content.contains("<configuration>"));
+            CLog.e("%s", content);
         } finally {
             FileUtil.deleteFile(test);
         }
diff --git a/tests/src/com/android/tradefed/device/DeviceManagerTest.java b/tests/src/com/android/tradefed/device/DeviceManagerTest.java
index 16dfbb9..41acc6b 100644
--- a/tests/src/com/android/tradefed/device/DeviceManagerTest.java
+++ b/tests/src/com/android/tradefed/device/DeviceManagerTest.java
@@ -40,6 +40,7 @@
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.io.PrintWriter;
+import java.util.ArrayList;
 import java.util.List;
 
 /**
@@ -115,7 +116,6 @@
         public int waitFor() throws InterruptedException {
             return 0;
         }
-
     }
 
     /**
@@ -890,6 +890,197 @@
     }
 
     /**
+     * Test freeing a device that was unable but showing in adb devices. Device will become
+     * Unavailable but still seen by the DeviceManager.
+     */
+    public void testFreeDevice_unavailable() {
+        EasyMock.expect(mMockIDevice.isEmulator()).andStubReturn(Boolean.FALSE);
+        EasyMock.expect(mMockIDevice.getState()).andReturn(DeviceState.ONLINE);
+        EasyMock.expect(mMockStateMonitor.waitForDeviceShell(EasyMock.anyLong()))
+                .andReturn(Boolean.TRUE);
+        mMockStateMonitor.setState(TestDeviceState.NOT_AVAILABLE);
+
+        CommandResult stubAdbDevices = new CommandResult(CommandStatus.SUCCESS);
+        stubAdbDevices.setStdout("List of devices attached\nserial\tdevice\n");
+        EasyMock.expect(
+                        mMockRunUtil.runTimedCmd(
+                                EasyMock.anyLong(), EasyMock.eq("adb"), EasyMock.eq("devices")))
+                .andReturn(stubAdbDevices);
+
+        replayMocks();
+        IManagedTestDevice testDevice = new TestDevice(mMockIDevice, mMockStateMonitor, null);
+        DeviceManager manager = createDeviceManagerNoInit();
+        manager.init(
+                null,
+                null,
+                new ManagedTestDeviceFactory(false, null, null) {
+                    @Override
+                    public IManagedTestDevice createDevice(IDevice idevice) {
+                        mMockTestDevice.setIDevice(idevice);
+                        return testDevice;
+                    }
+
+                    @Override
+                    protected CollectingOutputReceiver createOutputReceiver() {
+                        return new CollectingOutputReceiver() {
+                            @Override
+                            public String getOutput() {
+                                return "/system/bin/pm";
+                            }
+                        };
+                    }
+
+                    @Override
+                    public void setFastbootEnabled(boolean enable) {
+                        // ignore
+                    }
+                });
+
+        mDeviceListener.deviceConnected(mMockIDevice);
+
+        IManagedTestDevice device = (IManagedTestDevice) manager.allocateDevice();
+        assertNotNull(device);
+        // device becomes unavailable
+        device.setDeviceState(TestDeviceState.NOT_AVAILABLE);
+        // a freed 'unavailable' device becomes UNAVAILABLE state
+        manager.freeDevice(device, FreeDeviceState.UNAVAILABLE);
+        // ensure device cannot be allocated again
+        ITestDevice device2 = manager.allocateDevice();
+        assertNull(device2);
+        verifyMocks();
+        // We still have the device in the list
+        assertEquals(1, manager.getDeviceList().size());
+    }
+
+    /**
+     * Test that when freeing an Unavailable device that is not in 'adb devices' we correctly remove
+     * it from our tracking list.
+     */
+    public void testFreeDevice_unknown() {
+        EasyMock.expect(mMockIDevice.isEmulator()).andStubReturn(Boolean.FALSE);
+        EasyMock.expect(mMockIDevice.getState()).andReturn(DeviceState.ONLINE);
+        EasyMock.expect(mMockStateMonitor.waitForDeviceShell(EasyMock.anyLong()))
+                .andReturn(Boolean.TRUE);
+        mMockStateMonitor.setState(TestDeviceState.NOT_AVAILABLE);
+
+        CommandResult stubAdbDevices = new CommandResult(CommandStatus.SUCCESS);
+        // device serial is not in the list
+        stubAdbDevices.setStdout("List of devices attached\n");
+        EasyMock.expect(
+                        mMockRunUtil.runTimedCmd(
+                                EasyMock.anyLong(), EasyMock.eq("adb"), EasyMock.eq("devices")))
+                .andReturn(stubAdbDevices);
+
+        replayMocks();
+        IManagedTestDevice testDevice = new TestDevice(mMockIDevice, mMockStateMonitor, null);
+        DeviceManager manager = createDeviceManagerNoInit();
+        manager.init(
+                null,
+                null,
+                new ManagedTestDeviceFactory(false, null, null) {
+                    @Override
+                    public IManagedTestDevice createDevice(IDevice idevice) {
+                        mMockTestDevice.setIDevice(idevice);
+                        return testDevice;
+                    }
+
+                    @Override
+                    protected CollectingOutputReceiver createOutputReceiver() {
+                        return new CollectingOutputReceiver() {
+                            @Override
+                            public String getOutput() {
+                                return "/system/bin/pm";
+                            }
+                        };
+                    }
+
+                    @Override
+                    public void setFastbootEnabled(boolean enable) {
+                        // ignore
+                    }
+                });
+
+        mDeviceListener.deviceConnected(mMockIDevice);
+
+        IManagedTestDevice device = (IManagedTestDevice) manager.allocateDevice();
+        assertNotNull(device);
+        // device becomes unavailable
+        device.setDeviceState(TestDeviceState.NOT_AVAILABLE);
+        // a freed 'unavailable' device becomes UNAVAILABLE state
+        manager.freeDevice(device, FreeDeviceState.UNAVAILABLE);
+        // ensure device cannot be allocated again
+        ITestDevice device2 = manager.allocateDevice();
+        assertNull(device2);
+        verifyMocks();
+        // We have 0 device in the list since it was removed
+        assertEquals(0, manager.getDeviceList().size());
+    }
+
+    /**
+     * Test that when freeing an Unavailable device that is not in 'adb devices' we correctly remove
+     * it from our tracking list even if its serial is a substring of another serial.
+     */
+    public void testFreeDevice_unknown_subName() {
+        EasyMock.expect(mMockIDevice.isEmulator()).andStubReturn(Boolean.FALSE);
+        EasyMock.expect(mMockIDevice.getState()).andReturn(DeviceState.ONLINE);
+        EasyMock.expect(mMockStateMonitor.waitForDeviceShell(EasyMock.anyLong()))
+                .andReturn(Boolean.TRUE);
+        mMockStateMonitor.setState(TestDeviceState.NOT_AVAILABLE);
+
+        CommandResult stubAdbDevices = new CommandResult(CommandStatus.SUCCESS);
+        // device serial is not in the list
+        stubAdbDevices.setStdout("List of devices attached\n2serial\tdevice\n");
+        EasyMock.expect(
+                        mMockRunUtil.runTimedCmd(
+                                EasyMock.anyLong(), EasyMock.eq("adb"), EasyMock.eq("devices")))
+                .andReturn(stubAdbDevices);
+
+        replayMocks();
+        IManagedTestDevice testDevice = new TestDevice(mMockIDevice, mMockStateMonitor, null);
+        DeviceManager manager = createDeviceManagerNoInit();
+        manager.init(
+                null,
+                null,
+                new ManagedTestDeviceFactory(false, null, null) {
+                    @Override
+                    public IManagedTestDevice createDevice(IDevice idevice) {
+                        mMockTestDevice.setIDevice(idevice);
+                        return testDevice;
+                    }
+
+                    @Override
+                    protected CollectingOutputReceiver createOutputReceiver() {
+                        return new CollectingOutputReceiver() {
+                            @Override
+                            public String getOutput() {
+                                return "/system/bin/pm";
+                            }
+                        };
+                    }
+
+                    @Override
+                    public void setFastbootEnabled(boolean enable) {
+                        // ignore
+                    }
+                });
+
+        mDeviceListener.deviceConnected(mMockIDevice);
+
+        IManagedTestDevice device = (IManagedTestDevice) manager.allocateDevice();
+        assertNotNull(device);
+        // device becomes unavailable
+        device.setDeviceState(TestDeviceState.NOT_AVAILABLE);
+        // a freed 'unavailable' device becomes UNAVAILABLE state
+        manager.freeDevice(device, FreeDeviceState.UNAVAILABLE);
+        // ensure device cannot be allocated again
+        ITestDevice device2 = manager.allocateDevice();
+        assertNull(device2);
+        verifyMocks();
+        // We have 0 device in the list since it was removed
+        assertEquals(0, manager.getDeviceList().size());
+    }
+
+    /**
      * Helper to set the expectation when a {@link DeviceDescriptor} is expected.
      */
     private void setDeviceDescriptorExpectation() {
@@ -942,6 +1133,10 @@
                 + "\n", out.toString());
     }
 
+    /**
+     * Test that {@link DeviceManager#shouldAdbBridgeBeRestarted()} properly reports the flag state
+     * based on if it was requested or not.
+     */
     public void testAdbBridgeFlag() throws Exception {
         setCheckAvailableDeviceExpectations();
         replayMocks();
@@ -955,4 +1150,24 @@
 
         verifyMocks();
     }
+
+    /**
+     * Test that when a {@link IDeviceMonitor} is available in {@link DeviceManager} it properly
+     * goes through its life cycle.
+     */
+    public void testDeviceMonitorLifeCyle() throws Exception {
+        IDeviceMonitor mockMonitor = EasyMock.createMock(IDeviceMonitor.class);
+        List<IDeviceMonitor> monitors = new ArrayList<>();
+        monitors.add(mockMonitor);
+        setCheckAvailableDeviceExpectations();
+
+        mockMonitor.setDeviceLister(EasyMock.anyObject());
+        mockMonitor.run();
+        mockMonitor.stop();
+
+        replayMocks(mockMonitor);
+        DeviceManager manager = createDeviceManager(monitors, mMockIDevice);
+        manager.terminateDeviceMonitor();
+        verifyMocks(mockMonitor);
+    }
 }
diff --git a/tests/src/com/android/tradefed/device/TestDeviceFuncTest.java b/tests/src/com/android/tradefed/device/TestDeviceFuncTest.java
index 11307c9..ae4f1c7 100644
--- a/tests/src/com/android/tradefed/device/TestDeviceFuncTest.java
+++ b/tests/src/com/android/tradefed/device/TestDeviceFuncTest.java
@@ -15,6 +15,8 @@
  */
 package com.android.tradefed.device;
 
+import static org.junit.Assert.*;
+
 import com.android.ddmlib.IDevice;
 import com.android.ddmlib.Log;
 import com.android.ddmlib.testrunner.RemoteAndroidTestRunner;
@@ -24,13 +26,18 @@
 import com.android.tradefed.result.CollectingTestListener;
 import com.android.tradefed.result.FileInputStreamSource;
 import com.android.tradefed.result.InputStreamSource;
-import com.android.tradefed.testtype.DeviceTestCase;
+import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
+import com.android.tradefed.testtype.IDeviceTest;
 import com.android.tradefed.util.CommandStatus;
 import com.android.tradefed.util.FileUtil;
+import com.android.tradefed.util.KeyguardControllerState;
 import com.android.tradefed.util.RunUtil;
 import com.android.tradefed.util.StreamUtil;
 
 import org.easymock.EasyMock;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
 
 import java.awt.image.BufferedImage;
 import java.io.BufferedInputStream;
@@ -43,10 +50,11 @@
 
 /**
  * Functional tests for {@link TestDevice}.
- * <p/>
- * Requires a physical device to be connected.
+ *
+ * <p>Requires a physical device to be connected.
  */
-public class TestDeviceFuncTest extends DeviceTestCase {
+@RunWith(DeviceJUnit4ClassRunner.class)
+public class TestDeviceFuncTest implements IDeviceTest {
 
     private static final String LOG_TAG = "TestDeviceFuncTest";
     private TestDevice mTestDevice;
@@ -55,17 +63,24 @@
     private static final int MIN_BUGREPORT_BYTES = 1024 * 1024;
 
     @Override
-    protected void setUp() throws Exception {
-        super.setUp();
-        mTestDevice = (TestDevice)getDevice();
+    public void setDevice(ITestDevice device) {
+        mTestDevice = (TestDevice) device;
+    }
+
+    @Override
+    public ITestDevice getDevice() {
+        return mTestDevice;
+    }
+
+    @Before
+    public void setUp() throws Exception {
         mMonitor = mTestDevice.getDeviceStateMonitor();
         // Ensure at set-up that the device is available.
         mTestDevice.waitForDeviceAvailable();
     }
 
-    /**
-     * Simple testcase to ensure that the grabbing a bugreport from a real TestDevice works.
-     */
+    /** Simple testcase to ensure that the grabbing a bugreport from a real TestDevice works. */
+    @Test
     public void testBugreport() throws Exception {
         InputStreamSource bugreport = mTestDevice.getBugreport();
         try {
@@ -80,9 +95,8 @@
         }
     }
 
-    /**
-     * Simple testcase to ensure that the grabbing a bugreportz from a real TestDevice works.
-     */
+    /** Simple testcase to ensure that the grabbing a bugreportz from a real TestDevice works. */
+    @Test
     public void testBugreportz() throws Exception {
         if (mTestDevice.getApiLevel() < 24) {
             CLog.i("testBugreportz() not supported by this device, skipping.");
@@ -107,11 +121,11 @@
     }
 
     /**
-     * Simple normal case test for
-     * {@link TestDevice#executeShellCommand(String)}.
-     * <p/>
-     * Do a 'shell ls' command, and verify /data and /system are listed in result.
+     * Simple normal case test for {@link TestDevice#executeShellCommand(String)}.
+     *
+     * <p>Do a 'shell ls' command, and verify /data and /system are listed in result.
      */
+    @Test
     public void testExecuteShellCommand() throws DeviceNotAvailableException {
         Log.i(LOG_TAG, "testExecuteShellCommand");
         assertSimpleShellCommand();
@@ -126,9 +140,8 @@
         assertTrue(output.contains("system"));
     }
 
-    /**
-     * Test install and uninstall of package
-     */
+    /** Test install and uninstall of package */
+    @Test
     public void testInstallUninstall() throws IOException, DeviceNotAvailableException {
         Log.i(LOG_TAG, "testInstallUninstall");
         // use the wifi util apk
@@ -159,9 +172,8 @@
         }
     }
 
-    /**
-     * Test install and uninstall of package with spaces in file name
-     */
+    /** Test install and uninstall of package with spaces in file name */
+    @Test
     public void testInstallUninstall_space() throws IOException, DeviceNotAvailableException {
         Log.i(LOG_TAG, "testInstallUninstall_space");
 
@@ -177,9 +189,8 @@
         }
     }
 
-    /**
-     * Push and then pull a file from device, and verify contents are as expected.
-     */
+    /** Push and then pull a file from device, and verify contents are as expected. */
+    @Test
     public void testPushPull_normal() throws IOException, DeviceNotAvailableException {
         Log.i(LOG_TAG, "testPushPull");
         File tmpFile = null;
@@ -214,9 +225,10 @@
 
     /**
      * Push and then pull a file from device, and verify contents are as expected.
-     * <p />
-     * This variant of the test uses "${EXTERNAL_STORAGE}" in the pathname.
+     *
+     * <p>This variant of the test uses "${EXTERNAL_STORAGE}" in the pathname.
      */
+    @Test
     public void testPushPull_extStorageVariable() throws IOException, DeviceNotAvailableException {
         Log.i(LOG_TAG, "testPushPull");
         File tmpFile = null;
@@ -246,12 +258,8 @@
             assertTrue(compareFiles(tmpFile, tmpDestFile2));
         } finally {
             FileUtil.deleteFile(tmpFile);
-            if (tmpDestFile != null) {
-                tmpDestFile.delete();
-            }
-            if (tmpDestFile2 != null) {
-                tmpDestFile2.delete();
-            }
+            FileUtil.deleteFile(tmpDestFile);
+            FileUtil.deleteFile(tmpDestFile2);
             if (deviceFilePath != null) {
                 mTestDevice.executeShellCommand(String.format("rm %s", deviceFilePath));
             }
@@ -260,9 +268,10 @@
 
     /**
      * Test pulling a file from device that does not exist.
-     * <p/>
-     * Expect {@link TestDevice#pullFile(String)} to return <code>false</code>
+     *
+     * <p>Expect {@link TestDevice#pullFile(String)} to return <code>false</code>
      */
+    @Test
     public void testPull_noexist() throws DeviceNotAvailableException {
         Log.i(LOG_TAG, "testPull_noexist");
 
@@ -277,9 +286,10 @@
 
     /**
      * Test pulling a file from device into a local file that cannot be written to.
-     * <p/>
-     * Expect {@link TestDevice#pullFile(String, File)} to return <code>false</code>
+     *
+     * <p>Expect {@link TestDevice#pullFile(String, File)} to return <code>false</code>
      */
+    @Test
     public void testPull_nopermissions() throws IOException, DeviceNotAvailableException {
         CLog.i("testPull_nopermissions");
 
@@ -306,9 +316,10 @@
 
     /**
      * Test pushing a file onto device that does not exist.
-     * <p/>
-     * Expect {@link TestDevice#pushFile(File, String)} to return <code>false</code>
+     *
+     * <p>Expect {@link TestDevice#pushFile(File, String)} to return <code>false</code>
      */
+    @Test
     public void testPush_noexist() throws DeviceNotAvailableException {
         Log.i(LOG_TAG, "testPush_noexist");
 
@@ -355,19 +366,16 @@
             }
             return true;
         } finally {
-            if (stream1 != null) {
-                stream1.close();
-            }
-            if (stream2 != null) {
-                stream2.close();
-            }
+            StreamUtil.close(stream1);
+            StreamUtil.close(stream2);
         }
     }
 
     /**
-     * Make sure that we can correctly index directories that have a symlink in the middle.  This
+     * Make sure that we can correctly index directories that have a symlink in the middle. This
      * verifies a ddmlib bugfix which added/fixed this functionality.
      */
+    @Test
     public void testListSymlinkDir() throws Exception {
         final String extStore = "/data/local";
 
@@ -394,26 +402,24 @@
         }
     }
 
-    /**
-     * Test syncing a single file using {@link TestDevice#syncFiles(File, String)}.
-     */
+    /** Test syncing a single file using {@link TestDevice#syncFiles(File, String)}. */
+    @Test
     public void testSyncFiles_normal() throws Exception {
         doTestSyncFiles(mTestDevice.getMountPoint(IDevice.MNT_EXTERNAL_STORAGE));
     }
 
     /**
      * Test syncing a single file using {@link TestDevice#syncFiles(File, String)}.
-     * <p />
-     * This variant of the test uses "${EXTERNAL_STORAGE}" in the pathname.
+     *
+     * <p>This variant of the test uses "${EXTERNAL_STORAGE}" in the pathname.
      */
+    @Test
     public void testSyncFiles_extStorageVariable() throws Exception {
         doTestSyncFiles("${EXTERNAL_STORAGE}");
     }
 
-    /**
-     * Test syncing a single file using {@link TestDevice#syncFiles(File, String)}.
-     */
-    public void doTestSyncFiles(String externalStorePath) throws Exception {
+    /** Test syncing a single file using {@link TestDevice#syncFiles(File, String)}. */
+    private void doTestSyncFiles(String externalStorePath) throws Exception {
         String expectedDeviceFilePath = null;
 
         // create temp dir with one temp file
@@ -464,9 +470,8 @@
         }
     }
 
-    /**
-     * Test pushing a directory
-     */
+    /** Test pushing a directory */
+    @Test
     public void testPushDir() throws IOException, DeviceNotAvailableException {
         String expectedDeviceFilePath = null;
         String externalStorePath = null;
@@ -494,10 +499,11 @@
 
     /**
      * Test {@link TestDevice#executeFastbootCommand(String...)} when device is in adb mode.
-     * <p/>
-     * Expect fastboot recovery to be invoked, which will boot device back to fastboot mode and
+     *
+     * <p>Expect fastboot recovery to be invoked, which will boot device back to fastboot mode and
      * command will succeed.
      */
+    @Test
     public void testExecuteFastbootCommand_deviceInAdb() throws DeviceNotAvailableException {
         Log.i(LOG_TAG, "testExecuteFastbootCommand_deviceInAdb");
         if (!mTestDevice.isFastbootEnabled()) {
@@ -521,9 +527,10 @@
 
     /**
      * Test {@link TestDevice#executeFastbootCommand(String...)} when an invalid command is passed.
-     * <p/>
-     * Expect the result indicate failure, and recovery not to be invoked.
+     *
+     * <p>Expect the result indicate failure, and recovery not to be invoked.
      */
+    @Test
     public void testExecuteFastbootCommand_badCommand() throws DeviceNotAvailableException {
         Log.i(LOG_TAG, "testExecuteFastbootCommand_badCommand");
         if (!mTestDevice.isFastbootEnabled()) {
@@ -548,9 +555,8 @@
         }
     }
 
-    /**
-     * Verify device can be rebooted into bootloader and back to adb.
-     */
+    /** Verify device can be rebooted into bootloader and back to adb. */
+    @Test
     public void testRebootIntoBootloader() throws DeviceNotAvailableException {
         Log.i(LOG_TAG, "testRebootIntoBootloader");
         if (!mTestDevice.isFastbootEnabled()) {
@@ -566,9 +572,8 @@
         }
     }
 
-    /**
-     * Verify device can be rebooted into adb.
-     */
+    /** Verify device can be rebooted into adb. */
+    @Test
     public void testReboot() throws DeviceNotAvailableException {
         Log.i(LOG_TAG, "testReboot");
         mTestDevice.reboot();
@@ -577,9 +582,8 @@
         assertTrue(mTestDevice.executeShellCommand("id").contains("root"));
     }
 
-    /**
-     * Verify device can be rebooted into adb recovery.
-     */
+    /** Verify device can be rebooted into adb recovery. */
+    @Test
     public void testRebootIntoRecovery() throws Exception {
         Log.i(LOG_TAG, "testRebootIntoRecovery");
         if (!mTestDevice.isFastbootEnabled()) {
@@ -599,12 +603,13 @@
     /**
      * Verify that {@link TestDevice#clearErrorDialogs()} can successfully clear an error dialog
      * from screen.
-     * <p/>
-     * This is done by running a test app which will crash, then running another app that
-     * does UI based tests.
-     * <p/>
-     * Assumes DevTools and TradeFedUiApp are currently installed.
+     *
+     * <p>This is done by running a test app which will crash, then running another app that does UI
+     * based tests.
+     *
+     * <p>Assumes DevTools and TradeFedUiApp are currently installed.
      */
+    @Test
     public void testClearErrorDialogs_crash() throws DeviceNotAvailableException {
         Log.i(LOG_TAG, "testClearErrorDialogs_crash");
         // Ensure device is in a known state, we doing extra care here otherwise it may be flaky
@@ -633,22 +638,28 @@
 
     /**
      * Verify the steps taken to disable keyguard after reboot are successfully
-     * <p/>
-     * This is done by rebooting then run a app that does UI based tests.
-     * <p/>
-     * Assumes DevTools and TradeFedUiApp are currently installed.
+     *
+     * <p>This is done by rebooting then run a app that does UI based tests.
+     *
+     * <p>Assumes DevTools and TradeFedUiApp are currently installed.
      */
+    @Test
     public void testDisableKeyguard() throws DeviceNotAvailableException {
         Log.i(LOG_TAG, "testDisableKeyguard");
         getDevice().reboot();
         mTestDevice.waitForDeviceAvailable();
         RunUtil.getDefault().sleep(500);
-        assertTrue(runUITests());
+        KeyguardControllerState keyguard = mTestDevice.getKeyguardState();
+        if (keyguard == null) {
+            // If the getKeyguardState is not supported.
+            assertTrue(runUITests());
+        } else {
+            assertFalse(keyguard.isKeyguardShowing());
+        }
     }
 
-    /**
-     * Test that TradeFed can successfully recover from the adb host daemon process being killed
-     */
+    /** Test that TradeFed can successfully recover from the adb host daemon process being killed */
+    @Test
     public void testExecuteShellCommand_adbKilled() {
         // FIXME: adb typically does not recover, and this causes rest of tests to fail
         //Log.i(LOG_TAG, "testExecuteShellCommand_adbKilled");
@@ -659,9 +670,11 @@
 
     /**
      * Basic test for {@link TestDevice#getScreenshot()}.
-     * <p/>
-     * Grab a screenshot, save it to a file, and perform a cursory size check to ensure its valid.
+     *
+     * <p>Grab a screenshot, save it to a file, and perform a cursory size check to ensure its
+     * valid.
      */
+    @Test
     public void testGetScreenshot() throws DeviceNotAvailableException, IOException {
         CLog.i(LOG_TAG, "testGetScreenshot");
         InputStreamSource source = getDevice().getScreenshot();
@@ -684,10 +697,11 @@
 
     /**
      * Basic test for {@link TestDevice#getLogcat(int)}.
-     * <p/>
-     * Dumps a bunch of messages to logcat, calls getLogcat(), and verifies size of capture file is
-     * equal to provided data.
+     *
+     * <p>Dumps a bunch of messages to logcat, calls getLogcat(), and verifies size of capture file
+     * is equal to provided data.
      */
+    @Test
     public void testGetLogcat_size() throws DeviceNotAvailableException, IOException {
         CLog.i(LOG_TAG, "testGetLogcat_size");
         for (int i = 0; i < 100; i++) {
@@ -716,13 +730,14 @@
 
     /**
      * Basic test for encryption if encryption is supported.
-     * <p>
-     * Calls {@link TestDevice#encryptDevice(boolean)}, {@link TestDevice#unlockDevice()}, and
+     *
+     * <p>Calls {@link TestDevice#encryptDevice(boolean)}, {@link TestDevice#unlockDevice()}, and
      * {@link TestDevice#unencryptDevice()}, as well as reboots the device while the device is
      * encrypted.
-     * </p>
+     *
      * @throws DeviceNotAvailableException
      */
+    @Test
     public void testEncryption() throws DeviceNotAvailableException {
         CLog.i("testEncryption");
 
@@ -744,18 +759,16 @@
         assertFalse(getDevice().isDeviceEncrypted());
     }
 
-    /**
-     * Test that {@link TestDevice#getProperty(String)} works after a reboot.
-     */
+    /** Test that {@link TestDevice#getProperty(String)} works after a reboot. */
+    @Test
     public void testGetProperty() throws Exception {
         assertNotNull(getDevice().getProperty("ro.hardware"));
         getDevice().rebootUntilOnline();
         assertNotNull(getDevice().getProperty("ro.hardware"));
     }
 
-    /**
-     * Test that {@link TestDevice#getProperty(String)} works for volatile properties.
-     */
+    /** Test that {@link TestDevice#getProperty(String)} works for volatile properties. */
+    @Test
     public void testGetProperty_volatile() throws Exception {
         getDevice().executeShellCommand("setprop prop.test 0");
         assertEquals("0", getDevice().getProperty("prop.test"));
@@ -763,9 +776,8 @@
         assertEquals("1", getDevice().getProperty("prop.test"));
     }
 
-    /**
-     * Test that the recovery mechanism works in {@link TestDevice#getFileEntry(String)}
-     */
+    /** Test that the recovery mechanism works in {@link TestDevice#getFileEntry(String)} */
+    @Test
     public void testGetFileEntry_recovery() throws Exception {
         if (!mTestDevice.isFastbootEnabled()) {
             Log.i(LOG_TAG, "Fastboot not enabled skipping testGetFileEntry_recovery");
@@ -792,9 +804,8 @@
         return TestAppConstants.UI_TOTAL_TESTS == uilistener.getNumTestsInState(TestStatus.PASSED);
     }
 
-    /**
-     * Test for {@link NativeDevice#setSetting(int, String, String, String)}
-     */
+    /** Test for {@link NativeDevice#setSetting(int, String, String, String)} */
+    @Test
     public void testPutSettings() throws Exception {
         String initValue = mTestDevice.getSetting(0, "system", "screen_brightness");
         CLog.i("initial value was: %s", initValue);
diff --git a/tests/src/com/android/tradefed/device/TestDeviceTest.java b/tests/src/com/android/tradefed/device/TestDeviceTest.java
index f91ad34..31d17f7 100644
--- a/tests/src/com/android/tradefed/device/TestDeviceTest.java
+++ b/tests/src/com/android/tradefed/device/TestDeviceTest.java
@@ -565,6 +565,7 @@
         assertRecoverySuccess();
         mMockIDevice.executeShellCommand(EasyMock.eq(testCommand), EasyMock.eq(mMockReceiver),
                 EasyMock.anyLong(), (TimeUnit)EasyMock.anyObject());
+        EasyMock.expect(mMockStateMonitor.waitForDeviceOnline()).andReturn(mMockIDevice);
         injectSystemProperty("ro.build.version.sdk", "23");
         replayMocks();
         mTestDevice.executeShellCommand(testCommand, mMockReceiver);
@@ -632,6 +633,7 @@
         // now expect shellCommand to be executed again, and succeed
         mMockIDevice.executeShellCommand(EasyMock.eq(testCommand), EasyMock.eq(mMockReceiver),
                 EasyMock.anyLong(), (TimeUnit)EasyMock.anyObject());
+        EasyMock.expect(mMockStateMonitor.waitForDeviceOnline()).andReturn(mMockIDevice);
         injectSystemProperty("ro.build.version.sdk", "23");
         replayMocks();
         mTestDevice.executeShellCommand(testCommand, mMockReceiver);
@@ -682,6 +684,7 @@
             assertRecoverySuccess();
         }
         EasyMock.expect(mMockIDevice.getState()).andReturn(DeviceState.ONLINE).times(2);
+        EasyMock.expect(mMockStateMonitor.waitForDeviceOnline()).andReturn(mMockIDevice).times(3);
         injectSystemProperty("ro.build.version.sdk", "23").times(3);
         replayMocks();
         try {
@@ -3017,22 +3020,32 @@
      * image size with different encoding.
      */
     public void testCompressScreenshot() throws Exception {
-        File testImageFile = getTestImageResource();
+        InputStream imageData = getClass().getResourceAsStream("/testdata/SmallRawImage.raw");
+        File testImageFile = FileUtil.createTempFile("raw-to-buffered", ".raw");
+        FileUtil.writeToFile(imageData, testImageFile);
+        RawImage testImage = null;
         try {
-            RawImage testImage = prepareRawImage(testImageFile);
+            testImage = prepareRawImage(testImageFile);
+            // We used the small image so we adapt the size.
+            testImage.height = 25;
+            testImage.size = 2000;
+            testImage.width = 25;
             // Size of the raw test data
-            Assert.assertEquals(12441600, testImage.data.length);
+            Assert.assertEquals(3000, testImage.data.length);
             byte[] result = mTestDevice.compressRawImage(testImage, "PNG", true);
             // Size after compressing
-            Assert.assertEquals(4082, result.length);
+            Assert.assertEquals(107, result.length);
 
             // Do it again with JPEG encoding
-            Assert.assertEquals(12441600, testImage.data.length);
+            Assert.assertEquals(3000, testImage.data.length);
             result = mTestDevice.compressRawImage(testImage, "JPEG", true);
             // Size after compressing as JPEG
-            Assert.assertEquals(119998, result.length);
+            Assert.assertEquals(1041, result.length);
         } finally {
-            FileUtil.recursiveDelete(testImageFile.getParentFile());
+            if (testImage != null) {
+                testImage.data = null;
+            }
+            FileUtil.deleteFile(testImageFile);
         }
     }
 
@@ -3042,10 +3055,16 @@
      * @throws Exception
      */
     public void testRawImageToBufferedImage() throws Exception {
-        File testImageFile = getTestImageResource();
-
+        InputStream imageData = getClass().getResourceAsStream("/testdata/SmallRawImage.raw");
+        File testImageFile = FileUtil.createTempFile("raw-to-buffered", ".raw");
+        FileUtil.writeToFile(imageData, testImageFile);
+        RawImage testImage = null;
         try {
-            RawImage testImage = prepareRawImage(testImageFile);
+            testImage = prepareRawImage(testImageFile);
+            // We used the small image so we adapt the size.
+            testImage.height = 25;
+            testImage.size = 2000;
+            testImage.width = 25;
 
             // Test PNG format
             BufferedImage bufferedImage = mTestDevice.rawImageToBufferedImage(testImage, "PNG");
@@ -3059,7 +3078,10 @@
             assertEquals(testImage.height, bufferedImage.getHeight());
             assertEquals(BufferedImage.TYPE_3BYTE_BGR, bufferedImage.getType());
         } finally {
-            FileUtil.recursiveDelete(testImageFile.getParentFile());
+            if (testImage != null) {
+                testImage.data = null;
+            }
+            FileUtil.deleteFile(testImageFile);
         }
     }
 
@@ -3070,15 +3092,18 @@
      */
     public void testRescaleImage() throws Exception {
         File testImageFile = getTestImageResource();
-
+        RawImage testImage = null;
         try {
-            RawImage testImage = prepareRawImage(testImageFile);
+            testImage = prepareRawImage(testImageFile);
             BufferedImage bufferedImage = mTestDevice.rawImageToBufferedImage(testImage, "PNG");
 
             BufferedImage scaledImage = mTestDevice.rescaleImage(bufferedImage);
             assertEquals(bufferedImage.getWidth() / 2, scaledImage.getWidth());
             assertEquals(bufferedImage.getHeight() / 2, scaledImage.getHeight());
         } finally {
+            if (testImage != null) {
+                testImage.data = null;
+            }
             FileUtil.recursiveDelete(testImageFile.getParentFile());
         }
     }
@@ -3226,4 +3251,46 @@
         assertEquals("bbb/bbb", stringArgs.get(1));
         assertEquals(Integer.valueOf(10), intArgs.get(1));
     }
+
+    /** Test that the output of cryptfs allows for encryption for newest format. */
+    public void testIsEncryptionSupported_newformat() throws Exception {
+        mTestDevice =
+                new TestableTestDevice() {
+                    @Override
+                    public boolean isAdbRoot() throws DeviceNotAvailableException {
+                        return true;
+                    }
+
+                    @Override
+                    public boolean enableAdbRoot() throws DeviceNotAvailableException {
+                        return true;
+                    }
+                };
+        injectShellResponse(
+                "vdc cryptfs enablecrypto",
+                "500 8674 Usage with ext4crypt: cryptfs enablecrypto inplace default noui\r\n");
+        EasyMock.replay(mMockIDevice, mMockStateMonitor, mMockDvcMonitor);
+        assertTrue(mTestDevice.isEncryptionSupported());
+        EasyMock.verify(mMockIDevice, mMockStateMonitor, mMockDvcMonitor);
+    }
+
+    /** Test that the output of cryptfs does not allow for encryption. */
+    public void testIsEncryptionSupported_failure() throws Exception {
+        mTestDevice =
+                new TestableTestDevice() {
+                    @Override
+                    public boolean isAdbRoot() throws DeviceNotAvailableException {
+                        return true;
+                    }
+
+                    @Override
+                    public boolean enableAdbRoot() throws DeviceNotAvailableException {
+                        return true;
+                    }
+                };
+        injectShellResponse("vdc cryptfs enablecrypto", "500 8674 Command not recognized\r\n");
+        EasyMock.replay(mMockIDevice, mMockStateMonitor, mMockDvcMonitor);
+        assertFalse(mTestDevice.isEncryptionSupported());
+        EasyMock.verify(mMockIDevice, mMockStateMonitor, mMockDvcMonitor);
+    }
 }
diff --git a/tests/src/com/android/tradefed/device/WifiHelperTest.java b/tests/src/com/android/tradefed/device/WifiHelperTest.java
index a36be11..2d087e5 100644
--- a/tests/src/com/android/tradefed/device/WifiHelperTest.java
+++ b/tests/src/com/android/tradefed/device/WifiHelperTest.java
@@ -168,6 +168,35 @@
         EasyMock.verify(mMockDevice);
     }
 
+    @Test
+    public void testEnsureDeviceSetup_alternateWifiUtilAPKPath() throws Exception {
+        final String apkPath = "/path/to/WifiUtil.APK";
+        EasyMock.reset(mMockDevice);
+        EasyMock.expect(mMockDevice.executeShellCommand(WifiHelper.CHECK_PACKAGE_CMD))
+                .andReturn(String.format("versionCode=%d", WifiHelper.PACKAGE_VERSION_CODE - 1));
+        EasyMock.expect(mMockDevice.installPackage(EasyMock.<File>anyObject(), EasyMock.eq(true)))
+                .andReturn(null);
+        EasyMock.replay(mMockDevice);
+        WifiHelper wifiHelper = new WifiHelper(mMockDevice, apkPath);
+        File wifiUtilApkFile = wifiHelper.getWifiUtilApkFile();
+        assertEquals(wifiUtilApkFile.getPath(), apkPath);
+        EasyMock.verify(mMockDevice);
+    }
+
+    @Test
+    public void testEnsureDeviceSetup_deleteAPK() throws Exception {
+        EasyMock.reset(mMockDevice);
+        EasyMock.expect(mMockDevice.executeShellCommand(WifiHelper.CHECK_PACKAGE_CMD))
+                .andReturn(String.format("versionCode=%d", WifiHelper.PACKAGE_VERSION_CODE - 1));
+        EasyMock.expect(mMockDevice.installPackage(EasyMock.<File>anyObject(), EasyMock.eq(true)))
+                .andReturn(null);
+        EasyMock.replay(mMockDevice);
+        WifiHelper wifiHelper = new WifiHelper(mMockDevice);
+        File wifiUtilApkFile = wifiHelper.getWifiUtilApkFile();
+        assertFalse(wifiUtilApkFile.exists());
+        EasyMock.verify(mMockDevice);
+    }
+
     /** Test that {@link WifiHelper#cleanUp()} calls uninstall on the instrumentation package. */
     @Test
     public void testCleanPackage() throws Exception {
diff --git a/tests/src/com/android/tradefed/invoker/InvocationContextTest.java b/tests/src/com/android/tradefed/invoker/InvocationContextTest.java
index f3c8974..1ab1d7e 100644
--- a/tests/src/com/android/tradefed/invoker/InvocationContextTest.java
+++ b/tests/src/com/android/tradefed/invoker/InvocationContextTest.java
@@ -17,7 +17,14 @@
 
 import static org.junit.Assert.*;
 
+import com.android.tradefed.build.BuildInfo;
+import com.android.tradefed.build.IBuildInfo;
+import com.android.tradefed.config.ConfigurationDescriptor;
 import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.util.FileUtil;
+import com.android.tradefed.util.MultiMap;
+import com.android.tradefed.util.SerializationUtil;
+import com.android.tradefed.util.UniqueMultiMap;
 
 import org.easymock.EasyMock;
 import org.junit.Before;
@@ -25,6 +32,9 @@
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
+import java.io.File;
+import java.util.Arrays;
+
 /** Unit tests for {@link InvocationContext} */
 @RunWith(JUnit4.class)
 public class InvocationContextTest {
@@ -47,4 +57,66 @@
         assertEquals("test1", mContext.getDeviceName(device1));
         assertNull(mContext.getDeviceName(device2));
     }
+
+    /**
+     * Test adding attributes and querying them. The map returned is always a copy and does not
+     * affect the actual invocation attributes.
+     */
+    @Test
+    public void testGetAttributes() {
+        mContext.addInvocationAttribute("TEST_KEY", "TEST_VALUE");
+        assertEquals(Arrays.asList("TEST_VALUE"), mContext.getAttributes().get("TEST_KEY"));
+        MultiMap<String, String> map = mContext.getAttributes();
+        map.remove("TEST_KEY");
+        // assert that the key is still there in the map from the context
+        assertEquals(Arrays.asList("TEST_VALUE"), mContext.getAttributes().get("TEST_KEY"));
+    }
+
+    /** Test that once locked the invocation context does not accept more invocation attributes. */
+    @Test
+    public void testLockedContext() {
+        mContext.lockAttributes();
+        try {
+            mContext.addInvocationAttribute("test", "Test");
+            fail("Should have thrown an exception.");
+        } catch (IllegalStateException expected) {
+            // expected
+        }
+        try {
+            mContext.addInvocationAttributes(new UniqueMultiMap<>());
+            fail("Should have thrown an exception.");
+        } catch (IllegalStateException expected) {
+            // expected
+        }
+    }
+
+    /** Test that serializing and deserializing an {@link InvocationContext}. */
+    @Test
+    public void testSerialize() throws Exception {
+        assertNotNull(mContext.getDeviceBuildMap());
+        ITestDevice device = EasyMock.createMock(ITestDevice.class);
+        IBuildInfo info = new BuildInfo("1234", "test-target");
+        mContext.addAllocatedDevice("test-device", device);
+        mContext.addDeviceBuildInfo("test-device", info);
+        mContext.setConfigurationDescriptor(new ConfigurationDescriptor());
+        assertEquals(info, mContext.getBuildInfo(device));
+        File ser = SerializationUtil.serialize(mContext);
+        try {
+            InvocationContext deserialized =
+                    (InvocationContext) SerializationUtil.deserialize(ser, true);
+            // One consequence is that transient attribute will become null but our custom
+            // deserialization should fix that.
+            assertNotNull(deserialized.getDeviceBuildMap());
+            assertNotNull(deserialized.getConfigurationDescriptor());
+            assertEquals(info, deserialized.getBuildInfo("test-device"));
+
+            // The device are not carried
+            assertTrue(deserialized.getDevices().isEmpty());
+            // Re-assigning a device, recreate the previous relationships
+            deserialized.addAllocatedDevice("test-device", device);
+            assertEquals(info, mContext.getBuildInfo(device));
+        } finally {
+            FileUtil.deleteFile(ser);
+        }
+    }
 }
diff --git a/tests/src/com/android/tradefed/invoker/TestInvocationTest.java b/tests/src/com/android/tradefed/invoker/TestInvocationTest.java
index 11c4c49..d298d86 100644
--- a/tests/src/com/android/tradefed/invoker/TestInvocationTest.java
+++ b/tests/src/com/android/tradefed/invoker/TestInvocationTest.java
@@ -18,9 +18,11 @@
 import static org.mockito.Mockito.doReturn;
 
 import com.android.ddmlib.IDevice;
+import com.android.tradefed.build.BuildInfo;
 import com.android.tradefed.build.BuildRetrievalError;
 import com.android.tradefed.build.IBuildInfo;
 import com.android.tradefed.build.IBuildProvider;
+import com.android.tradefed.build.IDeviceBuildInfo;
 import com.android.tradefed.build.IDeviceBuildProvider;
 import com.android.tradefed.command.CommandOptions;
 import com.android.tradefed.command.CommandRunner.ExitCode;
@@ -70,6 +72,7 @@
 import com.android.tradefed.testtype.IRetriableTest;
 import com.android.tradefed.testtype.IShardableTest;
 import com.android.tradefed.testtype.IStrictShardableTest;
+import com.android.tradefed.util.FileUtil;
 
 import com.google.common.util.concurrent.SettableFuture;
 
@@ -80,6 +83,7 @@
 import org.easymock.EasyMock;
 import org.mockito.Mockito;
 
+import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.ArrayList;
@@ -1326,11 +1330,16 @@
         doReturn(future).when(idevice).getBattery(Mockito.anyLong(), Mockito.any());
         EasyMock.expect(device1.getSerialNumber()).andReturn("serial1");
         context.addAllocatedDevice("device1", device1);
+        context.addDeviceBuildInfo("device1", new BuildInfo());
         EasyMock.replay(device1);
         mTestInvocation.logDeviceBatteryLevel(context, fakeEvent);
         EasyMock.verify(device1);
-        assertEquals(1, context.getAttributes().size());
-        assertEquals("50", context.getAttributes().get("serial1-battery-" + fakeEvent).get(0));
+        assertEquals(1, context.getBuildInfo("device1").getBuildAttributes().size());
+        assertEquals(
+                "50",
+                context.getBuildInfo("device1")
+                        .getBuildAttributes()
+                        .get("serial1-battery-" + fakeEvent));
     }
 
     /**
@@ -1351,13 +1360,24 @@
         EasyMock.expect(device2.getIDevice()).andReturn(idevice);
         EasyMock.expect(device2.getSerialNumber()).andReturn("serial2");
         context.addAllocatedDevice("device1", device1);
+        context.addDeviceBuildInfo("device1", new BuildInfo());
         context.addAllocatedDevice("device2", device2);
+        context.addDeviceBuildInfo("device2", new BuildInfo());
         EasyMock.replay(device1, device2);
         mTestInvocation.logDeviceBatteryLevel(context, fakeEvent);
         EasyMock.verify(device1, device2);
-        assertEquals(2, context.getAttributes().size());
-        assertEquals("50", context.getAttributes().get("serial1-battery-" + fakeEvent).get(0));
-        assertEquals("50", context.getAttributes().get("serial2-battery-" + fakeEvent).get(0));
+        assertEquals(1, context.getBuildInfo("device1").getBuildAttributes().size());
+        assertEquals(1, context.getBuildInfo("device2").getBuildAttributes().size());
+        assertEquals(
+                "50",
+                context.getBuildInfo("device1")
+                        .getBuildAttributes()
+                        .get("serial1-battery-" + fakeEvent));
+        assertEquals(
+                "50",
+                context.getBuildInfo("device2")
+                        .getBuildAttributes()
+                        .get("serial2-battery-" + fakeEvent));
     }
 
     /**
@@ -1382,15 +1402,26 @@
         ITestDevice device4 = EasyMock.createMock(ITestDevice.class);
         EasyMock.expect(device1.getIDevice()).andStubReturn(new StubDevice("stub2"));
         context.addAllocatedDevice("device1", device1);
+        context.addDeviceBuildInfo("device1", new BuildInfo());
         context.addAllocatedDevice("device2", device2);
+        context.addDeviceBuildInfo("device2", new BuildInfo());
         context.addAllocatedDevice("device3", device3);
         context.addAllocatedDevice("device4", device4);
         EasyMock.replay(device1, device2);
         mTestInvocation.logDeviceBatteryLevel(context, fakeEvent);
         EasyMock.verify(device1, device2);
-        assertEquals(2, context.getAttributes().size());
-        assertEquals("50", context.getAttributes().get("serial1-battery-" + fakeEvent).get(0));
-        assertEquals("50", context.getAttributes().get("serial2-battery-" + fakeEvent).get(0));
+        assertEquals(1, context.getBuildInfo("device1").getBuildAttributes().size());
+        assertEquals(1, context.getBuildInfo("device2").getBuildAttributes().size());
+        assertEquals(
+                "50",
+                context.getBuildInfo("device1")
+                        .getBuildAttributes()
+                        .get("serial1-battery-" + fakeEvent));
+        assertEquals(
+                "50",
+                context.getBuildInfo("device2")
+                        .getBuildAttributes()
+                        .get("serial2-battery-" + fakeEvent));
     }
 
     /** Helper to set the expectation for N number of shards. */
@@ -1527,4 +1558,94 @@
         mTestInvocation.doSetup(mStubConfiguration, context, listener);
         EasyMock.verify(device1, listener);
     }
+
+    /**
+     * Test when a {@link IDeviceBuildInfo} is passing through we do not attempt to add any external
+     * directories when there is none coming from environment.
+     */
+    public void testInvoke_deviceInfoBuild_noEnv() throws Throwable {
+        mMockBuildInfo = EasyMock.createMock(IDeviceBuildInfo.class);
+        IRemoteTest test = EasyMock.createNiceMock(IRemoteTest.class);
+        ITargetCleaner mockCleaner = EasyMock.createMock(ITargetCleaner.class);
+        mockCleaner.setUp(mMockDevice, mMockBuildInfo);
+        mockCleaner.tearDown(mMockDevice, mMockBuildInfo, null);
+        mStubConfiguration.getTargetPreparers().add(mockCleaner);
+
+        File tmpTestsDir = FileUtil.createTempDir("invocation-tests-dir");
+        try {
+            EasyMock.expect(((IDeviceBuildInfo) mMockBuildInfo).getTestsDir())
+                    .andReturn(tmpTestsDir);
+            setupMockSuccessListeners();
+            setupNormalInvoke(test);
+            EasyMock.replay(mockCleaner, mockRescheduler);
+            mTestInvocation.invoke(mStubInvocationMetadata, mStubConfiguration, mockRescheduler);
+            verifyMocks(mockCleaner, mockRescheduler);
+            verifySummaryListener();
+        } finally {
+            FileUtil.recursiveDelete(tmpTestsDir);
+        }
+    }
+
+    /**
+     * Test when a {@link IDeviceBuildInfo} is passing through we attempt to add the external
+     * directories to it when they are available.
+     */
+    public void testInvoke_deviceInfoBuild_withEnv() throws Throwable {
+        File tmpTestsDir = FileUtil.createTempDir("invocation-tests-dir");
+        File tmpExternalTestsDir = FileUtil.createTempDir("external-tf-dir");
+        File tmpTestsFile = FileUtil.createTempFile("testsfile", "txt", tmpExternalTestsDir);
+        try {
+            mTestInvocation =
+                    new TestInvocation() {
+                        @Override
+                        ILogRegistry getLogRegistry() {
+                            return mMockLogRegistry;
+                        }
+
+                        @Override
+                        protected IShardHelper createShardHelper() {
+                            return new ShardHelper();
+                        }
+
+                        @Override
+                        protected void setExitCode(ExitCode code, Throwable stack) {
+                            // empty on purpose
+                        }
+
+                        @Override
+                        List<File> getExternalTestCasesDirs() {
+                            List<File> list = new ArrayList<>();
+                            list.add(tmpExternalTestsDir);
+                            return list;
+                        }
+                    };
+            mMockBuildInfo = EasyMock.createMock(IDeviceBuildInfo.class);
+            IRemoteTest test = EasyMock.createNiceMock(IRemoteTest.class);
+            ITargetCleaner mockCleaner = EasyMock.createMock(ITargetCleaner.class);
+            mockCleaner.setUp(mMockDevice, mMockBuildInfo);
+            mockCleaner.tearDown(mMockDevice, mMockBuildInfo, null);
+            mStubConfiguration.getTargetPreparers().add(mockCleaner);
+
+            EasyMock.expect(((IDeviceBuildInfo) mMockBuildInfo).getTestsDir())
+                    .andReturn(tmpTestsDir);
+
+            setupMockSuccessListeners();
+            setupNormalInvoke(test);
+            EasyMock.replay(mockCleaner, mockRescheduler);
+            mTestInvocation.invoke(mStubInvocationMetadata, mStubConfiguration, mockRescheduler);
+            verifyMocks(mockCleaner, mockRescheduler);
+            verifySummaryListener();
+            // Check that the external directory was copied in the testsDir.
+            assertTrue(tmpTestsDir.listFiles().length == 1);
+            // external-tf-dir
+            assertEquals(tmpExternalTestsDir.getName(), tmpTestsDir.listFiles()[0].getName());
+            // testsfile.txt
+            assertTrue(tmpTestsDir.listFiles()[0].listFiles().length == 1);
+            assertEquals(
+                    tmpTestsFile.getName(), tmpTestsDir.listFiles()[0].listFiles()[0].getName());
+        } finally {
+            FileUtil.recursiveDelete(tmpTestsDir);
+            FileUtil.recursiveDelete(tmpExternalTestsDir);
+        }
+    }
 }
diff --git a/tests/src/com/android/tradefed/invoker/shard/StrictShardHelperTest.java b/tests/src/com/android/tradefed/invoker/shard/StrictShardHelperTest.java
index 29f0ac4..1d46c84 100644
--- a/tests/src/com/android/tradefed/invoker/shard/StrictShardHelperTest.java
+++ b/tests/src/com/android/tradefed/invoker/shard/StrictShardHelperTest.java
@@ -15,14 +15,21 @@
  */
 package com.android.tradefed.invoker.shard;
 
-import static org.junit.Assert.*;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
 
 import com.android.tradefed.build.BuildInfo;
 import com.android.tradefed.command.CommandOptions;
 import com.android.tradefed.config.Configuration;
+import com.android.tradefed.config.ConfigurationException;
+import com.android.tradefed.config.ConfigurationFactory;
 import com.android.tradefed.config.IConfiguration;
 import com.android.tradefed.config.OptionSetter;
 import com.android.tradefed.device.DeviceNotAvailableException;
+import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.invoker.IInvocationContext;
 import com.android.tradefed.invoker.IRescheduler;
 import com.android.tradefed.invoker.InvocationContext;
@@ -30,7 +37,9 @@
 import com.android.tradefed.result.ITestInvocationListener;
 import com.android.tradefed.testtype.IRemoteTest;
 import com.android.tradefed.testtype.StubTest;
+import com.android.tradefed.testtype.suite.ITestSuite;
 
+import org.easymock.EasyMock;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -38,6 +47,10 @@
 import org.mockito.ArgumentMatcher;
 import org.mockito.Mockito;
 
+import java.util.ArrayList;
+import java.util.LinkedHashMap;
+import java.util.List;
+
 /** Unit tests for {@link StrictShardHelper}. */
 @RunWith(JUnit4.class)
 public class StrictShardHelperTest {
@@ -173,4 +186,109 @@
         // We have no tests to put in shard-index 1 so it's empty.
         assertEquals(0, mConfig.getTests().size());
     }
+
+    /** Test class to simulate an ITestSuite getting split. */
+    public static class SplitITestSuite extends ITestSuite {
+
+        private String mName;
+
+        public SplitITestSuite() {}
+
+        public SplitITestSuite(String name) {
+            mName = name;
+        }
+
+        @Override
+        public LinkedHashMap<String, IConfiguration> loadTests() {
+            LinkedHashMap<String, IConfiguration> configs = new LinkedHashMap<>();
+            IConfiguration configuration = null;
+            try {
+                configuration =
+                        ConfigurationFactory.getInstance()
+                                .createConfigurationFromArgs(
+                                        new String[] {"empty", "--num-shards", "2"});
+            } catch (ConfigurationException e) {
+                throw new RuntimeException(e);
+            }
+            configs.put(mName, configuration);
+            return configs;
+        }
+    }
+
+    private ITestSuite createFakeSuite(String name) throws Exception {
+        ITestSuite suite = new SplitITestSuite(name);
+        return suite;
+    }
+
+    private List<IRemoteTest> testShard(int shardIndex) throws Exception {
+        mContext.addAllocatedDevice("default", EasyMock.createMock(ITestDevice.class));
+        List<IRemoteTest> test = new ArrayList<>();
+        test.add(createFakeSuite("module2"));
+        test.add(createFakeSuite("module1"));
+        test.add(createFakeSuite("module3"));
+        test.add(createFakeSuite("module1"));
+        test.add(createFakeSuite("module1"));
+        test.add(createFakeSuite("module2"));
+        test.add(createFakeSuite("module3"));
+        CommandOptions options = new CommandOptions();
+        OptionSetter setter = new OptionSetter(options);
+        setter.setOptionValue("disable-strict-sharding", "true");
+        setter.setOptionValue("shard-count", "3");
+        setter.setOptionValue("shard-index", Integer.toString(shardIndex));
+        mConfig.setCommandOptions(options);
+        mConfig.setCommandLine(new String[] {"empty"});
+        mConfig.setTests(test);
+        mHelper.shardConfig(mConfig, mContext, mRescheduler);
+        return mConfig.getTests();
+    }
+
+    /**
+     * Total for all the _shardX test should be 14 tests (2 per modules). 6 for module1: 3 module1
+     * shard * 2 4 for module2: 2 module2 shard * 2 4 for module3: 2 module3 shard * 2
+     */
+    @Test
+    public void testMergeSuite_shard0() throws Exception {
+        List<IRemoteTest> res = testShard(0);
+        assertEquals(3, res.size());
+
+        assertTrue(res.get(0) instanceof ITestSuite);
+        assertEquals("module1", ((ITestSuite) res.get(0)).getDirectModule().getId());
+        assertEquals(3, ((ITestSuite) res.get(0)).getDirectModule().numTests());
+
+        assertTrue(res.get(1) instanceof ITestSuite);
+        assertEquals("module3", ((ITestSuite) res.get(1)).getDirectModule().getId());
+        assertEquals(1, ((ITestSuite) res.get(1)).getDirectModule().numTests());
+
+        assertTrue(res.get(2) instanceof ITestSuite);
+        assertEquals("module2", ((ITestSuite) res.get(2)).getDirectModule().getId());
+        assertEquals(1, ((ITestSuite) res.get(2)).getDirectModule().numTests());
+    }
+
+    @Test
+    public void testMergeSuite_shard1() throws Exception {
+        List<IRemoteTest> res = testShard(1);
+        assertEquals(2, res.size());
+
+        assertTrue(res.get(0) instanceof ITestSuite);
+        assertEquals("module3", ((ITestSuite) res.get(0)).getDirectModule().getId());
+        assertEquals(2, ((ITestSuite) res.get(0)).getDirectModule().numTests());
+
+        assertTrue(res.get(1) instanceof ITestSuite);
+        assertEquals("module2", ((ITestSuite) res.get(1)).getDirectModule().getId());
+        assertEquals(3, ((ITestSuite) res.get(1)).getDirectModule().numTests());
+    }
+
+    @Test
+    public void testMergeSuite_shard2() throws Exception {
+        List<IRemoteTest> res = testShard(2);
+        assertEquals(2, res.size());
+
+        assertTrue(res.get(0) instanceof ITestSuite);
+        assertEquals("module1", ((ITestSuite) res.get(0)).getDirectModule().getId());
+        assertEquals(3, ((ITestSuite) res.get(0)).getDirectModule().numTests());
+
+        assertTrue(res.get(1) instanceof ITestSuite);
+        assertEquals("module3", ((ITestSuite) res.get(1)).getDirectModule().getId());
+        assertEquals(1, ((ITestSuite) res.get(1)).getDirectModule().numTests());
+    }
 }
diff --git a/tests/src/com/android/tradefed/result/SnapshotInputStreamSourceTest.java b/tests/src/com/android/tradefed/result/SnapshotInputStreamSourceTest.java
index 0c5a559..943d400 100644
--- a/tests/src/com/android/tradefed/result/SnapshotInputStreamSourceTest.java
+++ b/tests/src/com/android/tradefed/result/SnapshotInputStreamSourceTest.java
@@ -47,12 +47,13 @@
             }
         };
 
-        InputStreamSource source = new SnapshotInputStreamSource(mInputStream) {
-            @Override
-            File createBackingFile(InputStream stream) {
-                return fakeFile;
-            }
-        };
+        InputStreamSource source =
+                new SnapshotInputStreamSource("SnapUnitTest", mInputStream) {
+                    @Override
+                    File createBackingFile(String name, InputStream stream) {
+                        return fakeFile;
+                    }
+                };
 
         try {
             source.cancel();
diff --git a/tests/src/com/android/tradefed/targetprep/AppSetupFuncTest.java b/tests/src/com/android/tradefed/targetprep/AppSetupFuncTest.java
index 278ce74..d8f5c6f 100644
--- a/tests/src/com/android/tradefed/targetprep/AppSetupFuncTest.java
+++ b/tests/src/com/android/tradefed/targetprep/AppSetupFuncTest.java
@@ -16,26 +16,45 @@
 
 package com.android.tradefed.targetprep;
 
+import static org.junit.Assert.*;
+
 import com.android.tradefed.build.AppBuildInfo;
 import com.android.tradefed.config.OptionSetter;
+import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.device.WifiHelper;
-import com.android.tradefed.testtype.DeviceTestCase;
+import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
+import com.android.tradefed.testtype.IDeviceTest;
 import com.android.tradefed.util.FileUtil;
 
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
 import java.io.File;
 
 /**
  * A functional test for {@link AppSetup}.
- * <p/>
- * Relies on WifiUtil.apk in tradefed.jar.
- * <p/>
- * 'aapt' must be in PATH.
+ *
+ * <p>Relies on WifiUtil.apk in tradefed.jar.
+ *
+ * <p>'aapt' must be in PATH.
  */
-public class AppSetupFuncTest extends DeviceTestCase {
+@RunWith(DeviceJUnit4ClassRunner.class)
+public class AppSetupFuncTest implements IDeviceTest {
 
-    /**
-     * Test end to end normal case for {@link AppSetup}.
-     */
+    private ITestDevice mDevice;
+
+    @Override
+    public void setDevice(ITestDevice device) {
+        mDevice = device;
+    }
+
+    @Override
+    public ITestDevice getDevice() {
+        return mDevice;
+    }
+
+    /** Test end to end normal case for {@link AppSetup}. */
+    @Test
     public void testSetupTeardown() throws Exception {
         // use wifiutil as a test apk since it already exists
         getDevice().uninstallPackage(WifiHelper.INSTRUMENTATION_PKG);
diff --git a/tests/src/com/android/tradefed/targetprep/DeviceFlashPreparerTest.java b/tests/src/com/android/tradefed/targetprep/DeviceFlashPreparerTest.java
index 232035a..2222a7d 100644
--- a/tests/src/com/android/tradefed/targetprep/DeviceFlashPreparerTest.java
+++ b/tests/src/com/android/tradefed/targetprep/DeviceFlashPreparerTest.java
@@ -16,6 +16,8 @@
 
 package com.android.tradefed.targetprep;
 
+import static org.junit.Assert.*;
+
 import com.android.tradefed.build.DeviceBuildInfo;
 import com.android.tradefed.build.IBuildInfo;
 import com.android.tradefed.build.IDeviceBuildInfo;
@@ -31,18 +33,20 @@
 import com.android.tradefed.util.FileUtil;
 import com.android.tradefed.util.RunUtil;
 
-import junit.framework.TestCase;
-
 import org.easymock.EasyMock;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
 
 import java.io.File;
 import java.util.Arrays;
 import java.util.concurrent.Semaphore;
 
-/**
- * Unit tests for {@link DeviceFlashPreparer}.
- */
-public class DeviceFlashPreparerTest extends TestCase {
+/** Unit tests for {@link DeviceFlashPreparer}. */
+@RunWith(JUnit4.class)
+public class DeviceFlashPreparerTest {
 
     private IDeviceFlasher mMockFlasher;
     private DeviceFlashPreparer mDeviceFlashPreparer;
@@ -51,12 +55,8 @@
     private IHostOptions mMockHostOptions;
     private File mTmpDir;
 
-    /**
-     * {@inheritDoc}
-     */
-    @Override
-    protected void setUp() throws Exception {
-        super.setUp();
+    @Before
+    public void setUp() throws Exception {
         mMockFlasher = EasyMock.createMock(IDeviceFlasher.class);
         mMockDevice = EasyMock.createMock(ITestDevice.class);
         EasyMock.expect(mMockDevice.getSerialNumber()).andReturn("foo").anyTimes();
@@ -89,18 +89,13 @@
         mTmpDir = FileUtil.createTempDir("tmp");
     }
 
-    /**
-     * {@inheritDoc}
-     */
-    @Override
-    protected void tearDown() throws Exception {
+    @After
+    public void tearDown() throws Exception {
         FileUtil.recursiveDelete(mTmpDir);
-        super.tearDown();
     }
 
-    /**
-     * Simple normal case test for {@link DeviceSetup#setUp(ITestDevice, IBuildInfo)}.
-     */
+    /** Simple normal case test for {@link DeviceFlashPreparer#setUp(ITestDevice, IBuildInfo)}. */
+    @Test
     public void testSetup() throws Exception {
         doSetupExpectations();
         EasyMock.replay(mMockFlasher, mMockDevice);
@@ -132,9 +127,10 @@
     }
 
     /**
-     * Test {@link DeviceSetup#setUp(ITestDevice, IBuildInfo)} when a non IDeviceBuildInfo type
-     * is provided
+     * Test {@link DeviceFlashPreparer#setUp(ITestDevice, IBuildInfo)} when a non IDeviceBuildInfo
+     * type is provided.
      */
+    @Test
     public void testSetUp_nonDevice() throws Exception {
         try {
             mDeviceFlashPreparer.setUp(mMockDevice, EasyMock.createMock(IBuildInfo.class));
@@ -144,9 +140,8 @@
         }
     }
 
-    /**
-     * Test {@link DeviceSetup#setUp(ITestDevice, IBuildInfo)} when build does not boot
-     */
+    /** Test {@link DeviceFlashPreparer#setUp(ITestDevice, IBuildInfo)} when build does not boot. */
+    @Test
     public void testSetup_buildError() throws Exception {
         mMockDevice.setRecoveryMode(RecoveryMode.ONLINE);
         mMockFlasher.overrideDeviceOptions(mMockDevice);
@@ -180,9 +175,8 @@
         EasyMock.verify(mMockFlasher, mMockDevice);
     }
 
-    /**
-     * Ensure that the flasher instance limiting machinery is working as expected.
-     */
+    /** Ensure that the flasher instance limiting machinery is working as expected. */
+    @Test
     public void testFlashLimit() throws Exception {
         final DeviceFlashPreparer dfp = mDeviceFlashPreparer;
         try {
@@ -216,9 +210,8 @@
         }
     }
 
-    /**
-     * Ensure that the flasher limiting respects {@link IHostOptions}.
-     */
+    /** Ensure that the flasher limiting respects {@link IHostOptions}. */
+    @Test
     public void testFlashLimit_withHostOptions() throws Exception {
         final DeviceFlashPreparer dfp = mDeviceFlashPreparer;
         try {
@@ -254,9 +247,8 @@
         }
     }
 
-    /**
-     * Ensure that the flasher instance limiting machinery is working as expected.
-     */
+    /** Ensure that the flasher instance limiting machinery is working as expected. */
+    @Test
     public void testUnlimitedFlashLimit() throws Exception {
         final DeviceFlashPreparer dfp = mDeviceFlashPreparer;
         try {
diff --git a/tests/src/com/android/tradefed/targetprep/DeviceSetupFuncTest.java b/tests/src/com/android/tradefed/targetprep/DeviceSetupFuncTest.java
index fc542be..bed2668 100644
--- a/tests/src/com/android/tradefed/targetprep/DeviceSetupFuncTest.java
+++ b/tests/src/com/android/tradefed/targetprep/DeviceSetupFuncTest.java
@@ -16,38 +16,51 @@
 
 package com.android.tradefed.targetprep;
 
+import static org.junit.Assert.*;
+
 import com.android.ddmlib.Log;
 import com.android.tradefed.build.DeviceBuildInfo;
 import com.android.tradefed.build.IBuildInfo;
 import com.android.tradefed.build.IDeviceBuildInfo;
 import com.android.tradefed.device.ITestDevice;
-import com.android.tradefed.testtype.DeviceTestCase;
+import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
+import com.android.tradefed.testtype.IDeviceTest;
 
-/**
- * Functional tests for {@link DeviceSetup}.
- */
-public class DeviceSetupFuncTest extends DeviceTestCase {
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+/** Functional tests for {@link DeviceSetup}. */
+@RunWith(DeviceJUnit4ClassRunner.class)
+public class DeviceSetupFuncTest implements IDeviceTest {
 
     private static final String LOG_TAG = "DeviceSetupFuncTest";
     private DeviceSetup mDeviceSetup;
+    private ITestDevice mDevice;
     private IDeviceBuildInfo mMockBuildInfo;
 
-    /**
-     * {@inheritDoc}
-     */
     @Override
-    protected void setUp() throws Exception {
-        super.setUp();
+    public void setDevice(ITestDevice device) {
+        mDevice = device;
+    }
 
+    @Override
+    public ITestDevice getDevice() {
+        return mDevice;
+    }
+
+    @Before
+    public void setUp() throws Exception {
         mMockBuildInfo = new DeviceBuildInfo("0", "");
         mDeviceSetup = new DeviceSetup();
     }
 
     /**
      * Simple normal case test for {@link DeviceSetup#setUp(ITestDevice, IBuildInfo)}.
-     * <p/>
-     * Do setup and verify a few expected properties
+     *
+     * <p>Do setup and verify a few expected properties
      */
+    @Test
     public void testSetup() throws Exception {
         Log.i(LOG_TAG, "testSetup()");
 
diff --git a/tests/src/com/android/tradefed/targetprep/InstallAllTestZipAppsSetupTest.java b/tests/src/com/android/tradefed/targetprep/InstallAllTestZipAppsSetupTest.java
new file mode 100644
index 0000000..daf3174
--- /dev/null
+++ b/tests/src/com/android/tradefed/targetprep/InstallAllTestZipAppsSetupTest.java
@@ -0,0 +1,270 @@
+/*
+ * Copyright (C) 2017 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.targetprep;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.fail;
+
+import com.android.tradefed.build.IBuildInfo;
+import com.android.tradefed.build.IDeviceBuildInfo;
+import com.android.tradefed.device.DeviceNotAvailableException;
+import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.util.FileUtil;
+
+import org.easymock.EasyMock;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+import java.io.File;
+import java.io.IOException;
+
+/** Unit tests for {@link InstallAllTestZipAppsSetupTest} */
+@RunWith(JUnit4.class)
+public class InstallAllTestZipAppsSetupTest {
+
+    private static final String SERIAL = "SERIAL";
+    private InstallAllTestZipAppsSetup mPrep;
+    private IBuildInfo mMockBuildInfo;
+    private ITestDevice mMockTestDevice;
+    private File mMockUnzipDir;
+    private boolean mFailUnzip;
+    private boolean mFailAapt;
+
+    @Before
+    public void setUp() throws Exception {
+        mPrep =
+                new InstallAllTestZipAppsSetup() {
+                    @Override
+                    File extractZip(File testsZip) throws IOException {
+                        if (mFailUnzip) {
+                            throw new IOException();
+                        }
+                        return mMockUnzipDir;
+                    }
+
+                    @Override
+                    String getAppPackageName(File appFile) {
+                        if (mFailAapt) {
+                            return null;
+                        }
+                        return "";
+                    }
+                };
+        mFailAapt = false;
+        mFailUnzip = false;
+        mMockUnzipDir = null;
+        mMockBuildInfo = EasyMock.createMock(IDeviceBuildInfo.class);
+        mMockTestDevice = EasyMock.createMock(ITestDevice.class);
+        EasyMock.expect(mMockTestDevice.getSerialNumber()).andStubReturn(SERIAL);
+        EasyMock.expect(mMockTestDevice.getDeviceDescriptor()).andStubReturn(null);
+    }
+
+    @After
+    public void tearDown() throws Exception {
+        if (mMockUnzipDir != null) {
+            FileUtil.recursiveDelete(mMockUnzipDir);
+        }
+    }
+
+    private void setMockUnzipDir() throws IOException {
+        File testDir = FileUtil.createTempDir("TestAppSetupTest");
+        // fake hierarchy of directory and files
+        FileUtil.createTempFile("fakeApk", ".apk", testDir);
+        FileUtil.createTempFile("fakeApk2", ".apk", testDir);
+        FileUtil.createTempFile("notAnApk", ".txt", testDir);
+        File subTestDir = FileUtil.createTempDir("SubTestAppSetupTest", testDir);
+        FileUtil.createTempFile("subfakeApk", ".apk", subTestDir);
+        mMockUnzipDir = testDir;
+    }
+
+    @Test
+    public void testGetZipFile() throws DeviceNotAvailableException, TargetSetupError {
+        String zip = "zip";
+        mPrep.setTestZipName(zip);
+        File file = new File(zip);
+        EasyMock.expect(mMockBuildInfo.getFile(zip)).andReturn(file);
+        EasyMock.replay(mMockBuildInfo, mMockTestDevice);
+        File ret = mPrep.getZipFile(mMockTestDevice, mMockBuildInfo);
+        assertEquals(file, ret);
+        EasyMock.verify(mMockBuildInfo);
+    }
+
+    @Test
+    public void testGetZipFileDoesntExist() throws DeviceNotAvailableException, TargetSetupError {
+        String zip = "zip";
+        mPrep.setTestZipName(zip);
+        EasyMock.expect(mMockBuildInfo.getFile(zip)).andReturn(null);
+        EasyMock.replay(mMockBuildInfo, mMockTestDevice);
+        File ret = mPrep.getZipFile(mMockTestDevice, mMockBuildInfo);
+        assertNull(ret);
+        EasyMock.verify(mMockBuildInfo);
+    }
+
+    @Test
+    public void testNullTestZipName() throws DeviceNotAvailableException {
+        EasyMock.replay(mMockBuildInfo, mMockTestDevice);
+        try {
+            mPrep.setUp(mMockTestDevice, mMockBuildInfo);
+            fail("Should have thrown a TargetSetupError");
+        } catch (TargetSetupError e) {
+            // expected
+        }
+        EasyMock.verify(mMockBuildInfo, mMockTestDevice);
+    }
+
+    @Test
+    public void testSuccess() throws Exception {
+        mPrep.setTestZipName("zip");
+
+        mMockBuildInfo.getFile((String) EasyMock.anyObject());
+        EasyMock.expectLastCall().andReturn(new File("zip"));
+
+        setMockUnzipDir();
+
+        mMockTestDevice.installPackage((File) EasyMock.anyObject(), EasyMock.anyBoolean());
+        EasyMock.expectLastCall().andReturn(null).times(3);
+        mMockTestDevice.uninstallPackage((String) EasyMock.anyObject());
+        EasyMock.expectLastCall().andReturn(null).times(3);
+        EasyMock.replay(mMockBuildInfo, mMockTestDevice);
+
+        mPrep.setUp(mMockTestDevice, mMockBuildInfo);
+        mPrep.tearDown(mMockTestDevice, mMockBuildInfo, null);
+
+        EasyMock.verify(mMockBuildInfo, mMockTestDevice);
+    }
+
+    @Test
+    public void testSuccessNoTearDown() throws Exception {
+        mPrep.setTestZipName("zip");
+        mPrep.setCleanup(false);
+
+        mMockBuildInfo.getFile((String) EasyMock.anyObject());
+        EasyMock.expectLastCall().andReturn(new File("zip"));
+
+        setMockUnzipDir();
+
+        mMockTestDevice.installPackage((File) EasyMock.anyObject(), EasyMock.anyBoolean());
+        EasyMock.expectLastCall().andReturn(null).times(3);
+        EasyMock.replay(mMockBuildInfo, mMockTestDevice);
+
+        mPrep.setUp(mMockTestDevice, mMockBuildInfo);
+        mPrep.tearDown(mMockTestDevice, mMockBuildInfo, null);
+
+        EasyMock.verify(mMockBuildInfo, mMockTestDevice);
+    }
+
+    @Test
+    public void testInstallFailure() throws DeviceNotAvailableException {
+        final String failure = "INSTALL_PARSE_FAILED_MANIFEST_MALFORMED";
+        final String file = "TEST";
+        EasyMock.expect(
+                        mMockTestDevice.installPackage(
+                                (File) EasyMock.anyObject(), EasyMock.eq(true)))
+                .andReturn(failure);
+        EasyMock.replay(mMockBuildInfo, mMockTestDevice);
+        try {
+            mPrep.installApk(new File(file), mMockTestDevice);
+            fail("Should have thrown an exception");
+        } catch (TargetSetupError e) {
+            String expected =
+                    String.format(
+                            "Failed to install %s on %s. Reason: '%s' " + "null",
+                            file, SERIAL, failure);
+            assertEquals(expected, e.getMessage());
+        }
+        EasyMock.verify(mMockBuildInfo, mMockTestDevice);
+    }
+
+    @Test
+    public void testInstallFailureNoStop() throws DeviceNotAvailableException, TargetSetupError {
+        final String failure = "INSTALL_PARSE_FAILED_MANIFEST_MALFORMED";
+        final String file = "TEST";
+        mPrep.setStopInstallOnFailure(false);
+        EasyMock.expect(
+                        mMockTestDevice.installPackage(
+                                (File) EasyMock.anyObject(), EasyMock.eq(true)))
+                .andReturn(failure);
+        EasyMock.replay(mMockBuildInfo, mMockTestDevice);
+        // should not throw exception
+        mPrep.installApk(new File(file), mMockTestDevice);
+        EasyMock.verify(mMockBuildInfo, mMockTestDevice);
+    }
+
+    @Test
+    public void testDisable() throws Exception {
+        EasyMock.replay(mMockBuildInfo, mMockTestDevice);
+        mPrep.setDisable(true);
+        mPrep.setUp(mMockTestDevice, mMockBuildInfo);
+        mPrep.tearDown(mMockTestDevice, mMockBuildInfo, null);
+        EasyMock.verify(mMockBuildInfo, mMockTestDevice);
+    }
+
+    @Test
+    public void testUnzipFail() throws Exception {
+        mFailUnzip = true;
+        mPrep.setTestZipName("zip");
+
+        mMockBuildInfo.getFile((String) EasyMock.anyObject());
+        EasyMock.expectLastCall().andReturn(new File("zip"));
+
+        EasyMock.replay(mMockBuildInfo, mMockTestDevice);
+
+        try {
+            mPrep.setUp(mMockTestDevice, mMockBuildInfo);
+            fail("Should have thrown an exception");
+        } catch (TargetSetupError e) {
+            TargetSetupError error =
+                    new TargetSetupError(
+                            "Failed to extract test zip.",
+                            e,
+                            mMockTestDevice.getDeviceDescriptor());
+            assertEquals(error.getMessage(), e.getMessage());
+        }
+        EasyMock.verify(mMockBuildInfo, mMockTestDevice);
+    }
+
+    @Test
+    public void testAaptFail() throws Exception {
+        mFailAapt = true;
+        mPrep.setTestZipName("zip");
+        setMockUnzipDir();
+
+        mMockBuildInfo.getFile((String) EasyMock.anyObject());
+        EasyMock.expectLastCall().andReturn(new File("zip"));
+        mMockTestDevice.installPackage((File) EasyMock.anyObject(), EasyMock.anyBoolean());
+        EasyMock.expectLastCall().andReturn(null);
+
+        EasyMock.replay(mMockBuildInfo, mMockTestDevice);
+
+        try {
+            mPrep.setUp(mMockTestDevice, mMockBuildInfo);
+            fail("Should have thrown an exception");
+        } catch (TargetSetupError e) {
+            TargetSetupError error =
+                    new TargetSetupError(
+                            "apk installed but AaptParser failed",
+                            e,
+                            mMockTestDevice.getDeviceDescriptor());
+            assertEquals(error.getMessage(), e.getMessage());
+        }
+        EasyMock.verify(mMockBuildInfo, mMockTestDevice);
+    }
+}
diff --git a/tests/src/com/android/tradefed/targetprep/PushFilePreparerTest.java b/tests/src/com/android/tradefed/targetprep/PushFilePreparerTest.java
index fc1ac6a..b59213a 100644
--- a/tests/src/com/android/tradefed/targetprep/PushFilePreparerTest.java
+++ b/tests/src/com/android/tradefed/targetprep/PushFilePreparerTest.java
@@ -26,14 +26,11 @@
 import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.util.FileUtil;
 
+import org.easymock.EasyMock;
 import org.junit.Before;
 import org.junit.Test;
 
-import org.easymock.EasyMock;
-import org.mockito.Mockito;
-
 import java.io.File;
-import java.util.Arrays;
 
 /** Unit tests for {@link PushFilePreparer} */
 public class PushFilePreparerTest {
@@ -115,23 +112,6 @@
         EasyMock.verify(mMockDevice);
     }
 
-    @Test
-    public void testPushFromTestCasesDir() throws Exception {
-        mOptionSetter.setOptionValue("push", "sh->/noexist/");
-        mOptionSetter.setOptionValue("abort-on-push-failure", "false");
-
-        PushFilePreparer spyPreparer = Mockito.spy(mPreparer);
-        Mockito.doReturn(Arrays.asList(new File("/bin"))).when(spyPreparer).getTestCasesDirs();
-
-        // expect a pushFile() call as /bin/sh should exist and return false (failed)
-        EasyMock.expect(mMockDevice.pushFile((File) EasyMock.anyObject(), EasyMock.eq("/noexist/")))
-                .andReturn(Boolean.FALSE);
-        EasyMock.replay(mMockDevice);
-
-        spyPreparer.setUp(mMockDevice, null);
-        EasyMock.verify(mMockDevice);
-    }
-
     /**
      * Test {@link PushFilePreparer#resolveRelativeFilePath(IBuildInfo, String)} do not search
      * additional tests directory if the given build if is not of IBuildInfo type.
diff --git a/tests/src/com/android/tradefed/targetprep/TestAppInstallSetupTest.java b/tests/src/com/android/tradefed/targetprep/TestAppInstallSetupTest.java
index 7d2d629..0c76ffc 100644
--- a/tests/src/com/android/tradefed/targetprep/TestAppInstallSetupTest.java
+++ b/tests/src/com/android/tradefed/targetprep/TestAppInstallSetupTest.java
@@ -16,7 +16,9 @@
 
 package com.android.tradefed.targetprep;
 
-import static org.junit.Assert.*;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 
 import com.android.tradefed.build.IBuildInfo;
 import com.android.tradefed.build.IDeviceBuildInfo;
@@ -26,7 +28,6 @@
 import com.android.tradefed.util.FileUtil;
 
 import org.easymock.EasyMock;
-
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -41,18 +42,29 @@
 
     private static final String SERIAL = "SERIAL";
     private static final String PACKAGE_NAME = "PACKAGE_NAME";
+    private static final String APK_NAME = "fakeApk.apk";
     private File fakeApk;
+    private File mFakeBuildApk;
     private TestAppInstallSetup mPrep;
     private IDeviceBuildInfo mMockBuildInfo;
     private ITestDevice mMockTestDevice;
-    private File testDir;
-    private OptionSetter setter;
+    private File mTestDir;
+    private File mBuildTestDir;
+    private OptionSetter mSetter;
 
     @Before
     public void setUp() throws Exception {
-        testDir = FileUtil.createTempDir("TestAppSetupTest");
+        mTestDir = FileUtil.createTempDir("TestAppSetupTest");
+        mBuildTestDir = FileUtil.createTempDir("TestAppBuildTestDir");
         // fake hierarchy of directory and files
-        fakeApk = FileUtil.createTempFile("fakeApk", ".apk", testDir);
+        fakeApk = FileUtil.createTempFile("fakeApk", ".apk", mTestDir);
+        FileUtil.copyFile(fakeApk, new File(mTestDir, APK_NAME));
+        fakeApk = new File(mTestDir, APK_NAME);
+
+        mFakeBuildApk = FileUtil.createTempFile("fakeApk", ".apk", mBuildTestDir);
+        new File(mBuildTestDir, "DATA/app").mkdirs();
+        FileUtil.copyFile(mFakeBuildApk, new File(mBuildTestDir, "DATA/app/" + APK_NAME));
+        mFakeBuildApk = new File(mBuildTestDir, "/DATA/app/" + APK_NAME);
 
         mPrep =
                 new TestAppInstallSetup() {
@@ -69,10 +81,10 @@
                         return fakeApk;
                     }
                 };
-        mPrep.addTestFileName("fakeApk.apk");
+        mPrep.addTestFileName(APK_NAME);
 
-        setter = new OptionSetter(mPrep);
-        setter.setOptionValue("cleanup-apks", "true");
+        mSetter = new OptionSetter(mPrep);
+        mSetter.setOptionValue("cleanup-apks", "true");
         mMockBuildInfo = EasyMock.createMock(IDeviceBuildInfo.class);
         mMockTestDevice = EasyMock.createMock(ITestDevice.class);
         EasyMock.expect(mMockTestDevice.getSerialNumber()).andStubReturn(SERIAL);
@@ -81,7 +93,8 @@
 
     @After
     public void tearDown() throws Exception {
-        FileUtil.recursiveDelete(testDir);
+        FileUtil.recursiveDelete(mTestDir);
+        FileUtil.recursiveDelete(mBuildTestDir);
     }
 
     @Test
@@ -133,7 +146,10 @@
         EasyMock.verify(mMockBuildInfo, mMockTestDevice);
     }
 
-    /** Test {@link TestAppInstallSetup#setUp()} with a missing apk. TargetSetupError expected. */
+    /**
+     * Test {@link TestAppInstallSetup#setUp(ITestDevice, IBuildInfo)} with a missing apk.
+     * TargetSetupError expected.
+     */
     @Test
     public void testMissingApk() throws Exception {
         fakeApk = null; // Apk doesn't exist
@@ -147,7 +163,8 @@
     }
 
     /**
-     * Test {@link TestAppInstallSetup#setUp()} with an unreadable apk. TargetSetupError expected.
+     * Test {@link TestAppInstallSetup#setUp(ITestDevice, IBuildInfo)} with an unreadable apk.
+     * TargetSetupError expected.
      */
     @Test
     public void testUnreadableApk() throws Exception {
@@ -162,27 +179,92 @@
     }
 
     /**
-     * Test {@link TestAppInstallSetup#setUp()} with a missing apk and ThrowIfNoFile=False. Silent
-     * skip expected.
+     * Test {@link TestAppInstallSetup#setUp(ITestDevice, IBuildInfo)} with a missing apk and
+     * ThrowIfNoFile=False. Silent skip expected.
      */
     @Test
     public void testMissingApk_silent() throws Exception {
         fakeApk = null; // Apk doesn't exist
-        setter.setOptionValue("throw-if-not-found", "false");
+        mSetter.setOptionValue("throw-if-not-found", "false");
 
         mPrep.setUp(mMockTestDevice, mMockBuildInfo);
     }
 
     /**
-     * Test {@link TestAppInstallsetup#setUp()} with an unreadable apk and ThrowIfNoFile=False.
-     * Silent skip expected.
+     * Test {@link TestAppInstallSetup#setUp(ITestDevice, IBuildInfo)} with an unreadable apk and
+     * ThrowIfNoFile=False. Silent skip expected.
      */
     @Test
     public void testUnreadableApk_silent() throws Exception {
         fakeApk = new File("/not/a/real/path"); // Apk cannot be read
-        setter.setOptionValue("throw-if-not-found", "false");
+        mSetter.setOptionValue("throw-if-not-found", "false");
 
         mPrep.setUp(mMockTestDevice, mMockBuildInfo);
     }
 
+    /**
+     * Tests that when in OVERRIDE mode we install first from alt-dirs, then from BuildInfo if not
+     * found.
+     */
+    @Test
+    public void testFindApk_override() throws Exception {
+        mPrep =
+                new TestAppInstallSetup() {
+                    @Override
+                    protected String parsePackageName(
+                            File testAppFile, DeviceDescriptor deviceDescriptor) {
+                        return PACKAGE_NAME;
+                    }
+                };
+        mPrep.addTestFileName("fakeApk.apk");
+        OptionSetter setter = new OptionSetter(mPrep);
+        setter.setOptionValue("alt-dir-behavior", "OVERRIDE");
+        setter.setOptionValue("alt-dir", mTestDir.getAbsolutePath());
+        setter.setOptionValue("install-arg", "-d");
+
+        EasyMock.expect(mMockTestDevice.getDeviceDescriptor()).andStubReturn(null);
+        EasyMock.expect(mMockBuildInfo.getTestsDir()).andStubReturn(mBuildTestDir);
+
+        EasyMock.expect(
+                        mMockTestDevice.installPackage(
+                                EasyMock.eq(fakeApk), EasyMock.anyBoolean(), EasyMock.eq("-d")))
+                .andReturn(null);
+
+        EasyMock.replay(mMockTestDevice, mMockBuildInfo);
+        mPrep.setUp(mMockTestDevice, mMockBuildInfo);
+        EasyMock.verify(mMockTestDevice, mMockBuildInfo);
+    }
+
+    /**
+     * Test when OVERRIDE is set but there is not alt-dir, in this case we still use the BuildInfo.
+     */
+    @Test
+    public void testFindApk_override_onlyInBuild() throws Exception {
+        mPrep =
+                new TestAppInstallSetup() {
+                    @Override
+                    protected String parsePackageName(
+                            File testAppFile, DeviceDescriptor deviceDescriptor) {
+                        return PACKAGE_NAME;
+                    }
+                };
+        mPrep.addTestFileName("fakeApk.apk");
+        OptionSetter setter = new OptionSetter(mPrep);
+        setter.setOptionValue("alt-dir-behavior", "OVERRIDE");
+        setter.setOptionValue("install-arg", "-d");
+
+        EasyMock.expect(mMockTestDevice.getDeviceDescriptor()).andStubReturn(null);
+        EasyMock.expect(mMockBuildInfo.getTestsDir()).andStubReturn(mBuildTestDir);
+
+        EasyMock.expect(
+                        mMockTestDevice.installPackage(
+                                EasyMock.eq(mFakeBuildApk),
+                                EasyMock.anyBoolean(),
+                                EasyMock.eq("-d")))
+                .andReturn(null);
+
+        EasyMock.replay(mMockTestDevice, mMockBuildInfo);
+        mPrep.setUp(mMockTestDevice, mMockBuildInfo);
+        EasyMock.verify(mMockTestDevice, mMockBuildInfo);
+    }
 }
diff --git a/tests/src/com/android/tradefed/targetprep/TestFilePushSetupTest.java b/tests/src/com/android/tradefed/targetprep/TestFilePushSetupTest.java
index 8ddd005..d693c6b 100644
--- a/tests/src/com/android/tradefed/targetprep/TestFilePushSetupTest.java
+++ b/tests/src/com/android/tradefed/targetprep/TestFilePushSetupTest.java
@@ -201,4 +201,39 @@
         assertEquals(mAltDirFile2.getAbsolutePath(), apk.getAbsolutePath());
     }
 
+    /**
+     * Test that an exception is thrown if the file doesn't exist in extracted test dir
+     */
+    public void testThrowIfNotFound() throws Exception {
+        TestFilePushSetup setup = new TestFilePushSetup();
+        setup.setThrowIfNoFile(true);
+        // Assuming that the "file-not-in-test-zip" file doesn't exist in the test zips folder.
+        setup.addTestFileName("file-not-in-test-zip");
+        DeviceBuildInfo stubBuild = new DeviceBuildInfo("0", "stub");
+        stubBuild.setTestsDir(mFakeTestsZipFolder.getBasePath(), "0");
+        try {
+            setup.setUp(mMockDevice, stubBuild);
+            fail("Should have thrown an exception");
+        } catch (TargetSetupError expected) {
+            assertEquals(
+                    "Could not find test file file-not-in-test-zip "
+                            + "directory in extracted tests.zip null",
+                    expected.getMessage());
+        }
+    }
+
+    /**
+     * Test that no exception is thrown if the file doesn't exist in extracted test dir
+     * given that the option "throw-if-not-found" is set to false.
+     */
+    public void testThrowIfNotFound_false() throws Exception {
+        TestFilePushSetup setup = new TestFilePushSetup();
+        setup.setThrowIfNoFile(false);
+        // Assuming that the "file-not-in-test-zip" file doesn't exist in the test zips folder.
+        setup.addTestFileName("file-not-in-test-zip");
+        DeviceBuildInfo stubBuild = new DeviceBuildInfo("0", "stub");
+        stubBuild.setTestsDir(mFakeTestsZipFolder.getBasePath(), "0");
+        // test that it does not throw
+        setup.setUp(mMockDevice, stubBuild);
+    }
 }
diff --git a/tests/src/com/android/tradefed/testtype/AndroidJUnitTestTest.java b/tests/src/com/android/tradefed/testtype/AndroidJUnitTestTest.java
index 49a85a1..072557b 100644
--- a/tests/src/com/android/tradefed/testtype/AndroidJUnitTestTest.java
+++ b/tests/src/com/android/tradefed/testtype/AndroidJUnitTestTest.java
@@ -28,6 +28,7 @@
 import junit.framework.TestCase;
 
 import java.io.File;
+import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 
@@ -79,6 +80,7 @@
         mAndroidJUnitTest.setTestTimeout(TEST_TIMEOUT);
         mAndroidJUnitTest.setShellTimeout(SHELL_TIMEOUT);
         mMockRemoteRunner.setMaxTimeToOutputResponse(SHELL_TIMEOUT, TimeUnit.MILLISECONDS);
+        mMockRemoteRunner.setMaxTimeout(0L, TimeUnit.MILLISECONDS);
         mMockRemoteRunner.addInstrumentationArg(InstrumentationTest.TEST_TIMEOUT_INST_ARGS_KEY,
                 Long.toString(SHELL_TIMEOUT));
     }
@@ -194,7 +196,8 @@
         EasyMock.expect(mMockTestDevice.pushFile(
                 EasyMock.<File>anyObject(), EasyMock.<String>anyObject())).andReturn(Boolean.TRUE);
         EasyMock.expect(mMockTestDevice.executeShellCommand(EasyMock.<String>anyObject()))
-                .andReturn("");
+                .andReturn("")
+                .times(2);
         EasyMock.replay(mMockRemoteRunner, mMockTestDevice);
 
         File tmpFile = FileUtil.createTempFile("testFile", ".txt");
@@ -218,7 +221,8 @@
         EasyMock.expect(mMockTestDevice.pushFile(
                 EasyMock.<File>anyObject(), EasyMock.<String>anyObject())).andReturn(Boolean.TRUE);
         EasyMock.expect(mMockTestDevice.executeShellCommand(EasyMock.<String>anyObject()))
-                .andReturn("");
+                .andReturn("")
+                .times(2);
         EasyMock.replay(mMockRemoteRunner, mMockTestDevice);
 
         File tmpFile = FileUtil.createTempFile("notTestFile", ".txt");
@@ -247,7 +251,7 @@
                 EasyMock.<String>anyObject())).andReturn(Boolean.TRUE).times(2);
         EasyMock.expect(mMockTestDevice.executeShellCommand(EasyMock.<String>anyObject()))
                 .andReturn("")
-                .times(2);
+                .times(4);
         EasyMock.replay(mMockRemoteRunner, mMockTestDevice);
 
         File tmpFileInclude = FileUtil.createTempFile("includeFile", ".txt");
@@ -266,6 +270,40 @@
     }
 
     /**
+     * Test that when pushing the filters fails, we have a test run failure since we were not able
+     * to run anything.
+     */
+    public void testRun_testFileAndFilters_fails() throws Exception {
+        mMockRemoteRunner = EasyMock.createMock(IRemoteAndroidTestRunner.class);
+        EasyMock.expect(
+                        mMockTestDevice.pushFile(
+                                EasyMock.<File>anyObject(), EasyMock.<String>anyObject()))
+                .andThrow(new DeviceNotAvailableException("failed to push", "device1"));
+
+        mMockListener.testRunStarted(EasyMock.anyObject(), EasyMock.eq(0));
+        mMockListener.testRunFailed("failed to push");
+        mMockListener.testRunEnded(0, Collections.emptyMap());
+
+        EasyMock.replay(mMockRemoteRunner, mMockTestDevice, mMockListener);
+        File tmpFileInclude = FileUtil.createTempFile("includeFile", ".txt");
+        File tmpFileExclude = FileUtil.createTempFile("excludeFile", ".txt");
+        try {
+            mAndroidJUnitTest.addIncludeFilter(TEST1.getClassName());
+            mAndroidJUnitTest.addExcludeFilter(TEST2.toString());
+            mAndroidJUnitTest.setIncludeTestFile(tmpFileInclude);
+            mAndroidJUnitTest.setExcludeTestFile(tmpFileExclude);
+            mAndroidJUnitTest.run(mMockListener);
+            fail("Should have thrown an exception.");
+        } catch (DeviceNotAvailableException expected) {
+            //expected
+        } finally {
+            FileUtil.deleteFile(tmpFileInclude);
+            FileUtil.deleteFile(tmpFileExclude);
+        }
+        EasyMock.verify(mMockRemoteRunner, mMockTestDevice, mMockListener);
+    }
+
+    /**
      * Test that setting option for "test-file-filter" works as intended
      */
     public void testRun_setTestFileOptions() throws Exception {
@@ -281,7 +319,7 @@
                 .times(2);
         EasyMock.expect(mMockTestDevice.executeShellCommand(EasyMock.<String>anyObject()))
                 .andReturn("")
-                .times(2);
+                .times(4);
         EasyMock.replay(mMockRemoteRunner, mMockTestDevice);
 
         File tmpFileInclude = FileUtil.createTempFile("includeFile", ".txt");
@@ -330,9 +368,7 @@
         assertNull(mAndroidJUnitTest.split());
     }
 
-    /**
-     * Test that {@link AndroidJUnitTest#split()} returns 3 shards when requested to do so.
-     */
+    /** Test that {@link AndroidJUnitTest#split(int)} returns 3 shards when requested to do so. */
     public void testSplit_threeShards() throws Exception {
         mAndroidJUnitTest = new AndroidJUnitTest();
         assertEquals(AndroidJUnitTest.AJUR, mAndroidJUnitTest.getRunnerName());
@@ -349,4 +385,24 @@
         assertNull(((AndroidJUnitTest) res.get(0)).split(2));
         assertNull(((AndroidJUnitTest) res.get(0)).split());
     }
+
+    /**
+     * Test that {@link AndroidJUnitTest#split(int)} can only split up to the ajur-max-shard option.
+     */
+    public void testSplit_maxShard() throws Exception {
+        mAndroidJUnitTest = new AndroidJUnitTest();
+        assertEquals(AndroidJUnitTest.AJUR, mAndroidJUnitTest.getRunnerName());
+        OptionSetter setter = new OptionSetter(mAndroidJUnitTest);
+        setter.setOptionValue("runtime-hint", "60s");
+        setter.setOptionValue("ajur-max-shard", "2");
+        List<IRemoteTest> res = (List<IRemoteTest>) mAndroidJUnitTest.split(3);
+        assertNotNull(res);
+        assertEquals(2, res.size());
+        // Third of the execution time on each shard.
+        assertEquals(30000L, ((AndroidJUnitTest) res.get(0)).getRuntimeHint());
+        assertEquals(30000L, ((AndroidJUnitTest) res.get(1)).getRuntimeHint());
+        // Make sure shards cannot be re-sharded
+        assertNull(((AndroidJUnitTest) res.get(0)).split(2));
+        assertNull(((AndroidJUnitTest) res.get(0)).split());
+    }
 }
diff --git a/tests/src/com/android/tradefed/testtype/HostTestTest.java b/tests/src/com/android/tradefed/testtype/HostTestTest.java
index 5405949..eed77d5 100644
--- a/tests/src/com/android/tradefed/testtype/HostTestTest.java
+++ b/tests/src/com/android/tradefed/testtype/HostTestTest.java
@@ -17,11 +17,18 @@
 
 import com.android.ddmlib.testrunner.TestIdentifier;
 import com.android.tradefed.build.IBuildInfo;
+import com.android.tradefed.config.ConfigurationException;
+import com.android.tradefed.config.Option;
 import com.android.tradefed.config.OptionSetter;
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.result.ByteArrayInputStreamSource;
 import com.android.tradefed.result.ITestInvocationListener;
+import com.android.tradefed.result.InputStreamSource;
+import com.android.tradefed.result.LogDataType;
+import com.android.tradefed.testtype.DeviceJUnit4ClassRunner.TestLogData;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner.TestMetrics;
+import com.android.tradefed.util.StreamUtil;
 
 import junit.framework.Test;
 import junit.framework.TestCase;
@@ -33,7 +40,9 @@
 import org.junit.Ignore;
 import org.junit.Rule;
 import org.junit.runner.RunWith;
+import org.junit.runners.BlockJUnit4ClassRunner;
 import org.junit.runners.Suite.SuiteClasses;
+import org.junit.runners.model.InitializationError;
 
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
@@ -59,8 +68,7 @@
     @MyAnnotation
     @MyAnnotation3
     public static class SuccessTestCase extends TestCase {
-        public SuccessTestCase() {
-        }
+        public SuccessTestCase() {}
 
         public SuccessTestCase(String name) {
             super(name);
@@ -78,12 +86,43 @@
 
     public static class TestMetricTestCase extends MetricTestCase {
 
+        @Option(name = "test-option")
+        public String testOption = null;
+
+        @Option(name = "list-option")
+        public List<String> listOption = new ArrayList<>();
+
+        @Option(name = "map-option")
+        public Map<String, String> mapOption = new HashMap<>();
+
         public void testPass() {
             addTestMetric("key1", "metric1");
         }
 
         public void testPass2() {
             addTestMetric("key2", "metric2");
+            if (testOption != null) {
+                addTestMetric("test-option", testOption);
+            }
+            if (!listOption.isEmpty()) {
+                addTestMetric("list-option", listOption.toString());
+            }
+            if (!mapOption.isEmpty()) {
+                addTestMetric("map-option", mapOption.toString());
+            }
+        }
+    }
+
+    public static class LogMetricTestCase extends MetricTestCase {
+
+        public void testPass() {}
+
+        public void testPass2() {
+            addTestLog(
+                    "test2_log",
+                    LogDataType.TEXT,
+                    new ByteArrayInputStreamSource("test_log".getBytes()));
+            addTestMetric("key2", "metric2");
         }
     }
 
@@ -114,6 +153,9 @@
 
         public Junit4TestClass() {}
 
+        @Option(name = "junit4-option")
+        public boolean mOption = false;
+
         @Rule public TestMetrics metrics = new TestMetrics();
 
         @MyAnnotation
@@ -128,6 +170,36 @@
         @org.junit.Test
         public void testPass6() {
             metrics.addTestMetric("key2", "value2");
+            if (mOption) {
+                metrics.addTestMetric("junit4-option", "true");
+            }
+        }
+    }
+
+    /**
+     * Test class, we have to annotate with full org.junit.Test to avoid name collision in import.
+     */
+    @RunWith(DeviceJUnit4ClassRunner.class)
+    public static class Junit4TestLogClass {
+
+        public Junit4TestLogClass() {}
+
+        @Rule public TestLogData logs = new TestLogData();
+
+        @org.junit.Test
+        public void testPass1() {
+            ByteArrayInputStreamSource source = new ByteArrayInputStreamSource("test".getBytes());
+            logs.addTestLog("TEST", LogDataType.TEXT, source);
+            // Always cancel streams.
+            StreamUtil.cancel(source);
+        }
+
+        @org.junit.Test
+        public void testPass2() {
+            ByteArrayInputStreamSource source = new ByteArrayInputStreamSource("test2".getBytes());
+            logs.addTestLog("TEST2", LogDataType.TEXT, source);
+            // Always cancel streams.
+            StreamUtil.cancel(source);
         }
     }
 
@@ -168,6 +240,26 @@
     }
 
     /**
+     * JUnit4 runner that implements {@link ISetOptionReceiver} but does not actually have the
+     * set-option.
+     */
+    public static class InvalidJunit4Runner extends BlockJUnit4ClassRunner
+            implements ISetOptionReceiver {
+        public InvalidJunit4Runner(Class<?> klass) throws InitializationError {
+            super(klass);
+        }
+    }
+
+    @RunWith(InvalidJunit4Runner.class)
+    public static class Junit4RegularClass {
+        @Option(name = "option")
+        private String mOption = null;
+
+        @org.junit.Test
+        public void testPass() {}
+    }
+
+    /**
      * Malformed on purpose test class.
      */
     public static class Junit4MalformedTestClass {
@@ -216,12 +308,19 @@
     }
 
     public static class SuccessDeviceTest extends DeviceTestCase {
+
+        @Option(name = "option")
+        public String mOption = null;
+
         public SuccessDeviceTest() {
             super();
         }
 
         public void testPass() {
             assertNotNull(getDevice());
+            if (mOption != null) {
+                addTestMetric("option", mOption);
+            }
         }
     }
 
@@ -344,6 +443,66 @@
     }
 
     /**
+     * Test a case where a test use {@link MetricTestCase#addTestLog(String, LogDataType,
+     * InputStreamSource)} in order to log data for all the reporters to know about.
+     */
+    public void testRun_LogMetricTestCase() throws Exception {
+        mHostTest.setClassName(LogMetricTestCase.class.getName());
+        TestIdentifier test1 = new TestIdentifier(LogMetricTestCase.class.getName(), "testPass");
+        TestIdentifier test2 = new TestIdentifier(LogMetricTestCase.class.getName(), "testPass2");
+        mListener.testRunStarted((String) EasyMock.anyObject(), EasyMock.eq(2));
+        mListener.testStarted(EasyMock.eq(test1));
+        // test1 should only have its metrics
+        mListener.testEnded(test1, Collections.emptyMap());
+        // test2 should only have its metrics
+        mListener.testStarted(EasyMock.eq(test2));
+        Map<String, String> metric2 = new HashMap<>();
+        metric2.put("key2", "metric2");
+        mListener.testLog(
+                EasyMock.eq("test2_log"), EasyMock.eq(LogDataType.TEXT), EasyMock.anyObject());
+        mListener.testEnded(test2, metric2);
+        mListener.testRunEnded(EasyMock.anyLong(), (Map<String, String>) EasyMock.anyObject());
+        EasyMock.replay(mListener);
+        mHostTest.run(mListener);
+        EasyMock.verify(mListener);
+    }
+
+    /**
+     * Test success case for {@link HostTest#run(ITestInvocationListener)}, where test to run is a
+     * {@link MetricTestCase} and where an option is set to get extra metrics.
+     */
+    public void testRun_MetricTestCase_withOption() throws Exception {
+        OptionSetter setter = new OptionSetter(mHostTest);
+        setter.setOptionValue("set-option", "test-option:test");
+        // List option can take several values.
+        setter.setOptionValue("set-option", "list-option:test1");
+        setter.setOptionValue("set-option", "list-option:test2");
+        // Map option
+        setter.setOptionValue("set-option", "map-option:key=value");
+        mHostTest.setClassName(TestMetricTestCase.class.getName());
+        TestIdentifier test1 = new TestIdentifier(TestMetricTestCase.class.getName(), "testPass");
+        TestIdentifier test2 = new TestIdentifier(TestMetricTestCase.class.getName(), "testPass2");
+        mListener.testRunStarted((String) EasyMock.anyObject(), EasyMock.eq(2));
+        mListener.testStarted(EasyMock.eq(test1));
+        // test1 should only have its metrics
+        Map<String, String> metric1 = new HashMap<>();
+        metric1.put("key1", "metric1");
+        mListener.testEnded(test1, metric1);
+        // test2 should only have its metrics
+        mListener.testStarted(EasyMock.eq(test2));
+        Map<String, String> metric2 = new HashMap<>();
+        metric2.put("key2", "metric2");
+        metric2.put("test-option", "test");
+        metric2.put("list-option", "[test1, test2]");
+        metric2.put("map-option", "{key=value}");
+        mListener.testEnded(test2, metric2);
+        mListener.testRunEnded(EasyMock.anyLong(), (Map<String, String>) EasyMock.anyObject());
+        EasyMock.replay(mListener);
+        mHostTest.run(mListener);
+        EasyMock.verify(mListener);
+    }
+
+    /**
      * Test success case for {@link HostTest#run(ITestInvocationListener)}, where test to run is a
      * {@link TestSuite}.
      */
@@ -517,11 +676,15 @@
         final ITestDevice device = EasyMock.createMock(ITestDevice.class);
         mHostTest.setClassName(SuccessDeviceTest.class.getName());
         mHostTest.setDevice(device);
+        OptionSetter setter = new OptionSetter(mHostTest);
+        setter.setOptionValue("set-option", "option:value");
 
         TestIdentifier test1 = new TestIdentifier(SuccessDeviceTest.class.getName(), "testPass");
         mListener.testRunStarted((String)EasyMock.anyObject(), EasyMock.eq(1));
         mListener.testStarted(EasyMock.eq(test1));
-        mListener.testEnded(EasyMock.eq(test1), (Map<String, String>)EasyMock.anyObject());
+        Map<String, String> expected = new HashMap<>();
+        expected.put("option", "value");
+        mListener.testEnded(EasyMock.eq(test1), EasyMock.eq(expected));
         mListener.testRunEnded(EasyMock.anyLong(), (Map<String, String>)EasyMock.anyObject());
         EasyMock.replay(mListener);
         mHostTest.run(mListener);
@@ -920,11 +1083,17 @@
     public void testRun_testcase_Junit4TestNotAnnotationFiltering() throws Exception {
         mHostTest.setClassName(Junit4TestClass.class.getName());
         mHostTest.addExcludeAnnotation("com.android.tradefed.testtype.HostTestTest$MyAnnotation2");
+        OptionSetter setter = new OptionSetter(mHostTest);
+        setter.setOptionValue("set-option", "junit4-option:true");
         TestIdentifier test1 = new TestIdentifier(Junit4TestClass.class.getName(), "testPass6");
         // Only test1 will run, test2 should be filtered out.
         mListener.testRunStarted((String)EasyMock.anyObject(), EasyMock.eq(1));
         mListener.testStarted(EasyMock.eq(test1));
-        mListener.testEnded(EasyMock.eq(test1), (Map<String, String>)EasyMock.anyObject());
+        Map<String, String> metrics = new HashMap<>();
+        metrics.put("key2", "value2");
+        // If the option was correctly set, this metric should be true.
+        metrics.put("junit4-option", "true");
+        mListener.testEnded(EasyMock.eq(test1), EasyMock.eq(metrics));
         mListener.testRunEnded(EasyMock.anyLong(), (Map<String, String>)EasyMock.anyObject());
         EasyMock.replay(mListener);
         mHostTest.run(mListener);
@@ -1005,7 +1174,7 @@
     }
 
     /**
-     * Test for {@link HostTest#split()} making sure each test type is properly handled and added
+     * Test for {@link HostTest#split(int)} making sure each test type is properly handled and added
      * with a container or directly.
      */
     public void testRun_junit_suite_split() throws Exception {
@@ -1015,7 +1184,8 @@
         setter.setOptionValue("class", Junit4SuiteClass.class.getName());
         setter.setOptionValue("class", SuccessTestSuite.class.getName());
         setter.setOptionValue("class", TestRemoteNotCollector.class.getName());
-        List<IRemoteTest> list = (ArrayList<IRemoteTest>) mHostTest.split();
+        List<IRemoteTest> list = (ArrayList<IRemoteTest>) mHostTest.split(1);
+        // split by class; numShards parameter should be ignored
         assertEquals(3, list.size());
         assertEquals("com.android.tradefed.testtype.HostTest",
                 list.get(0).getClass().getName());
@@ -1051,11 +1221,57 @@
     }
 
     /**
-     * Test for {@link HostTest#split()} when no class is specified throws an exception
+     * Similar to {@link #testRun_junit_suite_split()} but with shard-unit set to method
+     */
+    public void testRun_junit_suite_split_by_method() throws Exception {
+        OptionSetter setter = new OptionSetter(mHostTest);
+        mHostTest.setDevice(mMockDevice);
+        mHostTest.setBuild(mMockBuildInfo);
+        setter.setOptionValue("class", Junit4SuiteClass.class.getName());
+        setter.setOptionValue("class", SuccessTestSuite.class.getName());
+        setter.setOptionValue("class", TestRemoteNotCollector.class.getName());
+        setter.setOptionValue("shard-unit", "method");
+        final Class<?>[] expectedTestCaseClasses = new Class<?>[] {
+            Junit4TestClass.class,
+            Junit4TestClass.class,
+            SuccessTestCase.class,
+            SuccessTestCase.class,
+            SuccessTestSuite.class,
+            SuccessTestSuite.class,
+            TestRemoteNotCollector.class,
+        };
+        List<IRemoteTest> list = (ArrayList<IRemoteTest>) mHostTest.split(expectedTestCaseClasses.length);
+        assertEquals(expectedTestCaseClasses.length, list.size());
+        for (int i = 0; i < expectedTestCaseClasses.length; i++) {
+            IRemoteTest shard = list.get(i);
+            assertTrue(HostTest.class.isInstance(shard));
+            HostTest hostTest = (HostTest)shard;
+            assertEquals(1, hostTest.getClasses().size());
+            assertEquals(1, hostTest.countTestCases());
+            assertEquals(expectedTestCaseClasses[i], hostTest.getClasses().get(0));
+        }
+
+        // We expect all the test from the JUnit4 suite to run under the original suite classname
+        // not under the container class name.
+        TestIdentifier test = new TestIdentifier(Junit4TestClass.class.getName(), "testPass5");
+        mListener.testRunStarted(test.getClassName(), 1);
+        mListener.testStarted(test);
+        mListener.testEnded(EasyMock.eq(test), (Map<String, String>)EasyMock.anyObject());
+        mListener.testRunEnded(EasyMock.anyLong(), (Map<String, String>)EasyMock.anyObject());
+        EasyMock.replay(mListener);
+        // Run the JUnit4 Container
+        ((IBuildReceiver)list.get(0)).setBuild(mMockBuildInfo);
+        ((IDeviceTest)list.get(0)).setDevice(mMockDevice);
+        list.get(0).run(mListener);
+        EasyMock.verify(mListener);
+    }
+
+    /**
+     * Test for {@link HostTest#split(int)} when no class is specified throws an exception
      */
     public void testSplit_noClass() throws Exception {
         try {
-            mHostTest.split();
+            mHostTest.split(1);
             fail("Should have thrown an exception");
         } catch (IllegalArgumentException e) {
             assertEquals("Missing Test class name", e.getMessage());
@@ -1063,7 +1279,7 @@
     }
 
     /**
-     * Test for {@link HostTest#split()} when multiple classes are specified with a method option
+     * Test for {@link HostTest#split(int)} when multiple classes are specified with a method option
      * too throws an exception
      */
     public void testSplit_methodAndMultipleClass() throws Exception {
@@ -1072,7 +1288,7 @@
         setter.setOptionValue("class", SuccessTestSuite.class.getName());
         mHostTest.setMethodName("testPass2");
         try {
-            mHostTest.split();
+            mHostTest.split(1);
             fail("Should have thrown an exception");
         } catch (IllegalArgumentException e) {
             assertEquals("Method name given with multiple test classes", e.getMessage());
@@ -1080,14 +1296,14 @@
     }
 
     /**
-     * Test for {@link HostTest#split()} when a single class is specified, no splitting can occur
+     * Test for {@link HostTest#split(int)} when a single class is specified, no splitting can occur
      * and it returns null.
      */
     public void testSplit_singleClass() throws Exception {
         OptionSetter setter = new OptionSetter(mHostTest);
         setter.setOptionValue("class", SuccessTestSuite.class.getName());
         mHostTest.setMethodName("testPass2");
-        assertNull(mHostTest.split());
+        assertNull(mHostTest.split(1));
     }
 
     /**
@@ -1126,6 +1342,39 @@
     }
 
     /**
+     * Similar to {@link #testGetTestStrictShardable()} but with shard-unit set to method
+     */
+    public void testGetTestStrictShardable_shardUnit_method() throws Exception {
+        OptionSetter setter = new OptionSetter(mHostTest);
+        setter.setOptionValue("class", Junit4SuiteClass.class.getName());
+        setter.setOptionValue("class", SuccessTestSuite.class.getName());
+        setter.setOptionValue("class", TestRemoteNotCollector.class.getName());
+        setter.setOptionValue("runtime-hint", "2m");
+        setter.setOptionValue("shard-unit", "method");
+        final int numShards = mHostTest.countTestCases();
+        final long runtimeHint = 2 * 60 * 1000; // 2 minutes in microseconds
+        final Class<?>[] expectedTestCaseClasses = new Class<?>[] {
+            Junit4TestClass.class,
+            Junit4TestClass.class,
+            SuccessTestCase.class,
+            SuccessTestCase.class,
+            SuccessTestSuite.class,
+            SuccessTestSuite.class,
+            TestRemoteNotCollector.class,
+        };
+        assertEquals(expectedTestCaseClasses.length, numShards);
+        for (int i = 0; i < numShards; i++) {
+            IRemoteTest shard = mHostTest.getTestShard(numShards, i);
+            assertTrue(shard instanceof HostTest);
+            HostTest hostTest = (HostTest)shard;
+            assertEquals(1, hostTest.getClasses().size());
+            assertEquals(1, hostTest.countTestCases());
+            assertEquals(expectedTestCaseClasses[i], hostTest.getClasses().get(0));
+            assertEquals(runtimeHint / numShards, hostTest.getRuntimeHint());
+        }
+    }
+
+    /**
      * Test for {@link HostTest#getTestShard(int, int)} when more shard than classes are requested,
      * the empty shard will have no test (StubTest).
      */
@@ -1159,6 +1408,45 @@
     }
 
     /**
+     * Similar to {@link #testGetTestStrictShardable_tooManyShards()} but with shard-unit
+     * set to method
+     */
+    public void testGetTestStrictShardable_tooManyShards_shardUnit_method() throws Exception {
+        OptionSetter setter = new OptionSetter(mHostTest);
+        setter.setOptionValue("class", Junit4SuiteClass.class.getName());
+        setter.setOptionValue("class", SuccessTestSuite.class.getName());
+        setter.setOptionValue("shard-unit", "method");
+        int numTestCases = mHostTest.countTestCases();
+        final int numShards = numTestCases + 1;
+        final Class<?>[] expectedTestCaseClasses = new Class<?>[] {
+            Junit4TestClass.class,
+            Junit4TestClass.class,
+            SuccessTestCase.class,
+            SuccessTestCase.class,
+            SuccessTestSuite.class,
+            SuccessTestSuite.class,
+        };
+        assertEquals(expectedTestCaseClasses.length, numTestCases);
+        for (int i = 0; i < numTestCases ; i++) {
+            IRemoteTest shard = mHostTest.getTestShard(numShards, i);
+            assertTrue(shard instanceof HostTest);
+            HostTest hostTest = (HostTest)shard;
+            assertEquals(1, hostTest.getClasses().size());
+            assertEquals(1, hostTest.countTestCases());
+            assertEquals(expectedTestCaseClasses[i], hostTest.getClasses().get(0));
+        }
+        IRemoteTest lastShard = mHostTest.getTestShard(numShards, numTestCases);
+        assertTrue(lastShard instanceof HostTest);
+        assertEquals(0, ((HostTest)lastShard).getClasses().size());
+        assertEquals(0, ((HostTest)lastShard).countTestCases());
+        // empty shard that can run and be skipped without reporting anything
+        ITestInvocationListener mockListener = EasyMock.createMock(ITestInvocationListener.class);
+        EasyMock.replay(mockListener);
+        lastShard.run(mockListener);
+        EasyMock.verify(mockListener);
+    }
+
+    /**
      * Test for {@link HostTest#getTestShard(int, int)} with one shard per classes.
      */
     public void testGetTestStrictShardable_wrapping() throws Exception {
@@ -1193,6 +1481,76 @@
                 ((HostTest)shard2).getClasses().get(0).getName());
     }
 
+    /**
+     * Similar to {@link #testGetTestStrictShardable_wrapping()} but with shard-unit set to method
+     */
+    public void testGetTestStrictShardable_wrapping_shardUnit_method() throws Exception {
+        testGetTestShardable_wrapping_shardUnit_method(true);
+    }
+
+    /**
+     * Similar to {@link #testGetTestStrictShardable_wrapping_shardUnit_method()}
+     * but test {@link IShardableTest} interface
+     */
+    public void testGetTestShardable_wrapping_shardUnit_method() throws Exception {
+        testGetTestShardable_wrapping_shardUnit_method(false);
+    }
+
+    /**
+     * Shard by method and verify that each shard contains the expected classes
+     *
+     * @param strict test {@link IStrictShardableTest} interface if true,
+     * {@link IShardableTest} if false
+     * @throws Exception
+     */
+    private void testGetTestShardable_wrapping_shardUnit_method(boolean strict) throws Exception {
+        final ITestDevice device = EasyMock.createMock(ITestDevice.class);
+        mHostTest.setDevice(device);
+        OptionSetter setter = new OptionSetter(mHostTest);
+        setter.setOptionValue("class", Junit4SuiteClass.class.getName());
+        setter.setOptionValue("class", SuccessTestSuite.class.getName());
+        setter.setOptionValue("class", TestRemoteNotCollector.class.getName());
+        setter.setOptionValue("class", SuccessHierarchySuite.class.getName());
+        setter.setOptionValue("class", SuccessDeviceTest.class.getName());
+        setter.setOptionValue("runtime-hint", "2m");
+        setter.setOptionValue("shard-unit", "method");
+        final Class<?>[] expectedTestCaseClasses = new Class<?>[] {
+            Junit4TestClass.class,
+            SuccessTestCase.class,
+            TestRemoteNotCollector.class,
+            SuccessDeviceTest.class,
+            Junit4TestClass.class,
+            SuccessTestSuite.class,
+            SuccessHierarchySuite.class,
+            SuccessTestCase.class,
+            SuccessTestSuite.class,
+            SuccessHierarchySuite.class,
+        };
+        final int numShards = 3;
+        final long runtimeHint = 2 * 60 * 1000; // 2 minutes in microseconds
+        int numTestCases = mHostTest.countTestCases();
+        assertEquals(expectedTestCaseClasses.length, numTestCases);
+        for (int i = 0, j = 0; i < numShards ; i++) {
+            IRemoteTest shard;
+            if (strict) {
+                shard = mHostTest.getTestShard(numShards, i);
+            } else {
+                shard = new ArrayList<>(mHostTest.split(numShards)).get(i);
+            }
+            assertTrue(shard instanceof HostTest);
+            HostTest hostTest = (HostTest)shard;
+            int q = numTestCases / numShards;
+            int r = numTestCases % numShards;
+            int n = q + (i < r ? 1 : 0);
+            assertEquals(n, hostTest.countTestCases());
+            assertEquals(n, hostTest.getClasses().size());
+            assertEquals(runtimeHint * n / numTestCases, hostTest.getRuntimeHint());
+            for (int k = 0; k < n; k++) {
+                assertEquals(expectedTestCaseClasses[j++], hostTest.getClasses().get(k));
+            }
+        }
+    }
+
     /** An annotation on the class exclude it. All the method of the class should be excluded. */
     public void testClassAnnotation_excludeAll() throws Exception {
         mHostTest.setClassName(SuccessTestCase.class.getName());
@@ -1367,7 +1725,7 @@
     }
 
     /**
-     * Test for {@link HostTest#split()} when the exclude-filter is set, it should be carried over
+     * Test for {@link HostTest#split(int)} when the exclude-filter is set, it should be carried over
      * to shards.
      */
     public void testSplit_withExclude() throws Exception {
@@ -1376,7 +1734,8 @@
         setter.setOptionValue("class", AnotherTestCase.class.getName());
         mHostTest.addExcludeFilter(
                 "com.android.tradefed.testtype.HostTestTest$SuccessTestCase#testPass");
-        Collection<IRemoteTest> res = mHostTest.split();
+        Collection<IRemoteTest> res = mHostTest.split(1);
+        // split by class; numShards parameter should be ignored
         assertEquals(2, res.size());
 
         // only one tests in the SuccessTestCase because it's been filtered out.
@@ -1413,4 +1772,126 @@
         }
         EasyMock.verify(mListener, mMockDevice);
     }
+
+    /**
+     * Test that when the 'set-option' format is not respected, an exception is thrown. Only one '='
+     * is allowed in the value.
+     */
+    public void testRun_setOption_invalid() throws Exception {
+        OptionSetter setter = new OptionSetter(mHostTest);
+        // Map option with invalid format
+        setter.setOptionValue("set-option", "map-option:key=value=2");
+        mHostTest.setClassName(TestMetricTestCase.class.getName());
+        EasyMock.replay(mListener);
+        try {
+            mHostTest.run(mListener);
+            fail("Should have thrown an exception.");
+        } catch (RuntimeException expected) {
+            // expected
+        }
+        EasyMock.verify(mListener);
+    }
+
+    /**
+     * Test that when a JUnit runner implements {@link ISetOptionReceiver} we attempt to pass it the
+     * hostTest set-option.
+     */
+    public void testSetOption_regularJUnit4_fail() throws Exception {
+        OptionSetter setter = new OptionSetter(mHostTest);
+        // Map option with invalid format
+        setter.setOptionValue("set-option", "option:value");
+        mHostTest.setClassName(Junit4RegularClass.class.getName());
+        mListener.testRunStarted(
+                EasyMock.eq("com.android.tradefed.testtype.HostTestTest$Junit4RegularClass"),
+                EasyMock.eq(1));
+        EasyMock.replay(mListener);
+        try {
+            mHostTest.run(mListener);
+            fail("Should have thrown an exception.");
+        } catch (RuntimeException expected) {
+            // expected
+        }
+        EasyMock.verify(mListener);
+    }
+
+    /**
+     * Test for {@link HostTest#run(ITestInvocationListener)}, for test with Junit4 style that log
+     * some data.
+     */
+    public void testRun_junit4style_log() throws Exception {
+        mHostTest.setClassName(Junit4TestLogClass.class.getName());
+        TestIdentifier test1 = new TestIdentifier(Junit4TestLogClass.class.getName(), "testPass1");
+        TestIdentifier test2 = new TestIdentifier(Junit4TestLogClass.class.getName(), "testPass2");
+        mListener.testRunStarted((String) EasyMock.anyObject(), EasyMock.eq(2));
+        mListener.testStarted(EasyMock.eq(test1));
+        mListener.testLog(EasyMock.eq("TEST"), EasyMock.eq(LogDataType.TEXT), EasyMock.anyObject());
+        mListener.testEnded(test1, Collections.emptyMap());
+        mListener.testStarted(EasyMock.eq(test2));
+        // test cases do not share logs, only the second test logs are seen.
+        mListener.testLog(
+                EasyMock.eq("TEST2"), EasyMock.eq(LogDataType.TEXT), EasyMock.anyObject());
+        mListener.testEnded(test2, Collections.emptyMap());
+        mListener.testRunEnded(EasyMock.anyLong(), EasyMock.anyObject());
+        EasyMock.replay(mListener);
+        mHostTest.run(mListener);
+        EasyMock.verify(mListener);
+    }
+
+    /**
+     * Similar to {@link #testSplit_withExclude()} but with shard-unit set to method
+     */
+    public void testSplit_excludeTestCase_shardUnit_method() throws Exception {
+        OptionSetter setter = new OptionSetter(mHostTest);
+        setter.setOptionValue("class", SuccessTestCase.class.getName());
+        setter.setOptionValue("class", AnotherTestCase.class.getName());
+
+        // only one tests in the SuccessTestCase because it's been filtered out.
+        TestIdentifier tid2 = new TestIdentifier(SuccessTestCase.class.getName(), "testPass2");
+        TestIdentifier tid3 = new TestIdentifier(AnotherTestCase.class.getName(), "testPass3");
+        TestIdentifier tid4 = new TestIdentifier(AnotherTestCase.class.getName(), "testPass4");
+        testSplit_excludeFilter_shardUnit_Method(
+                SuccessTestCase.class.getName() + "#testPass",
+                new TestIdentifier[] {tid2, tid3, tid4});
+    }
+
+    /**
+     * Similar to {@link #testSplit_excludeTestCase_shardUnit_method()} but exclude class
+     */
+    public void testSplit_excludeTestClass_shardUnit_method() throws Exception {
+        OptionSetter setter = new OptionSetter(mHostTest);
+        setter.setOptionValue("class", SuccessTestCase.class.getName());
+        setter.setOptionValue("class", AnotherTestCase.class.getName());
+
+        TestIdentifier tid3 = new TestIdentifier(AnotherTestCase.class.getName(), "testPass3");
+        TestIdentifier tid4 = new TestIdentifier(AnotherTestCase.class.getName(), "testPass4");
+        testSplit_excludeFilter_shardUnit_Method(
+                SuccessTestCase.class.getName(),
+                new TestIdentifier[] {tid3, tid4});
+    }
+
+    private void testSplit_excludeFilter_shardUnit_Method(
+            String excludeFilter, TestIdentifier[] expectedTids)
+            throws DeviceNotAvailableException, ConfigurationException {
+        mHostTest.addExcludeFilter(excludeFilter);
+        OptionSetter setter = new OptionSetter(mHostTest);
+        setter.setOptionValue("shard-unit", "method");
+
+        Collection<IRemoteTest> res = mHostTest.split(expectedTids.length);
+        assertEquals(expectedTids.length, res.size());
+
+        for (TestIdentifier tid : expectedTids) {
+            mListener.testRunStarted(tid.getClassName(), 1);
+            mListener.testStarted(tid);
+            mListener.testEnded(tid, Collections.emptyMap());
+            mListener.testRunEnded(EasyMock.anyLong(), EasyMock.anyObject());
+        }
+
+        EasyMock.replay(mListener, mMockDevice);
+        for (IRemoteTest test : res) {
+            assertTrue(test instanceof HostTest);
+            ((HostTest) test).setDevice(mMockDevice);
+            test.run(mListener);
+        }
+        EasyMock.verify(mListener, mMockDevice);
+    }
 }
diff --git a/tests/src/com/android/tradefed/testtype/InstrumentationTestFuncTest.java b/tests/src/com/android/tradefed/testtype/InstrumentationTestFuncTest.java
index 11894cb..253f7eb 100644
--- a/tests/src/com/android/tradefed/testtype/InstrumentationTestFuncTest.java
+++ b/tests/src/com/android/tradefed/testtype/InstrumentationTestFuncTest.java
@@ -16,6 +16,10 @@
 
 package com.android.tradefed.testtype;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
 import com.android.ddmlib.Log;
 import com.android.ddmlib.testrunner.TestIdentifier;
 import com.android.ddmlib.testrunner.TestResult.TestStatus;
@@ -23,36 +27,48 @@
 import com.android.tradefed.config.OptionSetter;
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.device.DeviceUnresponsiveException;
+import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.device.ITestDevice.RecoveryMode;
-import com.android.tradefed.device.RemoteAndroidDevice;
 import com.android.tradefed.result.CollectingTestListener;
 import com.android.tradefed.result.ITestInvocationListener;
-import com.android.tradefed.util.KeyguardControllerState;
 import com.android.tradefed.util.RunUtil;
 
 import org.easymock.EasyMock;
+import org.junit.Before;
+import org.junit.Ignore;
+import org.junit.Test;
+import org.junit.runner.RunWith;
 
 import java.io.IOException;
 
-/**
- * Functional tests for {@link InstrumentationTest}.
- */
-public class InstrumentationTestFuncTest extends DeviceTestCase {
+/** Functional tests for {@link InstrumentationTest}. */
+@RunWith(DeviceJUnit4ClassRunner.class)
+public class InstrumentationTestFuncTest implements IDeviceTest {
 
     private static final String LOG_TAG = "InstrumentationTestFuncTest";
     private static final long SHELL_TIMEOUT = 2500;
     private static final int TEST_TIMEOUT = 2000;
     private static final long WAIT_FOR_DEVICE_AVAILABLE = 5 * 60 * 1000;
 
+    private ITestDevice mDevice;
+
     /** The {@link InstrumentationTest} under test */
     private InstrumentationTest mInstrumentationTest;
 
     private ITestInvocationListener mMockListener;
 
     @Override
-    protected void setUp() throws Exception {
-        super.setUp();
+    public void setDevice(ITestDevice device) {
+        mDevice = device;
+    }
 
+    @Override
+    public ITestDevice getDevice() {
+        return mDevice;
+    }
+
+    @Before
+    public void setUp() throws Exception {
         mInstrumentationTest = new InstrumentationTest();
         mInstrumentationTest.setPackageName(TestAppConstants.TESTAPP_PACKAGE);
         mInstrumentationTest.setDevice(getDevice());
@@ -64,9 +80,8 @@
         getDevice().disableKeyguard();
     }
 
-    /**
-     * Test normal run scenario with a single passed test result.
-     */
+    /** Test normal run scenario with a single passed test result. */
+    @Test
     public void testRun() throws DeviceNotAvailableException {
         Log.i(LOG_TAG, "testRun");
         TestIdentifier expectedTest = new TestIdentifier(TestAppConstants.TESTAPP_CLASS,
@@ -84,33 +99,29 @@
         EasyMock.verify(mMockListener);
     }
 
-    /**
-     * Test normal run scenario with a single failed test result.
-     */
+    /** Test normal run scenario with a single failed test result. */
+    @Test
     public void testRun_testFailed() throws DeviceNotAvailableException {
         Log.i(LOG_TAG, "testRun_testFailed");
-        TestIdentifier expectedTest = new TestIdentifier(TestAppConstants.TESTAPP_CLASS,
-                TestAppConstants.FAILED_TEST_METHOD);
         mInstrumentationTest.setClassName(TestAppConstants.TESTAPP_CLASS);
         mInstrumentationTest.setMethodName(TestAppConstants.FAILED_TEST_METHOD);
         mInstrumentationTest.setTestTimeout(TEST_TIMEOUT);
         mInstrumentationTest.setShellTimeout(SHELL_TIMEOUT);
-        mMockListener.testRunStarted(
-                EasyMock.eq(TestAppConstants.TESTAPP_PACKAGE), EasyMock.anyInt());
-        mMockListener.testStarted(EasyMock.eq(expectedTest));
-        mMockListener.testFailed(
-                EasyMock.eq(expectedTest),
-                EasyMock.contains("junit.framework.AssertionFailedError: test failed"));
-        mMockListener.testEnded(EasyMock.eq(expectedTest), EasyMock.anyObject());
-        mMockListener.testRunEnded(EasyMock.anyLong(), EasyMock.anyObject());
-        EasyMock.replay(mMockListener);
-        mInstrumentationTest.run(mMockListener);
-        EasyMock.verify(mMockListener);
+        String[] error = new String[1];
+        error[0] = null;
+        mInstrumentationTest.run(
+                new ITestInvocationListener() {
+                    @Override
+                    public void testFailed(TestIdentifier test, String trace) {
+                        error[0] = trace;
+                    }
+                });
+        assertNotNull("testFailed was not called", error[0]);
+        assertTrue(error[0].contains("junit.framework.AssertionFailedError: test failed"));
     }
 
-    /**
-     * Test run scenario where test process crashes.
-     */
+    /** Test run scenario where test process crashes. */
+    @Test
     public void testRun_testCrash() throws DeviceNotAvailableException {
         Log.i(LOG_TAG, "testRun_testCrash");
         TestIdentifier expectedTest = new TestIdentifier(TestAppConstants.TESTAPP_CLASS,
@@ -136,71 +147,59 @@
                     EasyMock.eq("Instrumentation run failed due to 'Process crashed.'"));
         }
         mMockListener.testRunEnded(EasyMock.anyLong(), EasyMock.anyObject());
-        EasyMock.replay(mMockListener);
-        mInstrumentationTest.run(mMockListener);
-        EasyMock.verify(mMockListener);
+        try {
+            EasyMock.replay(mMockListener);
+            mInstrumentationTest.run(mMockListener);
+            EasyMock.verify(mMockListener);
+        } finally {
+            getDevice().waitForDeviceAvailable();
+        }
     }
 
-    /**
-     * Test run scenario where test run hangs indefinitely, and times out.
-     */
+    /** Test run scenario where test run hangs indefinitely, and times out. */
+    @Test
     public void testRun_testTimeout() throws DeviceNotAvailableException {
         Log.i(LOG_TAG, "testRun_testTimeout");
         RecoveryMode initMode = getDevice().getRecoveryMode();
         getDevice().setRecoveryMode(RecoveryMode.NONE);
         try {
-            TestIdentifier expectedTest =
-                    new TestIdentifier(
-                            TestAppConstants.TESTAPP_CLASS, TestAppConstants.TIMEOUT_TEST_METHOD);
             mInstrumentationTest.setClassName(TestAppConstants.TESTAPP_CLASS);
             mInstrumentationTest.setMethodName(TestAppConstants.TIMEOUT_TEST_METHOD);
             mInstrumentationTest.setShellTimeout(SHELL_TIMEOUT);
             mInstrumentationTest.setTestTimeout(TEST_TIMEOUT);
 
-            mMockListener.testRunStarted(
-                    EasyMock.eq(TestAppConstants.TESTAPP_PACKAGE), EasyMock.anyInt());
-            mMockListener.testStarted(EasyMock.eq(expectedTest));
-            mMockListener.testFailed(
-                    EasyMock.eq(expectedTest),
-                    EasyMock.contains(
-                            String.format(
-                                    "Failed to receive adb shell test output within %s ms",
-                                    SHELL_TIMEOUT)));
-            mMockListener.testEnded(EasyMock.eq(expectedTest), EasyMock.anyObject());
-            mMockListener.testRunFailed(
-                    EasyMock.contains(
-                            String.format(
-                                    "Failed to receive adb shell test output within %s ms",
-                                    SHELL_TIMEOUT)));
-            mMockListener.testRunEnded(EasyMock.anyLong(), EasyMock.anyObject());
-            EasyMock.replay(mMockListener);
-            mInstrumentationTest.run(mMockListener);
-            EasyMock.verify(mMockListener);
+            String[] error = new String[1];
+            error[0] = null;
+            mInstrumentationTest.run(
+                    new ITestInvocationListener() {
+                        @Override
+                        public void testFailed(TestIdentifier test, String trace) {
+                            error[0] = trace;
+                        }
+                    });
+            assertEquals(
+                    "Test failed to run to completion. Reason: 'Failed to receive adb shell test "
+                            + "output within 2500 ms. Test may have timed out, or adb connection to device "
+                            + "became unresponsive'. Check device logcat for details",
+                    error[0]);
         } finally {
             getDevice().setRecoveryMode(initMode);
             RunUtil.getDefault().sleep(500);
         }
     }
 
-    /**
-     * Test run scenario where device reboots during test run.
-     */
+    /** Test run scenario where device reboots during test run. */
+    @Test
     public void testRun_deviceReboot() throws Exception {
         Log.i(LOG_TAG, "testRun_deviceReboot");
-
-        TestIdentifier expectedTest = new TestIdentifier(TestAppConstants.TESTAPP_CLASS,
-                TestAppConstants.TIMEOUT_TEST_METHOD);
         mInstrumentationTest.setClassName(TestAppConstants.TESTAPP_CLASS);
         mInstrumentationTest.setMethodName(TestAppConstants.TIMEOUT_TEST_METHOD);
+        mInstrumentationTest.setShellTimeout(0);
+        mInstrumentationTest.setTestTimeout(0);
+        // Set a max timeout to avoid hanging forever for safety
+        //OptionSetter setter = new OptionSetter(mInstrumentationTest);
+        //setter.setOptionValue("max-timeout", "600000");
 
-        mMockListener.testRunStarted(TestAppConstants.TESTAPP_PACKAGE, 1);
-        mMockListener.testStarted(EasyMock.eq(expectedTest));
-        mMockListener.testFailed(EasyMock.eq(expectedTest), EasyMock.anyObject());
-        mMockListener.testEnded(EasyMock.eq(expectedTest), EasyMock.anyObject());
-        mMockListener.testRunFailed(
-                EasyMock.eq("Test run failed to complete. Expected 1 tests, received 0"));
-        mMockListener.testRunEnded(EasyMock.anyLong(), EasyMock.anyObject());
-        EasyMock.replay(mMockListener);
         // fork off a thread to do the reboot
         Thread rebootThread =
                 new Thread() {
@@ -220,58 +219,59 @@
                 };
         rebootThread.setName("InstrumentationTestFuncTest#testRun_deviceReboot");
         rebootThread.start();
-        RecoveryMode initMode = getDevice().getRecoveryMode();
-        getDevice().setRecoveryMode(RecoveryMode.NONE);
         try {
-            mInstrumentationTest.run(mMockListener);
-            // Remote device will not throw the DUE because of the different recovery path.
-            if (!(getDevice() instanceof RemoteAndroidDevice)) {
-                fail("Should have thrown an exception.");
-            }
+            String[] error = new String[1];
+            error[0] = null;
+            mInstrumentationTest.run(
+                    new ITestInvocationListener() {
+                        @Override
+                        public void testRunFailed(String errorMessage) {
+                            error[0] = errorMessage;
+                        }
+                    });
+            assertEquals("Test run failed to complete. Expected 1 tests, received 0", error[0]);
         } catch (DeviceUnresponsiveException expected) {
             // expected
         } finally {
-            getDevice().setRecoveryMode(initMode);
-        }
-        rebootThread.join(WAIT_FOR_DEVICE_AVAILABLE);
-        EasyMock.verify(mMockListener);
-        // Give some time after device available so that keyguard disabled is picked up.
-        RunUtil.getDefault().sleep(5000);
-        // now we check that the keyguard is dismissed.
-        KeyguardControllerState kcs = getDevice().getKeyguardState();
-        if (kcs != null) {
-            assertFalse("Keyguard is showing when it should not.", kcs.isKeyguardShowing());
-        } else {
-            assertTrue(runUITests());
+            rebootThread.join(WAIT_FOR_DEVICE_AVAILABLE);
+            getDevice().waitForDeviceAvailable();
         }
     }
 
-    /**
-     * Test run scenario where device runtime resets during test run.
-     * <p/>
-     * TODO: this test probably belongs more in TestDeviceFuncTest
-     */
+    /** Test that when a max-timeout is set the instrumentation is stopped. */
+    @Test
+    public void testRun_maxTimeout() throws Exception {
+        Log.i(LOG_TAG, "testRun_maxTimeout");
+        mInstrumentationTest.setClassName(TestAppConstants.TESTAPP_CLASS);
+        mInstrumentationTest.setMethodName(TestAppConstants.TIMEOUT_TEST_METHOD);
+        mInstrumentationTest.setShellTimeout(0);
+        mInstrumentationTest.setTestTimeout(0);
+        OptionSetter setter = new OptionSetter(mInstrumentationTest);
+        setter.setOptionValue("max-timeout", "5000");
+        final String[] called = new String[1];
+        called[0] = null;
+        mInstrumentationTest.run(
+                new ITestInvocationListener() {
+                    @Override
+                    public void testRunFailed(String errorMessage) {
+                        called[0] = errorMessage;
+                    }
+                });
+        assertEquals(
+                "com.android.ddmlib.TimeoutException: executeRemoteCommand timed out after 5000ms",
+                called[0]);
+    }
+
+    /** Test run scenario where device runtime resets during test run. */
+    @Test
+    @Ignore
     public void testRun_deviceRuntimeReset() throws Exception {
         Log.i(LOG_TAG, "testRun_deviceRuntimeReset");
-        TestIdentifier expectedTest = new TestIdentifier(TestAppConstants.TESTAPP_CLASS,
-                TestAppConstants.TIMEOUT_TEST_METHOD);
         mInstrumentationTest.setShellTimeout(SHELL_TIMEOUT);
         mInstrumentationTest.setTestTimeout(TEST_TIMEOUT);
         mInstrumentationTest.setClassName(TestAppConstants.TESTAPP_CLASS);
         mInstrumentationTest.setMethodName(TestAppConstants.TIMEOUT_TEST_METHOD);
 
-        mMockListener.testRunStarted(
-                EasyMock.eq(TestAppConstants.TESTAPP_PACKAGE), EasyMock.anyInt());
-        mMockListener.testStarted(EasyMock.eq(expectedTest));
-        EasyMock.expectLastCall().anyTimes();
-        mMockListener.testFailed(EasyMock.eq(expectedTest), EasyMock.anyObject());
-        EasyMock.expectLastCall().anyTimes();
-        mMockListener.testEnded(EasyMock.eq(expectedTest), EasyMock.anyObject());
-        EasyMock.expectLastCall().anyTimes();
-        mMockListener.testRunFailed(EasyMock.anyObject());
-        mMockListener.testRunEnded(EasyMock.anyLong(), EasyMock.anyObject());
-
-        EasyMock.replay(mMockListener);
         // fork off a thread to do the runtime reset
         Thread resetThread =
                 new Thread() {
@@ -300,23 +300,25 @@
                 };
         resetThread.setName("InstrumentationTestFuncTest#testRun_deviceRuntimeReset");
         resetThread.start();
-        mInstrumentationTest.run(mMockListener);
-        resetThread.join(WAIT_FOR_DEVICE_AVAILABLE);
-        EasyMock.verify(mMockListener);
-        RunUtil.getDefault().sleep(5000);
-        getDevice().waitForDeviceAvailable();
-    }
-
-    /**
-     * Run the test app UI tests and return true if they all pass.
-     */
-    private boolean runUITests() throws DeviceNotAvailableException {
-        InstrumentationTest uiTest = new InstrumentationTest();
-        uiTest.setPackageName(TestAppConstants.UITESTAPP_PACKAGE);
-        uiTest.setDevice(getDevice());
-        CollectingTestListener uilistener = new CollectingTestListener();
-        uiTest.run(uilistener);
-        return TestAppConstants.UI_TOTAL_TESTS == uilistener.getNumTestsInState(TestStatus.PASSED);
+        try {
+            String[] error = new String[1];
+            error[0] = null;
+            mInstrumentationTest.run(
+                    new ITestInvocationListener() {
+                        @Override
+                        public void testRunFailed(String errorMessage) {
+                            error[0] = errorMessage;
+                        }
+                    });
+            assertEquals(
+                    "Failed to receive adb shell test output within 120000 ms. Test may have "
+                            + "timed out, or adb connection to device became unresponsive",
+                    error[0]);
+        } finally {
+            resetThread.join(WAIT_FOR_DEVICE_AVAILABLE);
+            RunUtil.getDefault().sleep(5000);
+            getDevice().waitForDeviceAvailable();
+        }
     }
 
     /**
@@ -324,12 +326,15 @@
      * (currently TIMEOUT_TEST_METHOD and CRASH_TEST_METHOD). Verify that results are recorded for
      * all tests in the suite.
      */
+    @Test
     public void testRun_rerun() throws Exception {
         Log.i(LOG_TAG, "testRun_rerun");
         // run all tests in class
         RecoveryMode initMode = getDevice().getRecoveryMode();
         getDevice().setRecoveryMode(RecoveryMode.NONE);
         try {
+            OptionSetter setter = new OptionSetter(mInstrumentationTest);
+            setter.setOptionValue("collect-tests-timeout", Long.toString(SHELL_TIMEOUT));
             mInstrumentationTest.setClassName(TestAppConstants.TESTAPP_CLASS);
             mInstrumentationTest.setRerunMode(true);
             mInstrumentationTest.setShellTimeout(SHELL_TIMEOUT);
@@ -347,9 +352,10 @@
 
     /**
      * Test a run that crashes when collecting tests.
-     * <p/>
-     * Expect run to proceed, but be reported as a run failure
+     *
+     * <p>Expect run to proceed, but be reported as a run failure
      */
+    @Test
     public void testRun_rerunCrash() throws Exception {
         Log.i(LOG_TAG, "testRun_rerunCrash");
         mInstrumentationTest.setClassName(TestAppConstants.CRASH_ON_INIT_TEST_CLASS);
@@ -368,9 +374,10 @@
 
     /**
      * Test a run that hangs when collecting tests.
-     * <p/>
-     * Expect a run failure to be reported
+     *
+     * <p>Expect a run failure to be reported
      */
+    @Test
     public void testRun_rerunHang() throws Exception {
         Log.i(LOG_TAG, "testRun_rerunHang");
         RecoveryMode initMode = getDevice().getRecoveryMode();
diff --git a/tests/src/com/android/tradefed/testtype/InstrumentationTestTest.java b/tests/src/com/android/tradefed/testtype/InstrumentationTestTest.java
index 0c02cbb..8fa6eb2 100644
--- a/tests/src/com/android/tradefed/testtype/InstrumentationTestTest.java
+++ b/tests/src/com/android/tradefed/testtype/InstrumentationTestTest.java
@@ -135,6 +135,7 @@
         mInstrumentationTest.setTestTimeout(TEST_TIMEOUT);
         mInstrumentationTest.setShellTimeout(SHELL_TIMEOUT);
         mMockRemoteRunner.setMaxTimeToOutputResponse(SHELL_TIMEOUT, TimeUnit.MILLISECONDS);
+        mMockRemoteRunner.setMaxTimeout(0L, TimeUnit.MILLISECONDS);
         mMockRemoteRunner.addInstrumentationArg(InstrumentationTest.TEST_TIMEOUT_INST_ARGS_KEY,
                 Long.toString(SHELL_TIMEOUT));
     }
@@ -288,6 +289,7 @@
         setCollectTestsExpectations(collectTestResponse);
         // expect normal mode to be turned back on
         mMockRemoteRunner.setMaxTimeToOutputResponse(TEST_TIMEOUT, TimeUnit.MILLISECONDS);
+        mMockRemoteRunner.setMaxTimeout(0L, TimeUnit.MILLISECONDS);
         mMockRemoteRunner.setTestCollection(false);
 
         // note: expect run to not be reported
@@ -364,6 +366,7 @@
         };
         setRerunExpectations(firstRunResponse, false);
         mMockRemoteRunner.setMaxTimeToOutputResponse(SHELL_TIMEOUT, TimeUnit.MILLISECONDS);
+        mMockRemoteRunner.setMaxTimeout(0L, TimeUnit.MILLISECONDS);
         mMockRemoteRunner.addInstrumentationArg(InstrumentationTest.TEST_TIMEOUT_INST_ARGS_KEY,
                 Long.toString(SHELL_TIMEOUT));
         EasyMock.replay(mMockRemoteRunner, mMockTestDevice, mMockListener);
@@ -421,6 +424,7 @@
 
         // now expect second run with test collection mode off
         mMockRemoteRunner.setMaxTimeToOutputResponse(TEST_TIMEOUT, TimeUnit.MILLISECONDS);
+        mMockRemoteRunner.setMaxTimeout(0L, TimeUnit.MILLISECONDS);
         mMockRemoteRunner.setTestCollection(false);
         setRunTestExpectations(firstRunAnswer);
 
@@ -584,6 +588,7 @@
         mMockRemoteRunner.setDebug(false);
         // expect normal mode to be turned back on
         mMockRemoteRunner.setMaxTimeToOutputResponse(TEST_TIMEOUT, TimeUnit.MILLISECONDS);
+        mMockRemoteRunner.setMaxTimeout(0L, TimeUnit.MILLISECONDS);
         mMockRemoteRunner.setTestCollection(false);
         // We collect successfully 5 tests
         CollectTestAnswer collectTestAnswer =
diff --git a/tests/src/com/android/tradefed/testtype/NoisyDryRunTestTest.java b/tests/src/com/android/tradefed/testtype/NoisyDryRunTestTest.java
index 33b53c1..f03088a 100644
--- a/tests/src/com/android/tradefed/testtype/NoisyDryRunTestTest.java
+++ b/tests/src/com/android/tradefed/testtype/NoisyDryRunTestTest.java
@@ -15,6 +15,11 @@
  */
 package com.android.tradefed.testtype;
 
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+
+import com.android.tradefed.config.OptionSetter;
 import com.android.tradefed.result.ITestInvocationListener;
 import com.android.tradefed.util.FileUtil;
 
@@ -26,6 +31,7 @@
 import org.junit.runners.JUnit4;
 
 import java.io.File;
+import java.io.IOException;
 
 /**
  * Unit test for {@link NoisyDryRunTest}.
@@ -66,7 +72,8 @@
         replayMocks();
 
         NoisyDryRunTest noisyDryRunTest = new NoisyDryRunTest();
-        noisyDryRunTest.setCmdFile(mFile.getAbsolutePath());
+        OptionSetter setter = new OptionSetter(noisyDryRunTest);
+        setter.setOptionValue("cmdfile", mFile.getAbsolutePath());
         noisyDryRunTest.run(mMockListener);
         verifyMocks();
     }
@@ -82,7 +89,8 @@
         replayMocks();
 
         NoisyDryRunTest noisyDryRunTest = new NoisyDryRunTest();
-        noisyDryRunTest.setCmdFile(mFile.getAbsolutePath());
+        OptionSetter setter = new OptionSetter(noisyDryRunTest);
+        setter.setOptionValue("cmdfile", mFile.getAbsolutePath());
         noisyDryRunTest.run(mMockListener);
         verifyMocks();
     }
@@ -107,11 +115,85 @@
         replayMocks();
 
         NoisyDryRunTest noisyDryRunTest = new NoisyDryRunTest();
-        noisyDryRunTest.setCmdFile(mFile.getAbsolutePath());
+        OptionSetter setter = new OptionSetter(noisyDryRunTest);
+        setter.setOptionValue("cmdfile", mFile.getAbsolutePath());
         noisyDryRunTest.run(mMockListener);
         verifyMocks();
     }
 
+    @Test
+    public void testCheckFileWithTimeout() throws Exception {
+        NoisyDryRunTest noisyDryRunTest = new NoisyDryRunTest() {
+            long mCurrentTime = 0;
+            @Override
+            void sleep() {
+            }
+
+            @Override
+            long currentTimeMillis() {
+                mCurrentTime += 5 * 1000;
+                return mCurrentTime;
+            }
+        };
+        OptionSetter setter = new OptionSetter(noisyDryRunTest);
+        setter.setOptionValue("timeout", "100");
+        noisyDryRunTest.checkFileWithTimeout(mFile);
+    }
+
+    @Test
+    public void testCheckFileWithTimeout_missingFile() throws Exception {
+        NoisyDryRunTest noisyDryRunTest = new NoisyDryRunTest() {
+            long mCurrentTime = 0;
+
+            @Override
+            void sleep() {
+            }
+
+            @Override
+            long currentTimeMillis() {
+                mCurrentTime += 5 * 1000;
+                return mCurrentTime;
+            }
+        };
+        OptionSetter setter = new OptionSetter(noisyDryRunTest);
+        setter.setOptionValue("timeout", "100");
+        File missingFile = new File("missing_file");
+        try {
+            noisyDryRunTest.checkFileWithTimeout(missingFile);
+            fail("Should have thrown IOException");
+        } catch (IOException e) {
+            assertEquals(missingFile.getAbsoluteFile() + " doesn't exist.", e.getMessage());
+            assertTrue(true);
+        }
+    }
+
+    @Test
+    public void testCheckFileWithTimeout_delayFile() throws Exception {
+        FileUtil.deleteFile(mFile);
+        NoisyDryRunTest noisyDryRunTest = new NoisyDryRunTest() {
+            long mCurrentTime = 0;
+
+            @Override
+            void sleep() {
+            }
+
+            @Override
+            long currentTimeMillis() {
+                mCurrentTime += 5 * 1000;
+                if (mCurrentTime > 10 * 1000) {
+                    try {
+                        mFile.createNewFile();
+                    } catch (IOException e) {
+                    }
+                }
+                return mCurrentTime;
+            }
+        };
+        OptionSetter setter = new OptionSetter(noisyDryRunTest);
+        setter.setOptionValue("timeout", "100000");
+        noisyDryRunTest.checkFileWithTimeout(mFile);
+    }
+
     private void replayMocks() {
         EasyMock.replay(mMockListener);
     }
diff --git a/tests/src/com/android/tradefed/testtype/VersionedTfLauncherTest.java b/tests/src/com/android/tradefed/testtype/VersionedTfLauncherTest.java
index 8cf702a..ebc6c61 100644
--- a/tests/src/com/android/tradefed/testtype/VersionedTfLauncherTest.java
+++ b/tests/src/com/android/tradefed/testtype/VersionedTfLauncherTest.java
@@ -27,6 +27,7 @@
 import com.android.tradefed.config.OptionSetter;
 import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.device.NullDevice;
+import com.android.tradefed.device.StubDevice;
 import com.android.tradefed.result.FileInputStreamSource;
 import com.android.tradefed.result.ITestInvocationListener;
 import com.android.tradefed.result.LogDataType;
@@ -140,7 +141,7 @@
         EasyMock.expect(mMockBuildInfo.getBuildId()).andReturn("FAKEID").times(2);
         EasyMock.expect(mMockBuildInfo.getFile("general-tests.zip"))
                 .andReturn(new File(ADDITIONAL_TEST_ZIP));
-        EasyMock.expect(mMockTestDevice.getIDevice()).andReturn(mMockIDevice).times(1);
+        EasyMock.expect(mMockTestDevice.getIDevice()).andReturn(mMockIDevice).times(2);
         EasyMock.expect(mMockTestDevice.getSerialNumber()).andReturn(FAKE_SERIAL).times(1);
         mMockListener.testLog((String)EasyMock.anyObject(), (LogDataType)EasyMock.anyObject(),
                 (FileInputStreamSource)EasyMock.anyObject());
@@ -211,6 +212,61 @@
         EasyMock.verify(mMockTestDevice, mMockBuildInfo, mMockRunUtil, mMockListener, mMockConfig);
     }
 
+    /** Test {@link VersionedTfLauncher#run(ITestInvocationListener)} for test with a StubDevice. */
+    @Test
+    public void testRun_DeviceNoPreSetup() {
+        CommandResult cr = new CommandResult(CommandStatus.SUCCESS);
+        mMockRunUtil.unsetEnvVariable(SubprocessTfLauncher.TF_GLOBAL_CONFIG);
+        mMockRunUtil.setEnvVariablePriority(EnvPriority.SET);
+        mMockRunUtil.setEnvVariable(
+                EasyMock.eq(SubprocessTfLauncher.TF_GLOBAL_CONFIG), (String) EasyMock.anyObject());
+
+        EasyMock.expect(
+                        mMockRunUtil.runTimedCmd(
+                                EasyMock.anyLong(),
+                                (FileOutputStream) EasyMock.anyObject(),
+                                (FileOutputStream) EasyMock.anyObject(),
+                                EasyMock.eq("java"),
+                                (String) EasyMock.anyObject(),
+                                EasyMock.eq("-cp"),
+                                (String) EasyMock.anyObject(),
+                                EasyMock.eq("com.android.tradefed.command.CommandRunner"),
+                                EasyMock.eq(CONFIG_NAME),
+                                EasyMock.eq(TF_COMMAND_LINE_TEMPLATE),
+                                EasyMock.eq(TF_COMMAND_LINE_TEST),
+                                EasyMock.eq(TF_COMMAND_LINE_OPTION),
+                                EasyMock.eq(TF_COMMAND_LINE_OPTION_VALUE),
+                                EasyMock.eq("--additional-tests-zip"),
+                                EasyMock.eq(ADDITIONAL_TEST_ZIP),
+                                EasyMock.eq("--subprocess-report-file"),
+                                (String) EasyMock.anyObject()))
+                .andReturn(cr);
+        Map<ITestDevice, IBuildInfo> deviceInfos = new HashMap<ITestDevice, IBuildInfo>();
+        deviceInfos.put(mMockTestDevice, null);
+        mVersionedTfLauncher.setDeviceInfos(deviceInfos);
+        EasyMock.expect(mMockBuildInfo.getRootDir()).andReturn(new File(""));
+        EasyMock.expect(mMockBuildInfo.getBuildId()).andReturn("FAKEID").times(2);
+        EasyMock.expect(mMockBuildInfo.getFile("general-tests.zip"))
+                .andReturn(new File(ADDITIONAL_TEST_ZIP));
+        EasyMock.expect(mMockTestDevice.getIDevice()).andReturn(new StubDevice("serial1")).times(2);
+        mMockListener.testLog(
+                (String) EasyMock.anyObject(),
+                (LogDataType) EasyMock.anyObject(),
+                (FileInputStreamSource) EasyMock.anyObject());
+        EasyMock.expectLastCall().times(3);
+        mMockListener.testRunStarted("StdErr", 1);
+        mMockListener.testStarted((TestIdentifier) EasyMock.anyObject());
+        mMockListener.testEnded(
+                (TestIdentifier) EasyMock.anyObject(),
+                EasyMock.eq(Collections.<String, String>emptyMap()));
+        mMockListener.testRunEnded(0, Collections.emptyMap());
+
+        EasyMock.expect(mMockConfig.getCommandOptions()).andReturn(new CommandOptions());
+        EasyMock.replay(mMockTestDevice, mMockBuildInfo, mMockRunUtil, mMockListener, mMockConfig);
+        mVersionedTfLauncher.run(mMockListener);
+        EasyMock.verify(mMockTestDevice, mMockBuildInfo, mMockRunUtil, mMockListener, mMockConfig);
+    }
+
     /**
      * Test that when a test is sharded, the instance of the implementation is used and options are
      * passed to the shard test.
@@ -265,7 +321,7 @@
         EasyMock.expect(mMockBuildInfo.getRootDir()).andReturn(new File(""));
         EasyMock.expect(mMockBuildInfo.getBuildId()).andReturn("FAKEID").times(2);
         EasyMock.expect(mMockBuildInfo.getFile("general-tests.zip")).andReturn(null);
-        EasyMock.expect(mMockTestDevice.getIDevice()).andReturn(mMockIDevice).times(1);
+        EasyMock.expect(mMockTestDevice.getIDevice()).andReturn(mMockIDevice).times(2);
         EasyMock.expect(mMockTestDevice.getSerialNumber()).andReturn(FAKE_SERIAL).times(1);
         mMockListener.testLog(
                 (String) EasyMock.anyObject(),
diff --git a/tests/src/com/android/tradefed/testtype/suite/ITestSuiteIntegrationTest.java b/tests/src/com/android/tradefed/testtype/suite/ITestSuiteIntegrationTest.java
index d4e66e1..33c5819 100644
--- a/tests/src/com/android/tradefed/testtype/suite/ITestSuiteIntegrationTest.java
+++ b/tests/src/com/android/tradefed/testtype/suite/ITestSuiteIntegrationTest.java
@@ -41,6 +41,7 @@
 import com.android.tradefed.result.suite.SuiteResultReporter;
 import com.android.tradefed.suite.checker.ISystemStatusChecker;
 import com.android.tradefed.suite.checker.ISystemStatusCheckerReceiver;
+import com.android.tradefed.testtype.IInvocationContextReceiver;
 import com.android.tradefed.testtype.IRemoteTest;
 import com.android.tradefed.util.FileUtil;
 
@@ -195,6 +196,7 @@
         suite.setDevice(mMockDevice);
         suite.setBuild(mMockBuildInfo);
         suite.setSystemStatusChecker(new ArrayList<ISystemStatusChecker>());
+        suite.setInvocationContext(mContext);
         mListener.invocationStarted(mContext);
         suite.run(mListener);
         mListener.invocationEnded(System.currentTimeMillis());
@@ -215,6 +217,7 @@
         suite.setDevice(mMockDevice);
         suite.setBuild(mMockBuildInfo);
         suite.setSystemStatusChecker(new ArrayList<ISystemStatusChecker>());
+        suite.setInvocationContext(mContext);
         mListener.invocationStarted(mContext);
         suite.run(mListener);
         mListener.invocationEnded(System.currentTimeMillis());
@@ -238,6 +241,7 @@
         suite.setDevice(mMockDevice);
         suite.setBuild(mMockBuildInfo);
         suite.setSystemStatusChecker(new ArrayList<ISystemStatusChecker>());
+        suite.setInvocationContext(mContext);
         mListener.invocationStarted(mContext);
         suite.run(mListener);
         mListener.invocationEnded(System.currentTimeMillis());
@@ -261,6 +265,7 @@
         suite.setDevice(mMockDevice);
         suite.setBuild(mMockBuildInfo);
         suite.setSystemStatusChecker(new ArrayList<ISystemStatusChecker>());
+        suite.setInvocationContext(mContext);
         mListener.invocationStarted(mContext);
         try {
             suite.run(mListener);
@@ -289,6 +294,9 @@
                     ((ISystemStatusCheckerReceiver) test)
                             .setSystemStatusChecker(config.getSystemStatusCheckers());
                 }
+                if (test instanceof IInvocationContextReceiver) {
+                    ((IInvocationContextReceiver) test).setInvocationContext(mContext);
+                }
                 try {
                     test.run(new ResultForwarder(config.getTestInvocationListeners()));
                 } catch (DeviceNotAvailableException e) {
@@ -325,6 +333,10 @@
                                                     .setSystemStatusChecker(
                                                             config.getSystemStatusCheckers());
                                         }
+                                        if (test instanceof IInvocationContextReceiver) {
+                                            ((IInvocationContextReceiver) test)
+                                                    .setInvocationContext(mContext);
+                                        }
                                         try {
                                             test.run(
                                                     new ResultForwarder(
@@ -362,6 +374,7 @@
         config.setTest(suite);
         config.setSystemStatusCheckers(new ArrayList<ISystemStatusChecker>());
         suite.setSystemStatusChecker(new ArrayList<ISystemStatusChecker>());
+        suite.setInvocationContext(mContext);
         config.setTestInvocationListener(mListener);
         config.getCommandOptions().setShardCount(5);
         IDeviceConfiguration deviceConfig = new DeviceConfigurationHolder("device");
@@ -398,6 +411,7 @@
         config.setTest(suite);
         config.setSystemStatusCheckers(new ArrayList<ISystemStatusChecker>());
         suite.setSystemStatusChecker(new ArrayList<ISystemStatusChecker>());
+        suite.setInvocationContext(mContext);
         config.setTestInvocationListener(mListener);
         config.getCommandOptions().setShardCount(5);
         IDeviceConfiguration deviceConfig = new DeviceConfigurationHolder("device");
@@ -438,6 +452,7 @@
         config.setTest(suite);
         config.setSystemStatusCheckers(new ArrayList<ISystemStatusChecker>());
         suite.setSystemStatusChecker(new ArrayList<ISystemStatusChecker>());
+        suite.setInvocationContext(mContext);
         config.setTestInvocationListener(mListener);
         config.getCommandOptions().setShardCount(2);
         config.getCommandOptions().setShardIndex(0);
@@ -460,6 +475,9 @@
                 ((ISystemStatusCheckerReceiver) test)
                         .setSystemStatusChecker(config.getSystemStatusCheckers());
             }
+            if (test instanceof IInvocationContextReceiver) {
+                ((IInvocationContextReceiver) test).setInvocationContext(mContext);
+            }
             test.run(new ResultForwarder(config.getTestInvocationListeners()));
         }
         new ResultForwarder(config.getTestInvocationListeners()).invocationEnded(500);
@@ -513,6 +531,7 @@
         config.setTest(suite);
         config.setSystemStatusCheckers(new ArrayList<ISystemStatusChecker>());
         suite.setSystemStatusChecker(new ArrayList<ISystemStatusChecker>());
+        suite.setInvocationContext(mContext);
         config.setTestInvocationListener(mListener);
         config.getCommandOptions().setShardCount(shardCount);
         config.getCommandOptions().setShardIndex(shardIndex);
@@ -535,6 +554,9 @@
                 ((ISystemStatusCheckerReceiver) test)
                         .setSystemStatusChecker(config.getSystemStatusCheckers());
             }
+            if (test instanceof IInvocationContextReceiver) {
+                ((IInvocationContextReceiver) test).setInvocationContext(mContext);
+            }
             test.run(new ResultForwarder(config.getTestInvocationListeners()));
         }
         new ResultForwarder(config.getTestInvocationListeners()).invocationEnded(500);
diff --git a/tests/src/com/android/tradefed/testtype/suite/ITestSuiteTest.java b/tests/src/com/android/tradefed/testtype/suite/ITestSuiteTest.java
index 9460d21..f16714a 100644
--- a/tests/src/com/android/tradefed/testtype/suite/ITestSuiteTest.java
+++ b/tests/src/com/android/tradefed/testtype/suite/ITestSuiteTest.java
@@ -27,6 +27,8 @@
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.device.DeviceUnresponsiveException;
 import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.invoker.IInvocationContext;
+import com.android.tradefed.invoker.InvocationContext;
 import com.android.tradefed.log.LogUtil.CLog;
 import com.android.tradefed.result.ByteArrayInputStreamSource;
 import com.android.tradefed.result.ITestInvocationListener;
@@ -62,6 +64,7 @@
     private ITestDevice mMockDevice;
     private IBuildInfo mMockBuildInfo;
     private ISystemStatusChecker mMockSysChecker;
+    private IInvocationContext mContext;
 
     /**
      * Very basic implementation of {@link ITestSuite} to test it.
@@ -143,6 +146,8 @@
         mMockSysChecker = EasyMock.createMock(ISystemStatusChecker.class);
         mTestSuite.setDevice(mMockDevice);
         mTestSuite.setBuild(mMockBuildInfo);
+        mContext = new InvocationContext();
+        mTestSuite.setInvocationContext(mContext);
     }
 
     /**
@@ -291,6 +296,7 @@
                 };
         mTestSuite.setDevice(mMockDevice);
         mTestSuite.setBuild(mMockBuildInfo);
+        mTestSuite.setInvocationContext(mContext);
         OptionSetter setter = new OptionSetter(mTestSuite);
         setter.setOptionValue("skip-all-system-status-check", "true");
         setter.setOptionValue("reboot-per-module", "true");
@@ -332,6 +338,7 @@
                 };
         mTestSuite.setDevice(mMockDevice);
         mTestSuite.setBuild(mMockBuildInfo);
+        mTestSuite.setInvocationContext(mContext);
         OptionSetter setter = new OptionSetter(mTestSuite);
         setter.setOptionValue("skip-all-system-status-check", "true");
         setter.setOptionValue("reboot-per-module", "true");
diff --git a/tests/src/com/android/tradefed/testtype/suite/ModuleDefinitionTest.java b/tests/src/com/android/tradefed/testtype/suite/ModuleDefinitionTest.java
index 49b1a9e..4a8cf31 100644
--- a/tests/src/com/android/tradefed/testtype/suite/ModuleDefinitionTest.java
+++ b/tests/src/com/android/tradefed/testtype/suite/ModuleDefinitionTest.java
@@ -15,18 +15,22 @@
  */
 package com.android.tradefed.testtype.suite;
 
-import static org.junit.Assert.*;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
 
 import com.android.ddmlib.testrunner.TestIdentifier;
 import com.android.tradefed.build.IBuildInfo;
 import com.android.tradefed.command.remote.DeviceDescriptor;
+import com.android.tradefed.config.ConfigurationDescriptor;
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.invoker.IInvocationContext;
 import com.android.tradefed.result.ITestInvocationListener;
 import com.android.tradefed.targetprep.BuildError;
 import com.android.tradefed.targetprep.ITargetCleaner;
 import com.android.tradefed.targetprep.ITargetPreparer;
 import com.android.tradefed.targetprep.TargetSetupError;
+import com.android.tradefed.testtype.Abi;
 import com.android.tradefed.testtype.IBuildReceiver;
 import com.android.tradefed.testtype.IDeviceTest;
 import com.android.tradefed.testtype.IRemoteTest;
@@ -115,7 +119,9 @@
         mTargetPrepList.add(mMockCleaner);
         mMockBuildInfo = EasyMock.createMock(IBuildInfo.class);
         mMockDevice = EasyMock.createMock(ITestDevice.class);
-        mModule = new ModuleDefinition(MODULE_NAME, mTestList, mTargetPrepList);
+        mModule =
+                new ModuleDefinition(
+                        MODULE_NAME, mTestList, mTargetPrepList, new ConfigurationDescriptor());
     }
 
     /**
@@ -189,7 +195,9 @@
                 throw new TargetSetupError(exceptionMessage, nullDescriptor);
             }
         });
-        mModule = new ModuleDefinition(MODULE_NAME, mTestList, mTargetPrepList);
+        mModule =
+                new ModuleDefinition(
+                        MODULE_NAME, mTestList, mTargetPrepList, new ConfigurationDescriptor());
         mMockCleaner.tearDown(EasyMock.eq(mMockDevice), EasyMock.eq(mMockBuildInfo),
                 EasyMock.isNull());
         mMockListener.testRunStarted(EasyMock.eq(MODULE_NAME), EasyMock.eq(1));
@@ -212,7 +220,9 @@
         final int testCount = 5;
         List<IRemoteTest> testList = new ArrayList<>();
         testList.add(new TestObject("run1", testCount, false));
-        mModule = new ModuleDefinition(MODULE_NAME, testList, mTargetPrepList);
+        mModule =
+                new ModuleDefinition(
+                        MODULE_NAME, testList, mTargetPrepList, new ConfigurationDescriptor());
         mModule.setBuild(mMockBuildInfo);
         mModule.setDevice(mMockDevice);
         mMockPrep.setUp(EasyMock.eq(mMockDevice), EasyMock.eq(mMockBuildInfo));
@@ -242,7 +252,9 @@
         final int testCount = 4;
         List<IRemoteTest> testList = new ArrayList<>();
         testList.add(new TestObject("run1", testCount, true));
-        mModule = new ModuleDefinition(MODULE_NAME, testList, mTargetPrepList);
+        mModule =
+                new ModuleDefinition(
+                        MODULE_NAME, testList, mTargetPrepList, new ConfigurationDescriptor());
         mModule.setBuild(mMockBuildInfo);
         mModule.setDevice(mMockDevice);
         mMockPrep.setUp(EasyMock.eq(mMockDevice), EasyMock.eq(mMockBuildInfo));
@@ -273,4 +285,24 @@
         assertEquals(2, mModule.getTestsResults().get(0).getNumCompleteTests());
         verifyMocks();
     }
+
+    /**
+     * Test that when a module is created with some particular informations, the resulting {@link
+     * IInvocationContext} of the module is properly populated.
+     */
+    @Test
+    public void testAbiSetting() {
+        final int testCount = 5;
+        ConfigurationDescriptor descriptor = new ConfigurationDescriptor();
+        descriptor.setAbi(new Abi("arm", "32"));
+        List<IRemoteTest> testList = new ArrayList<>();
+        testList.add(new TestObject("run1", testCount, false));
+        mModule = new ModuleDefinition(MODULE_NAME, testList, mTargetPrepList, descriptor);
+        // Check that the invocation module created has expected informations
+        IInvocationContext moduleContext = mModule.getModuleInvocationContext();
+        assertEquals(
+                MODULE_NAME,
+                moduleContext.getAttributes().get(ModuleDefinition.MODULE_NAME).get(0));
+        assertEquals("arm", moduleContext.getAttributes().get(ModuleDefinition.MODULE_ABI).get(0));
+    }
 }
diff --git a/tests/src/com/android/tradefed/testtype/suite/ModuleMergerTest.java b/tests/src/com/android/tradefed/testtype/suite/ModuleMergerTest.java
new file mode 100644
index 0000000..931646c
--- /dev/null
+++ b/tests/src/com/android/tradefed/testtype/suite/ModuleMergerTest.java
@@ -0,0 +1,153 @@
+/*
+ * Copyright (C) 2017 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.suite;
+
+import static org.junit.Assert.*;
+
+import com.android.tradefed.invoker.shard.StrictShardHelperTest.SplitITestSuite;
+import com.android.tradefed.testtype.IRemoteTest;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+import java.util.Collection;
+import java.util.Iterator;
+
+/** Unit tests for {@link ModuleMerger}. */
+@RunWith(JUnit4.class)
+public class ModuleMergerTest {
+
+    /**
+     * Test that {@link ModuleMerger#arePartOfSameSuite(ITestSuite, ITestSuite)} returns false when
+     * the first suite is not splitted yet.
+     */
+    @Test
+    public void testPartOfSameSuite_notSplittedYet() {
+        SplitITestSuite suite1 = new SplitITestSuite("module1");
+        SplitITestSuite suite2 = new SplitITestSuite("module2");
+        assertFalse(ModuleMerger.arePartOfSameSuite(suite1, suite2));
+    }
+
+    /**
+     * Test that {@link ModuleMerger#arePartOfSameSuite(ITestSuite, ITestSuite)} returns false when
+     * the second suite is not splitted yet.
+     */
+    @Test
+    public void testPartOfSameSuite_notSplittedYet2() {
+        SplitITestSuite suite1 = new SplitITestSuite("module1");
+        Collection<IRemoteTest> res1 = suite1.split(2);
+        SplitITestSuite suite2 = new SplitITestSuite("module2");
+        assertFalse(ModuleMerger.arePartOfSameSuite((ITestSuite) res1.iterator().next(), suite2));
+    }
+
+    /**
+     * Test that {@link ModuleMerger#arePartOfSameSuite(ITestSuite, ITestSuite)} returns true when
+     * the two suites are splitted and from the same module.
+     */
+    @Test
+    public void testPartOfSameSuite_sameSuite() {
+        SplitITestSuite suite1 = new SplitITestSuite("module1");
+        Collection<IRemoteTest> res1 = suite1.split(2);
+        Iterator<IRemoteTest> ite = res1.iterator();
+        assertTrue(
+                ModuleMerger.arePartOfSameSuite((ITestSuite) ite.next(), (ITestSuite) ite.next()));
+    }
+
+    /**
+     * Test that {@link ModuleMerger#arePartOfSameSuite(ITestSuite, ITestSuite)} returns false when
+     * the two suites are splitted but from different modules.
+     */
+    @Test
+    public void testPartOfSameSuite_notSameSuite() {
+        SplitITestSuite suite1 = new SplitITestSuite("module1");
+        Collection<IRemoteTest> res1 = suite1.split(2);
+        SplitITestSuite suite2 = new SplitITestSuite("module2");
+        Collection<IRemoteTest> res2 = suite2.split(2);
+        assertFalse(
+                ModuleMerger.arePartOfSameSuite(
+                        (ITestSuite) res1.iterator().next(), (ITestSuite) res2.iterator().next()));
+    }
+
+    /**
+     * Test that {@link ModuleMerger#mergeSplittedITestSuite(ITestSuite, ITestSuite)} throws an
+     * exception when the first suite is not splitted yet.
+     */
+    @Test
+    public void testMergeSplittedITestSuite_notSplittedYet() {
+        SplitITestSuite suite1 = new SplitITestSuite("module1");
+        SplitITestSuite suite2 = new SplitITestSuite("module2");
+        try {
+            ModuleMerger.mergeSplittedITestSuite(suite1, suite2);
+            fail("Should have thrown an exception.");
+        } catch (IllegalArgumentException expected) {
+            // expected
+        }
+    }
+
+    /**
+     * Test that {@link ModuleMerger#mergeSplittedITestSuite(ITestSuite, ITestSuite)} throws an
+     * exception when the second suite is not splitted yet.
+     */
+    @Test
+    public void testMergeSplittedITestSuite_notSplittedYet2() {
+        SplitITestSuite suite1 = new SplitITestSuite("module1");
+        Collection<IRemoteTest> res1 = suite1.split(2);
+        SplitITestSuite suite2 = new SplitITestSuite("module2");
+        try {
+            ModuleMerger.mergeSplittedITestSuite((ITestSuite) res1.iterator().next(), suite2);
+            fail("Should have thrown an exception.");
+        } catch (IllegalArgumentException expected) {
+            // expected
+        }
+    }
+
+    /**
+     * Test that {@link ModuleMerger#mergeSplittedITestSuite(ITestSuite, ITestSuite)} throws an
+     * exception when the two suites are splitted but coming from different modules.
+     */
+    @Test
+    public void testMergeSplittedITestSuite_splittedSuiteFromDifferentModules() {
+        SplitITestSuite suite1 = new SplitITestSuite("module1");
+        Collection<IRemoteTest> res1 = suite1.split(2);
+        SplitITestSuite suite2 = new SplitITestSuite("module2");
+        Collection<IRemoteTest> res2 = suite2.split(2);
+        try {
+            ModuleMerger.mergeSplittedITestSuite(
+                    (ITestSuite) res1.iterator().next(), (ITestSuite) res2.iterator().next());
+            fail("Should have thrown an exception.");
+        } catch (IllegalArgumentException expected) {
+            // expected
+        }
+    }
+
+    /**
+     * Test that {@link ModuleMerger#mergeSplittedITestSuite(ITestSuite, ITestSuite)} properly
+     * assigns tests from the second suite to the first since part of same module.
+     */
+    @Test
+    public void testMergeSplittedITestSuite() {
+        SplitITestSuite suite1 = new SplitITestSuite("module1");
+        Collection<IRemoteTest> res1 = suite1.split(2);
+        Iterator<IRemoteTest> ite = res1.iterator();
+        ITestSuite split1 = (ITestSuite) ite.next();
+        ITestSuite split2 = (ITestSuite) ite.next();
+        assertEquals(2, split1.getDirectModule().numTests());
+        assertEquals(2, split2.getDirectModule().numTests());
+        ModuleMerger.mergeSplittedITestSuite(split1, split2);
+        assertEquals(4, split1.getDirectModule().numTests());
+    }
+}
diff --git a/tests/src/com/android/tradefed/testtype/suite/ModuleSplitterTest.java b/tests/src/com/android/tradefed/testtype/suite/ModuleSplitterTest.java
index 6b72498..63257b1 100644
--- a/tests/src/com/android/tradefed/testtype/suite/ModuleSplitterTest.java
+++ b/tests/src/com/android/tradefed/testtype/suite/ModuleSplitterTest.java
@@ -50,6 +50,11 @@
 
         IConfiguration config = new Configuration("fakeConfig", "desc");
         config.setTargetPreparer(new StubTargetPreparer());
+        StubTest test = new StubTest();
+        OptionSetter setterTest = new OptionSetter(test);
+        // allow StubTest to shard in 6 sub tests
+        setterTest.setOptionValue("num-shards", "6");
+        config.setTest(test);
 
         OptionSetter setter = new OptionSetter(config.getConfigurationDescription());
         setter.setOptionValue("not-shardable", "true");
@@ -65,6 +70,64 @@
 
     /**
      * Tests that {@link ModuleSplitter#splitConfiguration(LinkedHashMap, int, boolean)} on a non
+     * strict shardable configuration and a dynamic context results in a matching ModuleDefinitions
+     * to be created for each shards since they will be sharded.
+     */
+    @Test
+    public void testSplitModule_configNotStrictShardable_dynamic() throws Exception {
+        LinkedHashMap<String, IConfiguration> runConfig = new LinkedHashMap<>();
+
+        IConfiguration config = new Configuration("fakeConfig", "desc");
+        config.setTargetPreparer(new StubTargetPreparer());
+        StubTest test = new StubTest();
+        OptionSetter setterTest = new OptionSetter(test);
+        // allow StubTest to shard in 6 sub tests
+        setterTest.setOptionValue("num-shards", "6");
+        config.setTest(test);
+
+        OptionSetter setter = new OptionSetter(config.getConfigurationDescription());
+        setter.setOptionValue("not-strict-shardable", "true");
+        runConfig.put("module1", config);
+        List<ModuleDefinition> res = ModuleSplitter.splitConfiguration(runConfig, 5, true);
+        // We are sharding since even if we are not-strict-shardable, we are in dynamic context
+        assertEquals(10, res.size());
+        // The original target preparer is changed since we split multiple <test> tags.
+        assertNotSame(config.getTargetPreparers().get(0), res.get(0).getTargetPreparers().get(0));
+        // The original IRemoteTest is changed since it was sharded
+        assertNotSame(config.getTests().get(0), res.get(0).getTests().get(0));
+    }
+
+    /**
+     * Tests that {@link ModuleSplitter#splitConfiguration(LinkedHashMap, int, boolean)} on a non
+     * strict shardable configuration and not dymaic results in a matching ModuleDefinition to be
+     * created with the same objects, since we will not be sharding.
+     */
+    @Test
+    public void testSplitModule_configNotStrictShardable_notDynamic() throws Exception {
+        LinkedHashMap<String, IConfiguration> runConfig = new LinkedHashMap<>();
+
+        IConfiguration config = new Configuration("fakeConfig", "desc");
+        config.setTargetPreparer(new StubTargetPreparer());
+        StubTest test = new StubTest();
+        OptionSetter setterTest = new OptionSetter(test);
+        // allow StubTest to shard in 6 sub tests
+        setterTest.setOptionValue("num-shards", "6");
+        config.setTest(test);
+
+        OptionSetter setter = new OptionSetter(config.getConfigurationDescription());
+        setter.setOptionValue("not-strict-shardable", "true");
+        runConfig.put("module1", config);
+        List<ModuleDefinition> res = ModuleSplitter.splitConfiguration(runConfig, 5, false);
+        // matching 1 for 1, config to ModuleDefinition since not shardable
+        assertEquals(1, res.size());
+        // The original target preparer is changed since we split multiple <test> tags.
+        assertNotSame(config.getTargetPreparers().get(0), res.get(0).getTargetPreparers().get(0));
+        // The original IRemoteTest is still there because we use a pool.
+        assertSame(config.getTests().get(0), res.get(0).getTests().get(0));
+    }
+
+    /**
+     * Tests that {@link ModuleSplitter#splitConfiguration(LinkedHashMap, int, boolean)} on a non
      * shardable test results in a matching ModuleDefinition created with the original IRemoteTest
      * and copied ITargetPreparers.
      */
diff --git a/tests/src/com/android/tradefed/testtype/suite/TfSuiteRunnerTest.java b/tests/src/com/android/tradefed/testtype/suite/TfSuiteRunnerTest.java
index ec37e86..245b916 100644
--- a/tests/src/com/android/tradefed/testtype/suite/TfSuiteRunnerTest.java
+++ b/tests/src/com/android/tradefed/testtype/suite/TfSuiteRunnerTest.java
@@ -27,6 +27,7 @@
 import com.android.tradefed.config.IConfiguration;
 import com.android.tradefed.config.OptionSetter;
 import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.invoker.InvocationContext;
 import com.android.tradefed.result.ITestInvocationListener;
 import com.android.tradefed.testtype.IRemoteTest;
 import com.android.tradefed.testtype.StubTest;
@@ -178,6 +179,7 @@
         mRunner.setDevice(mock(ITestDevice.class));
         mRunner.setBuild(mock(IBuildInfo.class));
         mRunner.setSystemStatusChecker(new ArrayList<>());
+        mRunner.setInvocationContext(new InvocationContext());
         // runs the expanded suite
         listener.testRunStarted("suite/stub1", 0);
         listener.testRunEnded(EasyMock.anyLong(), EasyMock.anyObject());
diff --git a/tests/src/com/android/tradefed/util/FileUtilFuncTest.java b/tests/src/com/android/tradefed/util/FileUtilFuncTest.java
index 2f69bca..683e347 100644
--- a/tests/src/com/android/tradefed/util/FileUtilFuncTest.java
+++ b/tests/src/com/android/tradefed/util/FileUtilFuncTest.java
@@ -342,6 +342,24 @@
         }
     }
 
+    /** Test that {@link FileUtil#recursiveSimlink(File, File)} properly simlink files. */
+    public void testRecursiveSimlink() throws IOException {
+        File dir1 = null;
+        File dest = null;
+        try {
+            dir1 = FileUtil.createTempDir("orig-dir");
+            File subdir1 = FileUtil.createTempDir("sub-dir", dir1);
+            File testFile = FileUtil.createTempFile("test", "file", subdir1);
+            dest = FileUtil.createTempDir("dest-dir");
+            FileUtil.recursiveSimlink(dir1, dest);
+            // check that file is in dest dir
+            assertNotNull(FileUtil.findFile(dest, testFile.getName()));
+        } finally {
+            FileUtil.recursiveDelete(dir1);
+            FileUtil.recursiveDelete(dest);
+        }
+    }
+
     // Assertions
     private String assertUnixPerms(File file, String expPerms) {
         String perms = ls(file.getPath());
@@ -352,7 +370,8 @@
 
     // Helpers
     private String ls(String path) {
-        CommandResult result = RunUtil.getDefault().runTimedCmd(10 * 1000, "ls", "-ld", path);
+        CommandResult result =
+                RunUtil.getDefault().runTimedCmdRetry(10 * 1000, 0, 3, "ls", "-ld", path);
         return result.getStdout();
     }
 
diff --git a/tests/src/com/android/tradefed/util/LogcatUpdaterEventParserTest.java b/tests/src/com/android/tradefed/util/LogcatUpdaterEventParserTest.java
index 231d87c..361e6aa 100644
--- a/tests/src/com/android/tradefed/util/LogcatUpdaterEventParserTest.java
+++ b/tests/src/com/android/tradefed/util/LogcatUpdaterEventParserTest.java
@@ -111,10 +111,8 @@
         assertEquals(mParser.parseEventType(unmappedLogLine), null);
     }
 
-    /**
-     * Test that a thread exits its wait loop when it sees any event.
-     */
-    public void testWaitForEvent_any() {
+    /** Test that a thread exits its wait loop when it sees any event. */
+    public void testWaitForEvent_any() throws Exception {
         Thread waitThread = new Thread(new Runnable() {
             @Override
             public void run() {
@@ -132,10 +130,8 @@
         waitAndAssertTerminated(waitThread, logLines);
     }
 
-    /**
-     * Test that a thread exits its wait loop when it sees a specific expected event.
-     */
-    public void testWaitForEvent_specific() {
+    /** Test that a thread exits its wait loop when it sees a specific expected event. */
+    public void testWaitForEvent_specific() throws Exception {
         Thread waitThread = new Thread(new Runnable() {
             @Override
             public void run() {
@@ -153,7 +149,7 @@
         waitAndAssertTerminated(waitThread, logLines);
     }
 
-    private void waitAndAssertTerminated(Thread waitThread, String[] logLines) {
+    private void waitAndAssertTerminated(Thread waitThread, String[] logLines) throws Exception {
         waitThread.start();
         for (String line : logLines) {
             try {
@@ -164,6 +160,7 @@
         }
         // Allow short time for thread to switch state.
         RunUtil.getDefault().sleep(SHORT_WAIT_MS);
+        waitThread.join(5000);
         assertEquals(Thread.State.TERMINATED, waitThread.getState());
     }
 }
diff --git a/tests/src/com/android/tradefed/util/RunUtilFuncTest.java b/tests/src/com/android/tradefed/util/RunUtilFuncTest.java
index 4913780..11fea45 100644
--- a/tests/src/com/android/tradefed/util/RunUtilFuncTest.java
+++ b/tests/src/com/android/tradefed/util/RunUtilFuncTest.java
@@ -31,8 +31,9 @@
  */
 public class RunUtilFuncTest extends TestCase {
 
-    private static final long VERY_SHORT_TIMEOUT_MS = 10;
-    private static final long SHORT_TIMEOUT_MS = 500;
+    private static final long VERY_SHORT_TIMEOUT_MS = 10l;
+    private static final long SHORT_TIMEOUT_MS = 500l;
+    private static final long LONG_TIMEOUT_MS = 5000l;
 
     private abstract class MyRunnable implements IRunUtil.IRunnableResult {
         boolean mCanceled = false;
@@ -97,8 +98,8 @@
      */
     public void testRunTimedCmd_repeatedOutput() {
         for (int i = 0; i < 1000; i++) {
-            CommandResult result = RunUtil.getDefault().runTimedCmd(SHORT_TIMEOUT_MS, "echo",
-                    "hello");
+            CommandResult result =
+                    RunUtil.getDefault().runTimedCmd(LONG_TIMEOUT_MS, "echo", "hello");
             assertTrue("Failed at iteration " + i,
                     CommandStatus.SUCCESS.equals(result.getStatus()));
             CLog.d(result.getStdout());
@@ -123,12 +124,17 @@
             }
             s.close();
 
-            final long timeOut = 5000;
             // FIXME: this test case is not ideal, as it will only work on platforms that support
             // cat command.
-            CommandResult result = RunUtil.getDefault().runTimedCmd(timeOut, "cat",
-                    f.getAbsolutePath());
-            assertTrue(result.getStatus() == CommandStatus.SUCCESS);
+            CommandResult result =
+                    RunUtil.getDefault()
+                            .runTimedCmd(3 * LONG_TIMEOUT_MS, "cat", f.getAbsolutePath());
+            assertEquals(
+                    String.format(
+                            "We expected SUCCESS but got %s, with stdout: '%s'\nstderr: %s",
+                            result.getStatus(), result.getStdout(), result.getStderr()),
+                    CommandStatus.SUCCESS,
+                    result.getStatus());
             assertTrue(result.getStdout().length() == dataSize);
         } finally {
             f.delete();
@@ -146,7 +152,12 @@
         // printenv
         CommandResult result =
                 runUtil.runTimedCmdRetry(SHORT_TIMEOUT_MS, SHORT_TIMEOUT_MS, 3, "printenv", "bar");
-        assertEquals(CommandStatus.SUCCESS, result.getStatus());
+        assertEquals(
+                String.format(
+                        "We expected SUCCESS but got %s, with stdout: '%s'\nstderr: %s",
+                        result.getStatus(), result.getStdout(), result.getStderr()),
+                CommandStatus.SUCCESS,
+                result.getStatus());
         assertEquals("foo", result.getStdout().trim());
 
         // remove env variable
@@ -165,7 +176,12 @@
         RunUtil runUtil = new RunUtil();
         String[] command = {"sleep", "10000"};
         CommandResult result = runUtil.runTimedCmd(VERY_SHORT_TIMEOUT_MS, command);
-        assertEquals(CommandStatus.TIMED_OUT, result.getStatus());
+        assertEquals(
+                String.format(
+                        "We expected TIMED_OUT but got %s, with stdout: '%s'\nstderr: %s",
+                        result.getStatus(), result.getStdout(), result.getStderr()),
+                CommandStatus.TIMED_OUT,
+                result.getStatus());
         assertEquals("", result.getStdout());
         assertEquals("", result.getStderr());
         // We give it some times to clean up the process
diff --git a/tests/src/com/android/tradefed/util/SerializationUtilTest.java b/tests/src/com/android/tradefed/util/SerializationUtilTest.java
index f992069a..aab84cb 100644
--- a/tests/src/com/android/tradefed/util/SerializationUtilTest.java
+++ b/tests/src/com/android/tradefed/util/SerializationUtilTest.java
@@ -18,6 +18,7 @@
 import static org.junit.Assert.*;
 
 import com.android.tradefed.build.BuildInfo;
+import com.android.tradefed.build.BuildSerializedVersion;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -39,6 +40,7 @@
      * Test class that implements {@link Serializable} but has an attribute that is not serializable
      */
     public static class SerialTestClass implements Serializable {
+        private static final long serialVersionUID = BuildSerializedVersion.VERSION;
         public InputStream mStream;
 
         public SerialTestClass() {
diff --git a/tests/src/com/android/tradefed/util/SystemUtilTest.java b/tests/src/com/android/tradefed/util/SystemUtilTest.java
index 3eb3c70..5b670af 100644
--- a/tests/src/com/android/tradefed/util/SystemUtilTest.java
+++ b/tests/src/com/android/tradefed/util/SystemUtilTest.java
@@ -34,8 +34,8 @@
 public class SystemUtilTest {
 
     /**
-     * test {@link SystemUtil#getTestCasesDirs()} to make sure it gets directories from environment
-     * variables.
+     * test {@link SystemUtil#getTestCasesDirs(IBuildInfo)} to make sure it gets directories from
+     * environment variables.
      */
     @Test
     public void testGetTestCasesDirs() throws IOException {
@@ -49,7 +49,7 @@
             Mockito.when(SystemUtil.singleton.getEnv(SystemUtil.ENV_ANDROID_TARGET_OUT_TESTCASES))
                     .thenReturn(targetOutDir.getAbsolutePath());
 
-            Set<File> testCasesDirs = new HashSet<File>(SystemUtil.getTestCasesDirs());
+            Set<File> testCasesDirs = new HashSet<File>(SystemUtil.getTestCasesDirs(null));
             assertTrue(testCasesDirs.contains(targetOutDir));
             assertTrue(!testCasesDirs.contains(hostOutDir));
         } finally {
@@ -59,8 +59,8 @@
     }
 
     /**
-     * test {@link SystemUtil#getTestCasesDirs()} to make sure no exception thrown if no environment
-     * variable is set or the directory does not exist.
+     * test {@link SystemUtil#getTestCasesDirs(IBuildInfo)} to make sure no exception thrown if no
+     * environment variable is set or the directory does not exist.
      */
     @Test
     public void testGetTestCasesDirsNoDir() {
@@ -70,7 +70,7 @@
         Mockito.when(SystemUtil.singleton.getEnv(SystemUtil.ENV_ANDROID_TARGET_OUT_TESTCASES))
                 .thenReturn(targetOutDir.getAbsolutePath());
 
-        Set<File> testCasesDirs = new HashSet<File>(SystemUtil.getTestCasesDirs());
+        Set<File> testCasesDirs = new HashSet<File>(SystemUtil.getTestCasesDirs(null));
         assertEquals(testCasesDirs.size(), 0);
     }
 }
diff --git a/tests/src/com/android/tradefed/util/TarUtilTest.java b/tests/src/com/android/tradefed/util/TarUtilTest.java
index 55231f3..cc9c267 100644
--- a/tests/src/com/android/tradefed/util/TarUtilTest.java
+++ b/tests/src/com/android/tradefed/util/TarUtilTest.java
@@ -15,6 +15,8 @@
  */
 package com.android.tradefed.util;
 
+import static org.junit.Assert.*;
+
 import com.android.tradefed.log.ITestLogger;
 import com.android.tradefed.result.LogDataType;
 
@@ -25,6 +27,7 @@
 import org.junit.Test;
 
 import java.io.File;
+import java.io.FileNotFoundException;
 import java.io.InputStream;
 import java.util.List;
 
@@ -111,4 +114,41 @@
             FileUtil.deleteFile(logTarGzFile);
         }
     }
+
+    /**
+     * Test that {@link TarUtil#gzip(File)} is properly zipping the file and can be unzipped to
+     * recover the original file.
+     */
+    @Test
+    public void testGzipDir_unGzip() throws Exception {
+        final String content = "I'LL BE ZIPPED";
+        File tmpFile = FileUtil.createTempFile("base_file", ".txt", mWorkDir);
+        File zipped = null;
+        File unzipped = FileUtil.createTempDir("unzip-test-dir", mWorkDir);
+        try {
+            FileUtil.writeToFile(content, tmpFile);
+            zipped = TarUtil.gzip(tmpFile);
+            assertTrue(zipped.exists());
+            assertTrue(zipped.getName().endsWith(".gz"));
+            // unzip the file to ensure our utility can go both way
+            TarUtil.unGzip(zipped, unzipped);
+            assertEquals(1, unzipped.list().length);
+            // the original file is found
+            assertEquals(content, FileUtil.readStringFromFile(unzipped.listFiles()[0]));
+        } finally {
+            FileUtil.recursiveDelete(zipped);
+            FileUtil.recursiveDelete(unzipped);
+        }
+    }
+
+    /** Test to ensure that {@link TarUtil#gzip(File)} properly throws if the file is not valid. */
+    @Test
+    public void testGzip_invalidFile() throws Exception {
+        try {
+            TarUtil.gzip(new File("i_do_not_exist"));
+            fail("Should have thrown an exception.");
+        } catch (FileNotFoundException expected) {
+            // expected
+        }
+    }
 }
diff --git a/tests/src/com/android/tradefed/util/net/HttpHelperTest.java b/tests/src/com/android/tradefed/util/net/HttpHelperTest.java
index e47edfc..bfb5a7b 100644
--- a/tests/src/com/android/tradefed/util/net/HttpHelperTest.java
+++ b/tests/src/com/android/tradefed/util/net/HttpHelperTest.java
@@ -16,12 +16,18 @@
 
 package com.android.tradefed.util.net;
 
+import static org.mockito.Mockito.doNothing;
+
+import com.android.tradefed.util.IRunUtil;
 import com.android.tradefed.util.MultiMap;
+import com.android.tradefed.util.RunUtil;
 import com.android.tradefed.util.StreamUtil;
 import com.android.tradefed.util.net.IHttpHelper.DataSizeException;
 
 import junit.framework.TestCase;
 
+import org.mockito.Mockito;
+
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
@@ -193,19 +199,28 @@
      */
     public void testDoGetWithRetry_retry() throws IOException, DataSizeException {
         mHelper.close();
-        mHelper = new TestHttpHelper() {
-            boolean mExceptionThrown = false;
+        RunUtil mockRunUtil = Mockito.spy(RunUtil.class);
+        mHelper =
+                new TestHttpHelper() {
+                    boolean mExceptionThrown = false;
 
-            @Override
-            public String doGet(String url) throws IOException, DataSizeException {
-                if (!mExceptionThrown) {
-                    mExceptionThrown = true;
-                    throw new IOException();
-                }
-                return super.doGet(url);
-            }
-        };
+                    @Override
+                    public IRunUtil getRunUtil() {
+                        return mockRunUtil;
+                    }
 
+                    @Override
+                    public String doGet(String url) throws IOException, DataSizeException {
+                        if (!mExceptionThrown) {
+                            mExceptionThrown = true;
+                            throw new IOException();
+                        }
+                        return super.doGet(url);
+                    }
+                };
+        mHelper.setMaxTime(5000);
+        // Avoid doing actual sleep in the retry
+        doNothing().when(mockRunUtil).sleep(Mockito.anyLong());
         assertEquals(TEST_DATA, mHelper.doGetWithRetry(TEST_URL_STRING));
     }