Merge "Update how to push a directory to remote"
diff --git a/src/com/android/tradefed/config/DynamicRemoteFileResolver.java b/src/com/android/tradefed/config/DynamicRemoteFileResolver.java
index 119302c..a10b6c7 100644
--- a/src/com/android/tradefed/config/DynamicRemoteFileResolver.java
+++ b/src/com/android/tradefed/config/DynamicRemoteFileResolver.java
@@ -30,6 +30,7 @@
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * Class that helps resolving path to remote files.
@@ -39,12 +40,14 @@
  */
 public class DynamicRemoteFileResolver {
 
+    public static final String DYNAMIC_RESOLVER = "dynamic-resolver";
     private static final Map<String, IRemoteFileResolver> PROTOCOL_SUPPORT = new HashMap<>();
 
     static {
-        // TODO: Have a way to dynamically specify more support
         PROTOCOL_SUPPORT.put(GcsRemoteFileResolver.PROTOCOL, new GcsRemoteFileResolver());
     }
+    // The configuration map being static, we only need to update it once per TF instance.
+    private static AtomicBoolean sIsUpdateDone = new AtomicBoolean(false);
 
     private Map<String, OptionFieldsForName> mOptionMap;
 
@@ -131,9 +134,32 @@
 
     @VisibleForTesting
     IRemoteFileResolver getResolver(String protocol) {
+        if (updateProtocols()) {
+            IGlobalConfiguration globalConfig = getGlobalConfig();
+            Object o = globalConfig.getConfigurationObject(DYNAMIC_RESOLVER);
+            if (o != null) {
+                if (o instanceof IRemoteFileResolver) {
+                    IRemoteFileResolver resolver = (IRemoteFileResolver) o;
+                    CLog.d("Adding %s to supported remote file resolver", resolver);
+                    PROTOCOL_SUPPORT.put(resolver.getSupportedProtocol(), resolver);
+                } else {
+                    CLog.e("%s is not of type IRemoteFileResolver", o);
+                }
+            }
+        }
         return PROTOCOL_SUPPORT.get(protocol);
     }
 
+    @VisibleForTesting
+    boolean updateProtocols() {
+        return sIsUpdateDone.compareAndSet(false, true);
+    }
+
+    @VisibleForTesting
+    IGlobalConfiguration getGlobalConfig() {
+        return GlobalConfiguration.getInstance();
+    }
+
     private File resolveRemoteFiles(File consideredFile, Option option)
             throws ConfigurationException {
         String path = consideredFile.getPath();
diff --git a/src/com/android/tradefed/config/remote/GcsRemoteFileResolver.java b/src/com/android/tradefed/config/remote/GcsRemoteFileResolver.java
index 248543a..afaa5be 100644
--- a/src/com/android/tradefed/config/remote/GcsRemoteFileResolver.java
+++ b/src/com/android/tradefed/config/remote/GcsRemoteFileResolver.java
@@ -24,6 +24,8 @@
 
 import java.io.File;
 
+import javax.annotation.Nonnull;
+
 /** Implementation of {@link IRemoteFileResolver} that allows downloading from a GCS bucket. */
 public class GcsRemoteFileResolver implements IRemoteFileResolver {
 
@@ -47,6 +49,11 @@
         }
     }
 
+    @Override
+    public @Nonnull String getSupportedProtocol() {
+        return PROTOCOL;
+    }
+
     @VisibleForTesting
     protected GCSDownloaderHelper getDownloader() {
         if (mHelper == null) {
diff --git a/src/com/android/tradefed/config/remote/IRemoteFileResolver.java b/src/com/android/tradefed/config/remote/IRemoteFileResolver.java
index a65a047..3d54ecb 100644
--- a/src/com/android/tradefed/config/remote/IRemoteFileResolver.java
+++ b/src/com/android/tradefed/config/remote/IRemoteFileResolver.java
@@ -38,4 +38,7 @@
      */
     public @Nonnull File resolveRemoteFiles(File consideredFile, Option option)
             throws ConfigurationException;
+
+    /** Returns the associated protocol supported for download. */
+    public @Nonnull String getSupportedProtocol();
 }
diff --git a/src/com/android/tradefed/testtype/suite/ModuleDefinition.java b/src/com/android/tradefed/testtype/suite/ModuleDefinition.java
index 4f54224..5dcd7e5 100644
--- a/src/com/android/tradefed/testtype/suite/ModuleDefinition.java
+++ b/src/com/android/tradefed/testtype/suite/ModuleDefinition.java
@@ -335,9 +335,26 @@
         // Setup
         long prepStartTime = getCurrentTime();
 
-        for (String deviceName : mModuleInvocationContext.getDeviceConfigNames()) {
+        for (int i = 0; i < mModuleInvocationContext.getDeviceConfigNames().size(); i++) {
+            String deviceName = mModuleInvocationContext.getDeviceConfigNames().get(i);
             ITestDevice device = mModuleInvocationContext.getDevice(deviceName);
-            for (ITargetPreparer preparer : mPreparersPerDevice.get(deviceName)) {
+            if (i >= mPreparersPerDevice.size()) {
+                CLog.d(
+                        "Main configuration has more devices than the module configuration. '%s' "
+                                + "will not run any preparation.",
+                        deviceName);
+                continue;
+            }
+            List<ITargetPreparer> preparers = mPreparersPerDevice.get(deviceName);
+            if (preparers == null) {
+                CLog.w(
+                        "Module configuration devices mismatch the main configuration "
+                                + "(Missing device '%s'), resolving preparers by index.",
+                        deviceName);
+                String key = new ArrayList<>(mPreparersPerDevice.keySet()).get(i);
+                preparers = mPreparersPerDevice.get(key);
+            }
+            for (ITargetPreparer preparer : preparers) {
                 preparationException =
                         runPreparerSetup(
                                 device,
@@ -743,9 +760,25 @@
             multiCleaner.tearDown(mModuleInvocationContext, setupException);
         }
 
-        for (String deviceName : mModuleInvocationContext.getDeviceConfigNames()) {
+        for (int i = 0; i < mModuleInvocationContext.getDeviceConfigNames().size(); i++) {
+            String deviceName = mModuleInvocationContext.getDeviceConfigNames().get(i);
             ITestDevice device = mModuleInvocationContext.getDevice(deviceName);
+            if (i >= mPreparersPerDevice.size()) {
+                CLog.d(
+                        "Main configuration has more devices than the module configuration. '%s' "
+                                + "will not run any tear down.",
+                        deviceName);
+                continue;
+            }
             List<ITargetPreparer> preparers = mPreparersPerDevice.get(deviceName);
+            if (preparers == null) {
+                CLog.w(
+                        "Module configuration devices mismatch the main configuration "
+                                + "(Missing device '%s'), resolving preparers by index.",
+                        deviceName);
+                String key = new ArrayList<>(mPreparersPerDevice.keySet()).get(i);
+                preparers = mPreparersPerDevice.get(key);
+            }
             ListIterator<ITargetPreparer> itr = preparers.listIterator(preparers.size());
             while (itr.hasPrevious()) {
                 ITargetPreparer preparer = itr.previous();
diff --git a/tests/src/com/android/tradefed/config/DynamicRemoteFileResolverTest.java b/tests/src/com/android/tradefed/config/DynamicRemoteFileResolverTest.java
index 7f0cc62..544cd72 100644
--- a/tests/src/com/android/tradefed/config/DynamicRemoteFileResolverTest.java
+++ b/tests/src/com/android/tradefed/config/DynamicRemoteFileResolverTest.java
@@ -35,6 +35,8 @@
 import java.util.Iterator;
 import java.util.Set;
 
+import javax.annotation.Nonnull;
+
 /** Unit tests for {@link DynamicRemoteFileResolver}. */
 @RunWith(JUnit4.class)
 public class DynamicRemoteFileResolverTest {
@@ -68,6 +70,12 @@
                         }
                         return null;
                     }
+
+                    @Override
+                    boolean updateProtocols() {
+                        // Do not set the static variable
+                        return false;
+                    }
                 };
     }
 
@@ -228,4 +236,83 @@
         }
         EasyMock.verify(mMockResolver);
     }
+
+    /** Allow extra resolver to be provided via global configuration. */
+    @Test
+    public void testResolve_addition() throws Exception {
+        mResolver =
+                new DynamicRemoteFileResolver() {
+                    @Override
+                    boolean updateProtocols() {
+                        // Do not set the static variable
+                        return true;
+                    }
+
+                    @Override
+                    IGlobalConfiguration getGlobalConfig() {
+                        IGlobalConfiguration config;
+                        try {
+                            config =
+                                    ConfigurationFactory.getInstance()
+                                            .createGlobalConfigurationFromArgs(
+                                                    new String[] {"empty"}, new ArrayList<>());
+                            // Add a custom object to the resolving map
+                            config.setConfigurationObject(
+                                    DynamicRemoteFileResolver.DYNAMIC_RESOLVER,
+                                    new IRemoteFileResolver() {
+
+                                        @Override
+                                        public @Nonnull File resolveRemoteFiles(
+                                                File consideredFile, Option option)
+                                                throws ConfigurationException {
+                                            return mMockResolver.resolveRemoteFiles(
+                                                    consideredFile, option);
+                                        }
+
+                                        @Override
+                                        public @Nonnull String getSupportedProtocol() {
+                                            return "fakeprotocol";
+                                        }
+                                    });
+
+                        } catch (ConfigurationException e) {
+                            throw new RuntimeException(e);
+                        }
+                        return config;
+                    }
+                };
+        RemoteFileOption object = new RemoteFileOption();
+        OptionSetter setter =
+                new OptionSetter(object) {
+                    @Override
+                    DynamicRemoteFileResolver createResolver() {
+                        return mResolver;
+                    }
+                };
+
+        File fake = FileUtil.createTempFile("gs-option-setter-test", "txt");
+
+        setter.setOptionValue("remote-file", "fakeprotocol://fake/path");
+        assertEquals("fakeprotocol:/fake/path", object.remoteFile.getPath());
+
+        EasyMock.expect(
+                        mMockResolver.resolveRemoteFiles(
+                                EasyMock.eq(new File("fakeprotocol:/fake/path")),
+                                EasyMock.anyObject()))
+                .andReturn(fake);
+        EasyMock.replay(mMockResolver);
+
+        Set<File> downloadedFile = setter.validateRemoteFilePath();
+        try {
+            assertEquals(1, downloadedFile.size());
+            File downloaded = downloadedFile.iterator().next();
+            // The file has been replaced by the downloaded one.
+            assertEquals(downloaded.getAbsolutePath(), object.remoteFile.getAbsolutePath());
+        } finally {
+            for (File f : downloadedFile) {
+                FileUtil.recursiveDelete(f);
+            }
+        }
+        EasyMock.verify(mMockResolver);
+    }
 }
diff --git a/tests/src/com/android/tradefed/testtype/suite/ModuleDefinitionMultiTest.java b/tests/src/com/android/tradefed/testtype/suite/ModuleDefinitionMultiTest.java
index c8c8db4..b456260 100644
--- a/tests/src/com/android/tradefed/testtype/suite/ModuleDefinitionMultiTest.java
+++ b/tests/src/com/android/tradefed/testtype/suite/ModuleDefinitionMultiTest.java
@@ -17,6 +17,7 @@
 
 import com.android.tradefed.build.IBuildInfo;
 import com.android.tradefed.config.Configuration;
+import com.android.tradefed.config.ConfigurationDef;
 import com.android.tradefed.config.DeviceConfigurationHolder;
 import com.android.tradefed.config.IConfiguration;
 import com.android.tradefed.config.IDeviceConfiguration;
@@ -108,9 +109,6 @@
     public void testCreateAndRun() throws Exception {
         // Add a preparer to the second device
         mMapDeviceTargetPreparer.get(DEVICE_NAME_2).add(mMockTargetPrep);
-        mMultiDeviceConfiguration
-                .getDeviceConfigByName(DEVICE_NAME_2)
-                .addSpecificConfig(mMockTargetPrep);
 
         mModule =
                 new ModuleDefinition(
@@ -126,7 +124,7 @@
         mModule.getModuleInvocationContext().addDeviceBuildInfo(DEVICE_NAME_2, mBuildInfo2);
 
         mListener.testRunStarted(MODULE_NAME, 0);
-        mListener.testRunEnded(EasyMock.anyLong(), (HashMap<String, Metric>) EasyMock.anyObject());
+        mListener.testRunEnded(EasyMock.anyLong(), EasyMock.<HashMap<String, Metric>>anyObject());
 
         // Target preparation is triggered against the preparer in the second device.
         EasyMock.expect(mMockTargetPrep.isDisabled()).andReturn(false);
@@ -136,4 +134,39 @@
         mModule.run(mListener);
         verifyMocks();
     }
+
+    /** Create a single device module running against a multi device main configuration. */
+    @Test
+    public void testPreparer_mismatch() throws Exception {
+        // The module configuration only contains the default device.
+        mMapDeviceTargetPreparer.clear();
+        List<ITargetPreparer> preparers = new ArrayList<>();
+        preparers.add(mMockTargetPrep);
+        mMapDeviceTargetPreparer.put(ConfigurationDef.DEFAULT_DEVICE_NAME, preparers);
+
+        mModule =
+                new ModuleDefinition(
+                        MODULE_NAME,
+                        mTestList,
+                        mMapDeviceTargetPreparer,
+                        new ArrayList<>(),
+                        mMultiDeviceConfiguration);
+        // Simulate injection of devices from ITestSuite
+        mModule.getModuleInvocationContext().addAllocatedDevice(DEVICE_NAME_1, mDevice1);
+        mModule.getModuleInvocationContext().addDeviceBuildInfo(DEVICE_NAME_1, mBuildInfo1);
+        mModule.getModuleInvocationContext().addAllocatedDevice(DEVICE_NAME_2, mDevice2);
+        mModule.getModuleInvocationContext().addDeviceBuildInfo(DEVICE_NAME_2, mBuildInfo2);
+
+        mListener.testRunStarted(MODULE_NAME, 0);
+        mListener.testRunEnded(EasyMock.anyLong(), EasyMock.<HashMap<String, Metric>>anyObject());
+
+        // Target preparation is of first device in module configuration is triggered against the
+        // first device from main configuration
+        EasyMock.expect(mMockTargetPrep.isDisabled()).andReturn(false);
+        mMockTargetPrep.setUp(mDevice1, mBuildInfo1);
+
+        replayMocks();
+        mModule.run(mListener);
+        verifyMocks();
+    }
 }