Allow Test Mapping to run tests based on additional test_mappings.zip.

Bug: 203176715
Bug: 201797814
Test: unittests
Change-Id: I41c56df7a33ee16897a71fe74bc8d045f25916b3
diff --git a/common_util/com/android/tradefed/result/error/InfraErrorIdentifier.java b/common_util/com/android/tradefed/result/error/InfraErrorIdentifier.java
index 883c8ac..9c5f1ca 100644
--- a/common_util/com/android/tradefed/result/error/InfraErrorIdentifier.java
+++ b/common_util/com/android/tradefed/result/error/InfraErrorIdentifier.java
@@ -80,6 +80,7 @@
     CONFIGURATION_NOT_FOUND(505_253, FailureStatus.CUSTOMER_ISSUE),
     UNEXPECTED_DEVICE_CONFIGURED(505_254, FailureStatus.CUSTOMER_ISSUE),
     KEYSTORE_CONFIG_ERROR(505_255, FailureStatus.DEPENDENCY_ISSUE),
+    TEST_MAPPING_PATH_COLLISION(505_256, FailureStatus.DEPENDENCY_ISSUE),
 
     UNDETERMINED(510_000, FailureStatus.UNSET);
 
diff --git a/javatests/com/android/tradefed/testtype/suite/TestMappingSuiteRunnerTest.java b/javatests/com/android/tradefed/testtype/suite/TestMappingSuiteRunnerTest.java
index ae4d94d..0862114 100644
--- a/javatests/com/android/tradefed/testtype/suite/TestMappingSuiteRunnerTest.java
+++ b/javatests/com/android/tradefed/testtype/suite/TestMappingSuiteRunnerTest.java
@@ -18,6 +18,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -35,6 +36,7 @@
 import com.android.tradefed.config.OptionSetter;
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.error.HarnessRuntimeException;
 import com.android.tradefed.invoker.IInvocationContext;
 import com.android.tradefed.invoker.InvocationContext;
 import com.android.tradefed.invoker.TestInformation;
@@ -941,7 +943,7 @@
             tempDir = FileUtil.createTempDir("test_mapping");
             tempTestsDir = FileUtil.createTempDir("test_mapping_testcases");
 
-            File zipFile = createTestMappingZip(tempDir);
+            File zipFile = createMainlineTestMappingZip(tempDir);
             createMainlineModuleConfig(tempTestsDir.getAbsolutePath());
 
             IDeviceBuildInfo mockBuildInfo = mock(IDeviceBuildInfo.class);
@@ -1100,6 +1102,96 @@
         }
     }
 
+    @Test
+    public void testLoadTests_WithCollisionAdditionalTestMappingZip() throws Exception {
+        File tempDir = null;
+        try {
+            mOptionSetter.setOptionValue("test-mapping-test-group", "presubmit");
+            mOptionSetter.setOptionValue("additional-test-mapping-zip", "extra-zip");
+
+            tempDir = FileUtil.createTempDir("test_mapping");
+            File zipFile = createTestMappingZip(tempDir);
+
+            IDeviceBuildInfo mockBuildInfo = mock(IDeviceBuildInfo.class);
+            when(mockBuildInfo.getFile(BuildInfoFileKey.TARGET_LINKED_DIR)).thenReturn(null);
+            when(mockBuildInfo.getTestsDir()).thenReturn(new File("non-existing-dir"));
+            when(mockBuildInfo.getFile(TEST_MAPPINGS_ZIP)).thenReturn(zipFile);
+            when(mockBuildInfo.getFile("extra-zip")).thenReturn(zipFile);
+            mRunner.setBuild(mockBuildInfo);
+
+            LinkedHashMap<String, IConfiguration> configMap = mRunner.loadTests();
+            fail("Should have thrown an exception.");
+        } catch (HarnessRuntimeException expected) {
+            // expected
+            assertTrue(expected.getMessage().contains("Collision of Test Mapping file"));
+        } finally {
+            FileUtil.recursiveDelete(tempDir);
+            TestMapping.setIgnoreTestMappingImports(true);
+            TestMapping.setTestMappingPaths(new ArrayList<String>());
+        }
+    }
+
+    @Test
+    public void testLoadTests_WithoutCollisionAdditionalTestMappingZip() throws Exception {
+        File tempDir = null;
+        File tempDir2 = null;
+        try {
+            mOptionSetter.setOptionValue("test-mapping-test-group", "presubmit");
+            mOptionSetter.setOptionValue("additional-test-mapping-zip", "extra-zip");
+
+            tempDir = FileUtil.createTempDir("test_mapping");
+            tempDir2 = FileUtil.createTempDir("test_mapping");
+            File zipFile = createTestMappingZip(tempDir);
+            File zipFile2 = createTestMappingZip(tempDir2);
+
+            IDeviceBuildInfo mockBuildInfo = mock(IDeviceBuildInfo.class);
+            when(mockBuildInfo.getFile(BuildInfoFileKey.TARGET_LINKED_DIR)).thenReturn(null);
+            when(mockBuildInfo.getTestsDir()).thenReturn(new File("non-existing-dir"));
+            when(mockBuildInfo.getFile(TEST_MAPPINGS_ZIP)).thenReturn(zipFile);
+            when(mockBuildInfo.getFile("extra-zip")).thenReturn(zipFile2);
+            mRunner.setBuild(mockBuildInfo);
+
+            LinkedHashMap<String, IConfiguration> configMap = mRunner.loadTests();
+            assertEquals(2, configMap.size());
+            verify(mockBuildInfo, times(1)).getFile(TEST_MAPPINGS_ZIP);
+            verify(mockBuildInfo, times(1)).getFile("extra-zip");
+        } finally {
+            FileUtil.recursiveDelete(tempDir);
+            FileUtil.recursiveDelete(tempDir2);
+            TestMapping.setIgnoreTestMappingImports(true);
+            TestMapping.setTestMappingPaths(new ArrayList<String>());
+        }
+    }
+
+    @Test
+    public void testLoadTests_WithMissingAdditionalTestMappingZips() throws Exception {
+        File tempDir = null;
+        try {
+            mOptionSetter.setOptionValue("test-mapping-test-group", "presubmit");
+            mOptionSetter.setOptionValue("additional-test-mapping-zip", "extra-zip");
+
+            tempDir = FileUtil.createTempDir("test_mapping");
+            File zipFile = createTestMappingZip(tempDir);
+
+            IDeviceBuildInfo mockBuildInfo = mock(IDeviceBuildInfo.class);
+            when(mockBuildInfo.getFile(BuildInfoFileKey.TARGET_LINKED_DIR)).thenReturn(null);
+            when(mockBuildInfo.getTestsDir()).thenReturn(new File("non-existing-dir"));
+            when(mockBuildInfo.getFile(TEST_MAPPINGS_ZIP)).thenReturn(zipFile);
+            when(mockBuildInfo.getFile("extra-zip")).thenReturn(null);
+            mRunner.setBuild(mockBuildInfo);
+            LinkedHashMap<String, IConfiguration> configMap = mRunner.loadTests();
+            fail("Should have thrown an exception.");
+        } catch (HarnessRuntimeException expected) {
+            // expected
+            assertEquals(
+                "Missing extra-zip in the BuildInfo file.", expected.getMessage());
+        } finally {
+            FileUtil.recursiveDelete(tempDir);
+            TestMapping.setIgnoreTestMappingImports(true);
+            TestMapping.setTestMappingPaths(new ArrayList<String>());
+        }
+    }
+
     /** Helper to create specific test infos. */
     private TestInfo createTestInfo(String name, String source) {
         TestInfo info = new TestInfo(name, source, false);
@@ -1109,6 +1201,28 @@
         return info;
     }
 
+    /** Helper to create test_mappings.zip for Mainline. */
+    private File createMainlineTestMappingZip(File tempDir) throws IOException {
+        File srcDir = FileUtil.createTempDir("src", tempDir);
+        String srcFile = File.separator + TEST_DATA_DIR + File.separator + DISABLED_PRESUBMIT_TESTS;
+        InputStream resourceStream = this.getClass().getResourceAsStream(srcFile);
+        FileUtil.saveResourceFile(resourceStream, srcDir, DISABLED_PRESUBMIT_TESTS);
+
+        srcFile = File.separator + TEST_DATA_DIR + File.separator + "test_mapping_with_mainline";
+        resourceStream = this.getClass().getResourceAsStream(srcFile);
+        FileUtil.saveResourceFile(resourceStream, srcDir, TEST_MAPPING);
+        File subDir = FileUtil.createTempDir("sub_dir", srcDir);
+        srcFile = File.separator + TEST_DATA_DIR + File.separator + "test_mapping_2";
+        resourceStream = this.getClass().getResourceAsStream(srcFile);
+        FileUtil.saveResourceFile(resourceStream, subDir, TEST_MAPPING);
+
+        List<File> filesToZip = Arrays.asList(srcDir, new File(tempDir, DISABLED_PRESUBMIT_TESTS));
+        File zipFile = Paths.get(tempDir.getAbsolutePath(), TEST_MAPPINGS_ZIP).toFile();
+        ZipUtil.createZip(filesToZip, zipFile);
+
+        return zipFile;
+    }
+
     /** Helper to create test_mappings.zip. */
     private File createTestMappingZip(File tempDir) throws IOException {
         File srcDir = FileUtil.createTempDir("src", tempDir);
@@ -1116,7 +1230,7 @@
         InputStream resourceStream = this.getClass().getResourceAsStream(srcFile);
         FileUtil.saveResourceFile(resourceStream, srcDir, DISABLED_PRESUBMIT_TESTS);
 
-        srcFile = File.separator + TEST_DATA_DIR + File.separator + "test_mapping_with_mainline";
+        srcFile = File.separator + TEST_DATA_DIR + File.separator + "test_mapping_1";
         resourceStream = this.getClass().getResourceAsStream(srcFile);
         FileUtil.saveResourceFile(resourceStream, srcDir, TEST_MAPPING);
         File subDir = FileUtil.createTempDir("sub_dir", srcDir);
diff --git a/javatests/com/android/tradefed/util/testmapping/TestMappingTest.java b/javatests/com/android/tradefed/util/testmapping/TestMappingTest.java
index b5bd16c..4b27c2e 100644
--- a/javatests/com/android/tradefed/util/testmapping/TestMappingTest.java
+++ b/javatests/com/android/tradefed/util/testmapping/TestMappingTest.java
@@ -759,6 +759,53 @@
         }
     }
 
+    /** Test for {@link TestMapping#getTests()} for loading tests from a given test_mappings.zip. */
+    @Test
+    public void testGetTestsWithGivenFilePath() throws Exception {
+        File tempDir = null;
+        IBuildInfo mockBuildInfo = mock(IBuildInfo.class);
+        try {
+            tempDir = FileUtil.createTempDir("test_mapping");
+            File srcDir = FileUtil.createTempDir("src", tempDir);
+            File subDir = FileUtil.createTempDir("sub_dir", srcDir);
+            createTestMapping(srcDir, "test_mapping_1", TEST_MAPPING);
+            createTestMapping(subDir, "test_mapping_2", TEST_MAPPING);
+            createTestMapping(tempDir, DISABLED_PRESUBMIT_TESTS, DISABLED_PRESUBMIT_TESTS);
+            List<File> filesToZip =
+                Arrays.asList(srcDir, new File(tempDir, DISABLED_PRESUBMIT_TESTS));
+            File zipFile = Paths.get(tempDir.getAbsolutePath(), TEST_MAPPINGS_ZIP).toFile();
+            ZipUtil.createZip(filesToZip, zipFile);
+
+            // Ensure the static variable doesn't have any relative path configured.
+            TestMapping.setTestMappingPaths(new ArrayList<String>());
+            Set<TestInfo> tests = TestMapping.getTests(
+                    mockBuildInfo, "presubmit", false, null, zipFile);
+            assertEquals(0, tests.size());
+
+            tests = TestMapping.getTests(mockBuildInfo, "presubmit", true, null, zipFile);
+            assertEquals(2, tests.size());
+            Set<String> names = new HashSet<String>();
+            for (TestInfo test : tests) {
+                names.add(test.getName());
+                if (test.getName().equals("test1")) {
+                    assertTrue(test.getHostOnly());
+                } else {
+                    assertFalse(test.getHostOnly());
+                }
+            }
+            assertTrue(!names.contains("suite/stub1"));
+            assertTrue(names.contains("test1"));
+        } finally {
+            FileUtil.recursiveDelete(tempDir);
+        }
+    }
+
+    private void createTestMapping(File srcDir, String srcName, String dstName) throws Exception {
+        String srcFile = File.separator + TEST_DATA_DIR + File.separator + srcName;
+        InputStream resourceStream = this.getClass().getResourceAsStream(srcFile);
+        FileUtil.saveResourceFile(resourceStream, srcDir, dstName);
+    }
+
     private String getJsonStringByName(String fileName) throws Exception {
         File tempDir = null;
         try {
diff --git a/src/com/android/tradefed/testtype/suite/TestMappingSuiteRunner.java b/src/com/android/tradefed/testtype/suite/TestMappingSuiteRunner.java
index b1c4185..3a51257 100644
--- a/src/com/android/tradefed/testtype/suite/TestMappingSuiteRunner.java
+++ b/src/com/android/tradefed/testtype/suite/TestMappingSuiteRunner.java
@@ -15,10 +15,13 @@
  */
 package com.android.tradefed.testtype.suite;
 
+import com.android.tradefed.build.IBuildInfo;
 import com.android.tradefed.config.ConfigurationDescriptor;
 import com.android.tradefed.config.IConfiguration;
 import com.android.tradefed.config.Option;
+import com.android.tradefed.error.HarnessRuntimeException;
 import com.android.tradefed.log.LogUtil.CLog;
+import com.android.tradefed.result.error.InfraErrorIdentifier;
 import com.android.tradefed.testtype.IAbi;
 import com.android.tradefed.testtype.IRemoteTest;
 import com.android.tradefed.util.FileUtil;
@@ -111,11 +114,21 @@
                             + "filtered by allowed tests.")
     private Set<String> mAllowedTestLists = new HashSet<>();
 
+    @Option(
+            name = "additional-test-mapping-zip",
+            description =
+                    "A list of additional test_mappings.zip that contains TEST_MAPPING files. The "
+                            + "runner will collect tests based on them. If none  is specified, "
+                            + "only the tests on the triggering device build will be run.")
+    private List<String> mAdditionalTestMappingZip = new ArrayList<>();
+
     /** Special definition in the test mapping structure. */
     private static final String TEST_MAPPING_INCLUDE_FILTER = "include-filter";
 
     private static final String TEST_MAPPING_EXCLUDE_FILTER = "exclude-filter";
 
+    private IBuildInfo mBuildInfo;
+
     /**
      * Load the tests configuration that will be run. Each tests is defined by a {@link
      * IConfiguration} and a unique name under which it will report results. There are 2 ways to
@@ -137,6 +150,7 @@
         // Name of the tests
         Set<String> testNames = new HashSet<>();
         Set<TestInfo> testInfosToRun = new HashSet<>();
+        mBuildInfo = getBuildInfo();
         if (mTestGroup == null && includeFilter.isEmpty()) {
             throw new RuntimeException(
                     "At least one of the options, --test-mapping-test-group or --include-filter, "
@@ -169,7 +183,28 @@
             }
             testInfosToRun =
                     TestMapping.getTests(
-                            getBuildInfo(), mTestGroup, getPrioritizeHostConfig(), mKeywords);
+                            mBuildInfo, mTestGroup, getPrioritizeHostConfig(), mKeywords);
+            if (!mAdditionalTestMappingZip.isEmpty()) {
+                for (String zipName : mAdditionalTestMappingZip) {
+                    File zipFile = mBuildInfo.getFile(zipName);
+                    if (zipFile == null) {
+                        throw new HarnessRuntimeException(
+                                String.format("Missing %s in the BuildInfo file.", zipName),
+                                InfraErrorIdentifier.ARTIFACT_NOT_FOUND);
+                    }
+                    CLog.i("Getting tests from additional test mapping zip: %s", zipName);
+                    Set<TestInfo> additionalTests =
+                            TestMapping.getTests(
+                                    mBuildInfo,
+                                    mTestGroup,
+                                    getPrioritizeHostConfig(),
+                                    mKeywords,
+                                    zipFile
+                            );
+                    validateTestMappingSource(testInfosToRun, additionalTests, zipName);
+                    testInfosToRun.addAll(additionalTests);
+                }
+            }
             if (!mTestModulesForced.isEmpty()) {
                 CLog.i("Filtering tests for the given names: %s", mTestModulesForced);
                 testInfosToRun =
@@ -247,6 +282,24 @@
         return mUseTestMappingPath;
     }
 
+    /** Ensure there are no collisions of TEST_MAPPING paths between different test mapping zips. */
+    private void validateTestMappingSource(Set<TestInfo> base, Set<TestInfo> target, String name) {
+        Set<String> baseSorces = new HashSet<>();
+        for (TestInfo testInfo : base) {
+            baseSorces.addAll(testInfo.getSources());
+        }
+        for (TestInfo testInfo : target) {
+            for (String src : testInfo.getSources()) {
+                if (baseSorces.contains(src)) {
+                    throw new HarnessRuntimeException(
+                            String.format("Collision of Test Mapping file: %s/TEST_MAPPING in " +
+                                    "artifact: %s.", src, name),
+                            InfraErrorIdentifier.TEST_MAPPING_PATH_COLLISION);
+                }
+            }
+        }
+    }
+
     /**
      * Create individual tests with test infos for a module.
      *
diff --git a/src/com/android/tradefed/util/testmapping/TestMapping.java b/src/com/android/tradefed/util/testmapping/TestMapping.java
index b62ef41..4dfc4b1 100644
--- a/src/com/android/tradefed/util/testmapping/TestMapping.java
+++ b/src/com/android/tradefed/util/testmapping/TestMapping.java
@@ -428,8 +428,7 @@
     }
 
     /**
-     * Helper to find all tests in all TEST_MAPPING files. This is needed when a suite run requires
-     * to run all tests in TEST_MAPPING files for a given group, e.g., presubmit.
+     * Helper to find all tests in all TEST_MAPPING files based on a artifact in the device build.
      *
      * @param buildInfo the {@link IBuildInfo} describing the build.
      * @param testGroup a {@link String} of the test group.
@@ -437,11 +436,33 @@
      *     returned. false to return tests that require device to run.
      * @return A {@code Set<TestInfo>} of tests set in the build artifact, test_mappings.zip.
      */
-    @SuppressWarnings("StreamResourceLeak")
     public static Set<TestInfo> getTests(
             IBuildInfo buildInfo, String testGroup, boolean hostOnly, Set<String> keywords) {
+        File zipFile = buildInfo.getFile(TEST_MAPPINGS_ZIP);
+        return getTests(buildInfo, testGroup, hostOnly, keywords, zipFile);
+    }
+
+    /**
+     * Helper to find all tests in all TEST_MAPPING files based on the given artifact. This is
+     * needed when a suite run requires to run all tests in TEST_MAPPING files for a given group,
+     * e.g., presubmit.
+     *
+     * @param buildInfo the {@link IBuildInfo} describing the build.
+     * @param testGroup a {@link String} of the test group.
+     * @param hostOnly true if only tests running on host and don't require device should be
+     *     returned. false to return tests that require device to run.
+     * @param zipFile the {@link File} of the test mapping zip.
+     * @return A {@code Set<TestInfo>} of tests set in the build artifact, test_mappings.zip.
+     */
+    @SuppressWarnings("StreamResourceLeak")
+    public static Set<TestInfo> getTests(
+        IBuildInfo buildInfo,
+        String testGroup,
+        boolean hostOnly,
+        Set<String> keywords,
+        File zipFile) {
         Set<TestInfo> tests = new HashSet<TestInfo>();
-        File testMappingsDir = extractTestMappingsZip(buildInfo.getFile(TEST_MAPPINGS_ZIP));
+        File testMappingsDir = extractTestMappingsZip(zipFile);
         Stream<Path> stream = null;
         try {
             Path testMappingsRootPath = Paths.get(testMappingsDir.getAbsolutePath());