Add a package name provider interface

Use an interface to provide package names to the module generator to make adding support for new sources easier.

Test: atest :presubmit
Bug: 178392533
Change-Id: I5e32d0d240994b532f492b915f7b591da7a43edd
Merged-In: I5e32d0d240994b532f492b915f7b591da7a43edd
(cherry picked from commit 912ad5c665e60a23b172ddc78820dc92058c38a5)

7 files changed, 223 insertions, 37 deletions

diff --git a/harness/src/main/java/com/android/csuite/config/ModuleGenerator.java b/harness/src/main/java/com/android/csuite/config/ModuleGenerator.java
index b37e7ef..0760085 100644
--- a/harness/src/main/java/com/android/csuite/config/ModuleGenerator.java
+++ b/harness/src/main/java/com/android/csuite/config/ModuleGenerator.java
@@ -17,7 +17,10 @@
 package com.android.csuite.config;
 
 import com.android.compatibility.common.tradefed.build.CompatibilityBuildHelper;
+import com.android.csuite.core.PackageNameProvider;
 import com.android.tradefed.build.IBuildInfo;
+import com.android.tradefed.config.IConfiguration;
+import com.android.tradefed.config.IConfigurationReceiver;
 import com.android.tradefed.config.Option;
 import com.android.tradefed.config.Option.Importance;
 import com.android.tradefed.device.DeviceNotAvailableException;
@@ -33,6 +36,7 @@ import com.google.common.annotations.VisibleForTesting;
 import com.google.common.io.Resources;
 
 import java.io.IOException;
+import java.io.UncheckedIOException;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.FileSystem;
 import java.nio.file.FileSystems;
@@ -58,13 +62,33 @@ import java.util.Set;
  * go/sharding-hack-for-module-gen. Note that since the generate step is executed as a test instance
  * and cleanup step is executed as a target preparer, there should be no saved states between
  * generating and cleaning up module files.
+ *
+ * <p>This module generator collects package names from all PackageNameProvider objects specified in
+ * the test configs.
+ *
+ * <h2>Syntax and usage</h2>
+ *
+ * <p>References to package name providers in TradeFed test configs must have the following syntax:
+ *
+ * <blockquote>
+ *
+ * <b>&lt;object type="PACKAGE_NAME_PROVIDER" class="</b><i>provider_class_name</i><b>"/&gt;</b>
+ *
+ * </blockquote>
+ *
+ * where <i>provider_class_name</i> is the fully-qualified class name of an PackageNameProvider
+ * implementation class.
  */
 public final class ModuleGenerator
-        implements IRemoteTest, IShardableTest, IBuildReceiver, ITargetPreparer {
+        implements IRemoteTest,
+                IShardableTest,
+                IBuildReceiver,
+                ITargetPreparer,
+                IConfigurationReceiver {
 
     @VisibleForTesting static final String MODULE_FILE_EXTENSION = ".config";
     @VisibleForTesting static final String OPTION_TEMPLATE = "template";
-    @VisibleForTesting static final String OPTION_PACKAGE = "package";
+    @VisibleForTesting static final String PACKAGE_NAME_PROVIDER = "PACKAGE_NAME_PROVIDER";
     private static final String TEMPLATE_PACKAGE_PATTERN = "\\{package\\}";
     private static final Collection<IRemoteTest> NOT_SPLITABLE = null;
 
@@ -74,16 +98,15 @@ public final class ModuleGenerator
             importance = Importance.ALWAYS)
     private String mTemplate;
 
-    @Option(name = OPTION_PACKAGE, description = "App package names.")
-    private final Set<String> mPackages = new HashSet<>();
-
     private final TestDirectoryProvider mTestDirectoryProvider;
     private final ResourceLoader mResourceLoader;
     private final FileSystem mFileSystem;
     private IBuildInfo mBuildInfo;
+    private IConfiguration mConfiguration;
 
-    public ModuleGenerator(Set<String> mPackages) {
-        this(FileSystems.getDefault());
+    @Override
+    public void setConfiguration(IConfiguration configuration) {
+        mConfiguration = configuration;
     }
 
     public ModuleGenerator() {
@@ -132,7 +155,7 @@ public final class ModuleGenerator
             // Executes the generate step.
             generateModules();
         } catch (IOException e) {
-            throw new RuntimeException("Failed to generate modules", e);
+            throw new UncheckedIOException("Failed to generate modules", e);
         }
 
         return NOT_SPLITABLE;
@@ -145,14 +168,26 @@ public final class ModuleGenerator
         // preparer, it is not considered as a IBuildReceiver instance.
         mBuildInfo = testInfo.getBuildInfo();
 
-        // Executes the clean up step.
-        cleanUpModules();
+        try {
+            // Executes the clean up step.
+            cleanUpModules();
+        } catch (IOException ioException) {
+            throw new UncheckedIOException("Failed to clean up generated modules", ioException);
+        }
+    }
+
+    private Set<String> getPackageNames() throws IOException {
+        Set<String> packages = new HashSet<>();
+        for (Object provider : mConfiguration.getConfigurationObjectList(PACKAGE_NAME_PROVIDER)) {
+            packages.addAll(((PackageNameProvider) provider).get());
+        }
+        return packages;
     }
 
     private void generateModules() throws IOException {
         String templateContent = mResourceLoader.load(mTemplate);
 
-        for (String packageName : mPackages) {
+        for (String packageName : getPackageNames()) {
             validatePackageName(packageName);
             Files.write(
                     getModulePath(packageName),
@@ -160,17 +195,19 @@ public final class ModuleGenerator
         }
     }
 
-    private void cleanUpModules() {
-        mPackages.forEach(
-                packageName -> {
-                    try {
-                        Files.delete(getModulePath(packageName));
-                    } catch (IOException ioException) {
-                        CLog.e(
-                                "Failed to delete the generated module for package " + packageName,
-                                ioException);
-                    }
-                });
+    private void cleanUpModules() throws IOException {
+        getPackageNames()
+                .forEach(
+                        packageName -> {
+                            try {
+                                Files.delete(getModulePath(packageName));
+                            } catch (IOException ioException) {
+                                CLog.e(
+                                        "Failed to delete the generated module for package "
+                                                + packageName,
+                                        ioException);
+                            }
+                        });
     }
 
     private Path getModulePath(String packageName) throws IOException {
diff --git a/harness/src/main/java/com/android/csuite/core/CommandLinePackageNameProvider.java b/harness/src/main/java/com/android/csuite/core/CommandLinePackageNameProvider.java
new file mode 100644
index 0000000..33426ed
--- /dev/null
+++ b/harness/src/main/java/com/android/csuite/core/CommandLinePackageNameProvider.java
@@ -0,0 +1,39 @@
+/*
+ * Copyright (C) 2021 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.csuite.core;
+
+import com.android.tradefed.config.Option;
+import com.android.tradefed.config.OptionClass;
+
+import com.google.common.annotations.VisibleForTesting;
+
+import java.util.HashSet;
+import java.util.Set;
+
+/** A package name provider that accepts package names via a command line option. */
+@OptionClass(alias = "command-line-package-name-provider")
+public final class CommandLinePackageNameProvider implements PackageNameProvider {
+    @VisibleForTesting static final String PACKAGE = "package";
+
+    @Option(name = PACKAGE, description = "App package names.")
+    private final Set<String> mPackages = new HashSet<>();
+
+    @Override
+    public Set<String> get() {
+        return mPackages;
+    }
+}
diff --git a/harness/src/main/java/com/android/csuite/core/PackageNameProvider.java b/harness/src/main/java/com/android/csuite/core/PackageNameProvider.java
new file mode 100644
index 0000000..05709a1
--- /dev/null
+++ b/harness/src/main/java/com/android/csuite/core/PackageNameProvider.java
@@ -0,0 +1,31 @@
+/*
+ * Copyright (C) 2021 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.csuite.core;
+
+import java.io.IOException;
+import java.util.Set;
+
+/** Provides a list of package names. */
+public interface PackageNameProvider {
+    /**
+     * Returns a set of package names.
+     *
+     * @return the package names. An empty set is returned if no package names are to be provided.
+     * @throws IOException if failed to get package names.
+     */
+    Set<String> get() throws IOException;
+}
diff --git a/harness/src/main/resources/config/csuite-base.xml b/harness/src/main/resources/config/csuite-base.xml
index 77b804c..33788ab 100644
--- a/harness/src/main/resources/config/csuite-base.xml
+++ b/harness/src/main/resources/config/csuite-base.xml
@@ -33,4 +33,6 @@
   </target_preparer>
   <!-- Cleans generated module files after test -->
   <target_preparer class="com.android.csuite.config.ModuleGenerator" />
+
+  <object type="PACKAGE_NAME_PROVIDER" class="com.android.csuite.core.CommandLinePackageNameProvider" />
 </configuration>
diff --git a/harness/src/test/java/com/android/csuite/CSuiteUnitTests.java b/harness/src/test/java/com/android/csuite/CSuiteUnitTests.java
index 096f882..c97736d 100644
--- a/harness/src/test/java/com/android/csuite/CSuiteUnitTests.java
+++ b/harness/src/test/java/com/android/csuite/CSuiteUnitTests.java
@@ -27,6 +27,7 @@ import org.junit.runners.Suite.SuiteClasses;
     com.android.compatibility.testtype.AppLaunchTestTest.class,
     com.android.csuite.config.AppRemoteFileResolverTest.class,
     com.android.csuite.config.ModuleGeneratorTest.class,
+    com.android.csuite.core.CommandLinePackageNameProviderTest.class,
     com.android.csuite.testing.CorrespondencesTest.class,
     com.android.csuite.testing.MoreAssertsTest.class,
 })
diff --git a/harness/src/test/java/com/android/csuite/config/ModuleGeneratorTest.java b/harness/src/test/java/com/android/csuite/config/ModuleGeneratorTest.java
index 6834bd8..a679828 100644
--- a/harness/src/test/java/com/android/csuite/config/ModuleGeneratorTest.java
+++ b/harness/src/test/java/com/android/csuite/config/ModuleGeneratorTest.java
@@ -20,7 +20,10 @@ import static com.google.common.truth.Truth.assertThat;
 
 import static org.testng.Assert.assertThrows;
 
+import com.android.csuite.core.PackageNameProvider;
 import com.android.tradefed.build.BuildInfo;
+import com.android.tradefed.config.Configuration;
+import com.android.tradefed.config.IConfiguration;
 import com.android.tradefed.config.OptionSetter;
 import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.invoker.IInvocationContext;
@@ -29,8 +32,8 @@ import com.android.tradefed.invoker.TestInformation;
 
 import com.google.common.collect.ArrayListMultimap;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ListMultimap;
-import com.google.common.jimfs.Configuration;
 import com.google.common.jimfs.Jimfs;
 import com.google.common.truth.IterableSubject;
 import com.google.common.truth.StringSubject;
@@ -41,12 +44,15 @@ import org.junit.runners.JUnit4;
 import org.mockito.Mockito;
 
 import java.io.IOException;
+import java.io.UncheckedIOException;
 import java.nio.file.FileSystem;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 @RunWith(JUnit4.class)
 public final class ModuleGeneratorTest {
@@ -55,16 +61,16 @@ public final class ModuleGeneratorTest {
     private static final String PACKAGE_PLACEHOLDER = "{package}";
     private static final Exception NO_EXCEPTION = null;
 
-    private final FileSystem mFileSystem = Jimfs.newFileSystem(Configuration.unix());
+    private final FileSystem mFileSystem =
+            Jimfs.newFileSystem(com.google.common.jimfs.Configuration.unix());
 
     @Test
-    public void tearDown_packageOptionIsSet_deletesGeneratedModules() throws Exception {
+    public void tearDown_packageNamesProvided_deletesGeneratedModules() throws Exception {
         Path testsDir = createTestsDir();
         ModuleGenerator generator1 =
                 createGeneratorBuilder()
                         .setTestsDir(testsDir)
                         .addPackage(TEST_PACKAGE_NAME1)
-                        .addPackage(TEST_PACKAGE_NAME1) // Simulate duplicate package option
                         .addPackage(TEST_PACKAGE_NAME2)
                         .build();
         generator1.split();
@@ -72,7 +78,6 @@ public final class ModuleGeneratorTest {
                 createGeneratorBuilder()
                         .setTestsDir(testsDir)
                         .addPackage(TEST_PACKAGE_NAME1)
-                        .addPackage(TEST_PACKAGE_NAME1) // Simulate duplicate package option
                         .addPackage(TEST_PACKAGE_NAME2)
                         .build();
 
@@ -82,7 +87,7 @@ public final class ModuleGeneratorTest {
     }
 
     @Test
-    public void tearDown_packageOptionIsNotSet_doesNotThrowError() throws Exception {
+    public void tearDown_packageNamesNotProvided_doesNotThrowError() throws Exception {
         ModuleGenerator generator = createGeneratorBuilder().setTestsDir(createTestsDir()).build();
         generator.split();
 
@@ -105,14 +110,15 @@ public final class ModuleGeneratorTest {
     }
 
     @Test
-    public void split_packageOptionContainsDuplicates_ignoreDuplicates() throws Exception {
+    public void split_multiplePackageNameProviders_generateModulesForAll() throws Exception {
         Path testsDir = createTestsDir();
         ModuleGenerator generator =
                 createGeneratorBuilder()
                         .setTestsDir(testsDir)
-                        .addPackage(TEST_PACKAGE_NAME1)
-                        .addPackage(TEST_PACKAGE_NAME1) // Simulate duplicate package option
-                        .addPackage(TEST_PACKAGE_NAME2)
+                        .addPackageNameProvider(() -> ImmutableSet.of(TEST_PACKAGE_NAME1))
+                        // Simulates package providers providing duplicated package names.
+                        .addPackageNameProvider(() -> ImmutableSet.of(TEST_PACKAGE_NAME1))
+                        .addPackageNameProvider(() -> ImmutableSet.of(TEST_PACKAGE_NAME2))
                         .build();
 
         generator.split();
@@ -124,7 +130,22 @@ public final class ModuleGeneratorTest {
     }
 
     @Test
-    public void split_packageOptionNotSet_doesNotGenerate() throws Exception {
+    public void split_packageNameProviderThrowsException_throwsException() throws Exception {
+        Path testsDir = createTestsDir();
+        ModuleGenerator generator =
+                createGeneratorBuilder()
+                        .setTestsDir(testsDir)
+                        .addPackageNameProvider(
+                                () -> {
+                                    throw new IOException();
+                                })
+                        .build();
+
+        assertThrows(UncheckedIOException.class, () -> generator.split());
+    }
+
+    @Test
+    public void split_packageNamesNotProvided_doesNotGenerate() throws Exception {
         Path testsDir = createTestsDir();
         ModuleGenerator generator = createGeneratorBuilder().setTestsDir(testsDir).build();
 
@@ -230,7 +251,8 @@ public final class ModuleGeneratorTest {
 
     private static final class GeneratorBuilder {
         private final ListMultimap<String, String> mOptions = ArrayListMultimap.create();
-        private final List<String> mPackages = new ArrayList<>();
+        private final Set<String> mPackages = new HashSet<>();
+        private final List<PackageNameProvider> mPackageNameProviders = new ArrayList<>();
         private Path mTestsDir;
         private String mTemplateContent;
         private FileSystem mFileSystem;
@@ -240,6 +262,11 @@ public final class ModuleGeneratorTest {
             return this;
         }
 
+        GeneratorBuilder addPackageNameProvider(PackageNameProvider packageNameProvider) {
+            mPackageNameProviders.add(packageNameProvider);
+            return this;
+        }
+
         GeneratorBuilder setFileSystem(FileSystem fileSystem) {
             mFileSystem = fileSystem;
             return this;
@@ -270,9 +297,13 @@ public final class ModuleGeneratorTest {
                 optionSetter.setOptionValue(entry.getKey(), entry.getValue());
             }
 
-            for (String packageName : mPackages) {
-                optionSetter.setOptionValue(ModuleGenerator.OPTION_PACKAGE, packageName);
-            }
+            List<PackageNameProvider> packageNameProviders = new ArrayList<>(mPackageNameProviders);
+            packageNameProviders.add(() -> mPackages);
+
+            IConfiguration configuration = new Configuration("name", "description");
+            configuration.setConfigurationObjectList(
+                    ModuleGenerator.PACKAGE_NAME_PROVIDER, packageNameProviders);
+            generator.setConfiguration(configuration);
 
             return generator;
         }
diff --git a/harness/src/test/java/com/android/csuite/core/CommandLinePackageNameProviderTest.java b/harness/src/test/java/com/android/csuite/core/CommandLinePackageNameProviderTest.java
new file mode 100644
index 0000000..41e6dc1
--- /dev/null
+++ b/harness/src/test/java/com/android/csuite/core/CommandLinePackageNameProviderTest.java
@@ -0,0 +1,45 @@
+/*
+ * Copyright (C) 2021 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.csuite.core;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.android.tradefed.config.OptionSetter;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+import java.util.Set;
+
+@RunWith(JUnit4.class)
+public final class CommandLinePackageNameProviderTest {
+
+    @Test
+    public void get_packageNamesProvided_returnsPackageNames() throws Exception {
+        CommandLinePackageNameProvider provider = new CommandLinePackageNameProvider();
+        String package1 = "package.name1";
+        String package2 = "package.name2";
+        OptionSetter optionSetter = new OptionSetter(provider);
+        optionSetter.setOptionValue(CommandLinePackageNameProvider.PACKAGE, package1);
+        optionSetter.setOptionValue(CommandLinePackageNameProvider.PACKAGE, package2);
+
+        Set<String> packageNames = provider.get();
+
+        assertThat(packageNames).containsExactly(package1, package2);
+    }
+}