Allow extra args to the dynamic resolver queries

Test: unit tests
Bug: 149516126
Change-Id: I622c3a81fd000f9b8fa63ece1c3a5d4a6acae37c
diff --git a/src/com/android/tradefed/command/CommandOptions.java b/src/com/android/tradefed/command/CommandOptions.java
index f529844..7ce3d1a 100644
--- a/src/com/android/tradefed/command/CommandOptions.java
+++ b/src/com/android/tradefed/command/CommandOptions.java
@@ -23,7 +23,9 @@
 import com.android.tradefed.log.LogUtil.CLog;
 import com.android.tradefed.util.UniqueMultiMap;
 
+import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
+import java.util.Map;
 import java.util.Set;
 
 /**
@@ -201,6 +203,13 @@
             description = "Feature flag to release the device as soon as done with it.")
     private boolean mEnableEarlyDeviceRelease = false;
 
+    @Option(
+            name = "dynamic-download-args",
+            description =
+                    "Extra args passed to the IRemoteFileResolver interface for dynamic download "
+                            + "in the queryArgs.")
+    private Map<String, String> mDynamicDownloadArgs = new LinkedHashMap<>();
+
     /**
      * Set the help mode for the config.
      * <p/>
@@ -511,4 +520,10 @@
     public boolean earlyDeviceRelease() {
         return mEnableEarlyDeviceRelease;
     }
+
+    /** {@inheritDoc} */
+    @Override
+    public Map<String, String> getDynamicDownloadArgs() {
+        return mDynamicDownloadArgs;
+    }
 }
diff --git a/src/com/android/tradefed/command/ICommandOptions.java b/src/com/android/tradefed/command/ICommandOptions.java
index 4064904..16d30f0 100644
--- a/src/com/android/tradefed/command/ICommandOptions.java
+++ b/src/com/android/tradefed/command/ICommandOptions.java
@@ -19,6 +19,7 @@
 import com.android.tradefed.device.metric.AutoLogCollector;
 import com.android.tradefed.util.UniqueMultiMap;
 
+import java.util.Map;
 import java.util.Set;
 
 /**
@@ -192,4 +193,7 @@
 
     /** Whether or not to release the device early when done with it. */
     public boolean earlyDeviceRelease();
+
+    /** Returns the map of args to pass to the dynamic download query. */
+    public Map<String, String> getDynamicDownloadArgs();
 }
diff --git a/src/com/android/tradefed/config/DynamicRemoteFileResolver.java b/src/com/android/tradefed/config/DynamicRemoteFileResolver.java
index 56467a6..3b9e778 100644
--- a/src/com/android/tradefed/config/DynamicRemoteFileResolver.java
+++ b/src/com/android/tradefed/config/DynamicRemoteFileResolver.java
@@ -74,6 +74,8 @@
     public static final String OPTIONAL_KEY = "optional";
 
     private Map<String, OptionFieldsForName> mOptionMap;
+    // Populated from {@link ICommandOptions#getDynamicDownloadArgs()}
+    private Map<String, String> mExtraArgs = new LinkedHashMap<>();
     private ITestDevice mDevice;
 
     /** Sets the map of options coming from {@link OptionSetter} */
@@ -86,6 +88,11 @@
         mDevice = device;
     }
 
+    /** Add extra args for the query. */
+    public void addExtraArgs(Map<String, String> extraArgs) {
+        mExtraArgs.putAll(extraArgs);
+    }
+
     /**
      * Runs through all the {@link File} option type and check if their path should be resolved.
      *
@@ -357,6 +364,8 @@
                 CLog.d(
                         "Considering option '%s' with path: '%s' for download.",
                         option.name(), path);
+                // Overrides query args
+                query.putAll(mExtraArgs);
                 return resolver.resolveRemoteFiles(fileToResolve, query);
             } catch (BuildRetrievalError e) {
                 if (isOptional(query)) {
diff --git a/src/com/android/tradefed/invoker/TestInvocation.java b/src/com/android/tradefed/invoker/TestInvocation.java
index 0e1e8ac..a2737f6 100644
--- a/src/com/android/tradefed/invoker/TestInvocation.java
+++ b/src/com/android/tradefed/invoker/TestInvocation.java
@@ -694,6 +694,7 @@
             if (!RunMode.REMOTE_INVOCATION.equals(mode)) {
                 DynamicRemoteFileResolver resolver = new DynamicRemoteFileResolver();
                 resolver.setDevice(context.getDevices().get(0));
+                resolver.addExtraArgs(config.getCommandOptions().getDynamicDownloadArgs());
                 config.resolveDynamicOptions(resolver);
             }
             return true;
diff --git a/src/com/android/tradefed/invoker/shard/ShardHelper.java b/src/com/android/tradefed/invoker/shard/ShardHelper.java
index 205b6b3..591f82e 100644
--- a/src/com/android/tradefed/invoker/shard/ShardHelper.java
+++ b/src/com/android/tradefed/invoker/shard/ShardHelper.java
@@ -229,6 +229,7 @@
         config.validateOptions();
         DynamicRemoteFileResolver resolver = new DynamicRemoteFileResolver();
         resolver.setDevice(testInfo.getDevice());
+        resolver.addExtraArgs(config.getCommandOptions().getDynamicDownloadArgs());
         config.resolveDynamicOptions(resolver);
     }
 
diff --git a/src/com/android/tradefed/invoker/shard/TestsPoolPoller.java b/src/com/android/tradefed/invoker/shard/TestsPoolPoller.java
index 4344ffe..2f80337 100644
--- a/src/com/android/tradefed/invoker/shard/TestsPoolPoller.java
+++ b/src/com/android/tradefed/invoker/shard/TestsPoolPoller.java
@@ -205,6 +205,8 @@
                     validationConfig.setTest(test);
                     validationConfig.validateOptions();
                     DynamicRemoteFileResolver resolver = new DynamicRemoteFileResolver();
+                    resolver.addExtraArgs(
+                            validationConfig.getCommandOptions().getDynamicDownloadArgs());
                     resolver.setDevice(info.getDevice());
                     validationConfig.resolveDynamicOptions(resolver);
                     // Set the configuration after the validation, otherwise we override the config
diff --git a/src/com/android/tradefed/testtype/suite/ITestSuite.java b/src/com/android/tradefed/testtype/suite/ITestSuite.java
index 02c69f4..068607b 100644
--- a/src/com/android/tradefed/testtype/suite/ITestSuite.java
+++ b/src/com/android/tradefed/testtype/suite/ITestSuite.java
@@ -466,6 +466,8 @@
         // Ignore config file as it's part of config zip artifact that's staged already.
         List<String> excludeFilters = Arrays.asList("[.]config$");
         mDynamicResolver.setDevice(device);
+        mDynamicResolver.addExtraArgs(
+                mMainConfiguration.getCommandOptions().getDynamicDownloadArgs());
         for (File remoteFile : mBuildInfo.getRemoteFiles()) {
             try {
                 mDynamicResolver.resolvePartialDownloadZip(
@@ -778,6 +780,8 @@
         module.setRetryDecision(decision);
 
         module.setEnableDynamicDownload(mEnableDynamicDownload);
+        module.addDynamicDownloadArgs(
+                mMainConfiguration.getCommandOptions().getDynamicDownloadArgs());
         // Actually run the module
         module.run(
                 moduleInfo,
diff --git a/src/com/android/tradefed/testtype/suite/ModuleDefinition.java b/src/com/android/tradefed/testtype/suite/ModuleDefinition.java
index 17ab3d5..9008ecb 100644
--- a/src/com/android/tradefed/testtype/suite/ModuleDefinition.java
+++ b/src/com/android/tradefed/testtype/suite/ModuleDefinition.java
@@ -399,6 +399,10 @@
                             .setInvocationContext(mModuleInvocationContext);
                 }
                 mInternalTestConfiguration = new Configuration("tmp-download", "tmp-download");
+                mInternalTestConfiguration
+                        .getCommandOptions()
+                        .getDynamicDownloadArgs()
+                        .putAll(mModuleConfiguration.getCommandOptions().getDynamicDownloadArgs());
                 // We do it before the official set, otherwise the IConfiguration will not be the
                 // right one.
                 mInternalTestConfiguration.setTest(test);
@@ -958,6 +962,10 @@
         mEnableDynamicDownload = enableDynamicDownload;
     }
 
+    public void addDynamicDownloadArgs(Map<String, String> extraArgs) {
+        mModuleConfiguration.getCommandOptions().getDynamicDownloadArgs().putAll(extraArgs);
+    }
+
     /**
      * Allow to load a module_controller object to tune how should a particular module run.
      *
@@ -1036,6 +1044,7 @@
             CLog.d("Attempting to resolve dynamic files from %s", getId());
             DynamicRemoteFileResolver resolver = new DynamicRemoteFileResolver();
             resolver.setDevice(device);
+            resolver.addExtraArgs(moduleConfiguration.getCommandOptions().getDynamicDownloadArgs());
             moduleConfiguration.resolveDynamicOptions(resolver);
             return null;
         } catch (RuntimeException | ConfigurationException | BuildRetrievalError e) {