Merge "Start removing some of the GCS test resource logic"
diff --git a/src/com/android/tradefed/device/NativeDevice.java b/src/com/android/tradefed/device/NativeDevice.java
index fd00a3c..45124bd 100644
--- a/src/com/android/tradefed/device/NativeDevice.java
+++ b/src/com/android/tradefed/device/NativeDevice.java
@@ -3161,8 +3161,8 @@
         // Track Tradefed reboot time
         mLastTradefedRebootTime = System.currentTimeMillis();
 
-        if (TestDeviceState.FASTBOOT == getDeviceState()) {
-            CLog.i("device %s in fastboot. Rebooting to userspace.", getSerialNumber());
+        if (isStateBootloaderOrFastbootd()) {
+            CLog.i("device %s in %s. Rebooting to userspace.", getSerialNumber(), getDeviceState());
             executeFastbootCommand("reboot");
         } else {
             if (mOptions.shouldDisableReboot()) {
@@ -3730,7 +3730,7 @@
         if (!deviceState.equals(getDeviceState())) {
             // disable state changes while fastboot lock is held, because issuing fastboot command
             // will disrupt state
-            if (getDeviceState().equals(TestDeviceState.FASTBOOT) && mFastbootLock.isLocked()) {
+            if (isStateBootloaderOrFastbootd() && mFastbootLock.isLocked()) {
                 return;
             }
             mState = deviceState;
@@ -4869,7 +4869,7 @@
         if (getIDevice() instanceof StubDevice) {
             return null;
         }
-        if (TestDeviceState.FASTBOOT.equals(getDeviceState())) {
+        if (isStateBootloaderOrFastbootd()) {
             return null;
         }
         try {
diff --git a/src/com/android/tradefed/device/WaitDeviceRecovery.java b/src/com/android/tradefed/device/WaitDeviceRecovery.java
index 66d25c3..e163025 100644
--- a/src/com/android/tradefed/device/WaitDeviceRecovery.java
+++ b/src/com/android/tradefed/device/WaitDeviceRecovery.java
@@ -120,10 +120,13 @@
         // ensure bootloader state is updated
         monitor.waitForDeviceBootloaderStateUpdate();
 
-        if (monitor.getDeviceState().equals(TestDeviceState.FASTBOOT)) {
-            Log.i(LOG_TAG, String.format(
-                    "Found device %s in fastboot but expected online. Rebooting...",
-                    monitor.getSerialNumber()));
+        TestDeviceState state = monitor.getDeviceState();
+        if (TestDeviceState.FASTBOOT.equals(state) || TestDeviceState.FASTBOOTD.equals(state)) {
+            Log.i(
+                    LOG_TAG,
+                    String.format(
+                            "Found device %s in %s but expected online. Rebooting...",
+                            monitor.getSerialNumber(), state));
             // TODO: retry if failed
             getRunUtil().runTimedCmd(mFastbootWaitTime, mFastbootPath, "-s",
                     monitor.getSerialNumber(), "reboot");
diff --git a/src/com/android/tradefed/invoker/InvocationExecution.java b/src/com/android/tradefed/invoker/InvocationExecution.java
index a1dcecf..78a1efa 100644
--- a/src/com/android/tradefed/invoker/InvocationExecution.java
+++ b/src/com/android/tradefed/invoker/InvocationExecution.java
@@ -42,6 +42,7 @@
 import com.android.tradefed.invoker.logger.InvocationMetricLogger.InvocationMetricKey;
 import com.android.tradefed.invoker.logger.TfObjectTracker;
 import com.android.tradefed.invoker.shard.IShardHelper;
+import com.android.tradefed.invoker.shard.TestsPoolPoller;
 import com.android.tradefed.log.ITestLogger;
 import com.android.tradefed.log.LogUtil.CLog;
 import com.android.tradefed.result.ByteArrayInputStreamSource;
@@ -585,7 +586,9 @@
                 // Handle the no-retry use case
                 if (!decision.isAutoRetryEnabled()
                         || RetryStrategy.NO_RETRY.equals(decision.getRetryStrategy())
-                        || test instanceof ITestSuite) {
+                        || test instanceof ITestSuite
+                        // TODO: Handle auto-retry in local-sharding for non-suite
+                        || test instanceof TestsPoolPoller) {
                     runTest(config, info, listener, test);
                     remainingTests.remove(test);
                     continue;
diff --git a/test_framework/com/android/tradefed/targetprep/InstallApexModuleTargetPreparer.java b/test_framework/com/android/tradefed/targetprep/InstallApexModuleTargetPreparer.java
index 9edbd10..31ffcc6 100644
--- a/test_framework/com/android/tradefed/targetprep/InstallApexModuleTargetPreparer.java
+++ b/test_framework/com/android/tradefed/targetprep/InstallApexModuleTargetPreparer.java
@@ -222,15 +222,14 @@
      * Extracts and returns splits for the specified apks.
      *
      * @param testInfo the {@link TestInformation}
-     * @param apksName The name of the apks file to extract splits from.
+     * @param moduleFile The module file to extract the splits from.
      * @return a File[] containing the splits.
      * @throws TargetSetupError if bundletool cannot be found or device spec file fails to generate.
      */
-    private File[] getSplitsForApks(TestInformation testInfo, String apksName)
+    private File[] getSplitsForApks(TestInformation testInfo, File moduleFile)
             throws TargetSetupError {
         initBundletoolUtil(testInfo);
         initDeviceSpecFilePath(testInfo.getDevice());
-        File moduleFile = getLocalPathForFilename(testInfo, apksName);
         File splitsDir =
                 getBundletoolUtil()
                         .extractSplitsFromApks(
@@ -274,7 +273,7 @@
             }
             String modulePackageName = "";
             if (moduleFile.getName().endsWith(SPLIT_APKS_SUFFIX)) {
-                File[] splits = getSplitsForApks(testInfo, moduleFileName);
+                File[] splits = getSplitsForApks(testInfo, moduleFile);
                 if (splits == null) {
                     // Bundletool failed to extract splits.
                     CLog.w(
@@ -437,7 +436,7 @@
             throws TargetSetupError, DeviceNotAvailableException {
         File apks = getLocalPathForFilename(testInfo, apksName);
         // Rename the extracted files and add the file to filename list.
-        File[] splits = getSplitsForApks(testInfo, apks.getName());
+        File[] splits = getSplitsForApks(testInfo, apks);
         ITestDevice device = testInfo.getDevice();
         if (splits.length == 0) {
             throw new TargetSetupError(
@@ -471,7 +470,7 @@
         for (String moduleFileName : testAppFileNames) {
             File moduleFile = getLocalPathForFilename(testInfo, moduleFileName);
             if (moduleFileName.endsWith(SPLIT_APKS_SUFFIX)) {
-                File[] splits = getSplitsForApks(testInfo, moduleFileName);
+                File[] splits = getSplitsForApks(testInfo, moduleFile);
                 String splitsArgs = createInstallArgsForSplit(splits, device);
                 mSplitsInstallArgs.add(splitsArgs);
             } else {