Support importing other TEST_MAPPING files

This change allows atest to handle TEST_MAPPING files containing imports,
for example:
  "imports": [
    {
      "path": "../folder2"
    }
  ]
atest shall search for TEST_MAPPING files in the imported path and its
parent directories and include all tests found for the given test group.

Bug: 110166535
Test: unittest
Change-Id: I73bd80534bceefb5c9c688369306f85156e2de8e
diff --git a/atest/cli_translator.py b/atest/cli_translator.py
index aae117d..fb2461d 100644
--- a/atest/cli_translator.py
+++ b/atest/cli_translator.py
@@ -100,18 +100,27 @@
             test_mapping_file: Path to a TEST_MAPPING file.
 
         Returns:
-            A dictionary of all tests in the TEST_MAPPING file, grouped by test
-            group.
+            A tuple of (all_tests, imports), where
+            all_tests is a dictionary of all tests in the TEST_MAPPING file,
+                grouped by test group.
+            imports is a list of test_mapping.Import to include other test
+                mapping files.
         """
         all_tests = {}
+        imports = []
         test_mapping_dict = None
         with open(test_mapping_file) as json_file:
             test_mapping_dict = json.load(json_file)
         for test_group_name, test_list in test_mapping_dict.items():
-            grouped_tests = all_tests.setdefault(test_group_name, set())
-            grouped_tests.update(
-                [test_mapping.TestDetail(test) for test in test_list])
-        return all_tests
+            if test_group_name == constants.TEST_MAPPING_IMPORTS:
+                for import_detail in test_list:
+                    imports.append(
+                        test_mapping.Import(test_mapping_file, import_detail))
+            else:
+                grouped_tests = all_tests.setdefault(test_group_name, set())
+                grouped_tests.update(
+                    [test_mapping.TestDetail(test) for test in test_list])
+        return all_tests, imports
 
     def _find_files(self, path, file_name=TEST_MAPPING):
         """Find all files with given name under the given path.
@@ -138,17 +147,22 @@
             test_mapping_files: A list of path of TEST_MAPPING files.
 
         Returns:
-            A tuple of (tests, all_tests), where,
+            A tuple of (tests, all_tests, imports), where,
             tests is a set of tests (test_mapping.TestDetail) defined in
             TEST_MAPPING file of the given path, and its parent directories,
             with matching test_group.
             all_tests is a dictionary of all tests in TEST_MAPPING files,
             grouped by test group.
+            imports is a list of test_mapping.Import objects that contains the
+            details of where to import a TEST_MAPPING file.
         """
+        all_imports = []
         # Read and merge the tests in all TEST_MAPPING files.
         merged_all_tests = {}
         for test_mapping_file in test_mapping_files:
-            all_tests = self._read_tests_in_test_mapping(test_mapping_file)
+            all_tests, imports = self._read_tests_in_test_mapping(
+                test_mapping_file)
+            all_imports.extend(imports)
             for test_group_name, test_list in all_tests.items():
                 grouped_tests = merged_all_tests.setdefault(
                     test_group_name, set())
@@ -162,11 +176,13 @@
         elif test_group == constants.TEST_GROUP_ALL:
             for grouped_tests in merged_all_tests.values():
                 tests.update(grouped_tests)
-        return tests, merged_all_tests
+        return tests, merged_all_tests, all_imports
 
+    # pylint: disable=too-many-arguments
+    # pylint: disable=too-many-locals
     def _find_tests_by_test_mapping(
             self, path='', test_group=constants.TEST_GROUP_PRESUBMIT,
-            file_name=TEST_MAPPING, include_subdirs=False):
+            file_name=TEST_MAPPING, include_subdirs=False, checked_files=None):
         """Find tests defined in TEST_MAPPING in the given path.
 
         Args:
@@ -176,6 +192,7 @@
                 `TEST_MAPPING`. The argument is added for testing purpose.
             include_subdirs: True to include tests in TEST_MAPPING files in sub
                 directories.
+            checked_files: Paths of TEST_MAPPING files that have been checked.
 
         Returns:
             A tuple of (tests, all_tests), where,
@@ -187,23 +204,54 @@
         """
         path = os.path.realpath(path)
         test_mapping_files = set()
+        all_tests = {}
         test_mapping_file = os.path.join(path, file_name)
         if os.path.exists(test_mapping_file):
             test_mapping_files.add(test_mapping_file)
-        # Include all TEST_MAPPING files in parent directories if
-        # `include_subdirs` is set to True.
+        # Include all TEST_MAPPING files in sub-directories if `include_subdirs`
+        # is set to True.
         if include_subdirs:
             test_mapping_files.update(self._find_files(path, file_name))
         # Include all possible TEST_MAPPING files in parent directories.
-        while path != constants.ANDROID_BUILD_TOP and path != os.sep:
+        root_dir = os.environ.get(constants.ANDROID_BUILD_TOP, os.sep)
+        while path != root_dir and path != os.sep:
             path = os.path.dirname(path)
             test_mapping_file = os.path.join(path, file_name)
             if os.path.exists(test_mapping_file):
                 test_mapping_files.add(test_mapping_file)
 
-        return self._get_tests_from_test_mapping_files(
+        if checked_files is None:
+            checked_files = set()
+        test_mapping_files.difference_update(checked_files)
+        checked_files.update(test_mapping_files)
+        if not test_mapping_files:
+            return test_mapping_files, all_tests
+
+        tests, all_tests, imports = self._get_tests_from_test_mapping_files(
             test_group, test_mapping_files)
 
+        # Load TEST_MAPPING files from imports recursively.
+        if imports:
+            for import_detail in imports:
+                path = import_detail.get_path()
+                # (b/110166535 #19) Import path might not exist if a project is
+                # located in different directory in different branches.
+                if path is None:
+                    logging.warn(
+                        'Failed to import TEST_MAPPING at %s', import_detail)
+                    continue
+                # Search for tests based on the imported search path.
+                import_tests, import_all_tests = (
+                    self._find_tests_by_test_mapping(
+                        path, test_group, file_name, include_subdirs,
+                        checked_files))
+                # Merge the collections
+                tests.update(import_tests)
+                for group, grouped_tests in import_all_tests.items():
+                    all_tests.setdefault(group, set()).update(grouped_tests)
+
+        return tests, all_tests
+
     def _gather_build_targets(self, test_infos):
         targets = set()
         for test_info in test_infos:
@@ -233,7 +281,7 @@
 
         test_details, all_test_details = self._find_tests_by_test_mapping(
             path=src_path, test_group=test_group,
-            include_subdirs=args.include_subdirs)
+            include_subdirs=args.include_subdirs, checked_files=set())
         test_details_list = list(test_details)
         if not test_details_list:
             logging.warn(
diff --git a/atest/cli_translator_unittest.py b/atest/cli_translator_unittest.py
index fa99be2..893afe6 100755
--- a/atest/cli_translator_unittest.py
+++ b/atest/cli_translator_unittest.py
@@ -37,6 +37,12 @@
 TEST_2 = test_mapping.TestDetail({'name': 'test2'})
 TEST_3 = test_mapping.TestDetail({'name': 'test3'})
 TEST_4 = test_mapping.TestDetail({'name': 'test4'})
+TEST_5 = test_mapping.TestDetail({'name': 'test5'})
+TEST_6 = test_mapping.TestDetail({'name': 'test6'})
+TEST_7 = test_mapping.TestDetail({'name': 'test7'})
+TEST_8 = test_mapping.TestDetail({'name': 'test8'})
+TEST_9 = test_mapping.TestDetail({'name': 'test9'})
+TEST_10 = test_mapping.TestDetail({'name': 'test10'})
 
 SEARCH_DIR_RE = re.compile(r'^find ([^ ]*).*$')
 
@@ -180,49 +186,69 @@
 
     def test_find_tests_by_test_mapping_presubmit(self):
         """Test _find_tests_by_test_mapping method to locate presubmit tests."""
-        tests, all_tests = self.ctr._find_tests_by_test_mapping(
-            path=TEST_MAPPING_DIR, file_name='test_mapping_sample')
-        expected = set([TEST_1, TEST_2])
+        os_environ_mock = {constants.ANDROID_BUILD_TOP: uc.TEST_DATA_DIR}
+        with mock.patch.dict('os.environ', os_environ_mock, clear=True):
+            tests, all_tests = self.ctr._find_tests_by_test_mapping(
+                path=TEST_MAPPING_DIR, file_name='test_mapping_sample',
+                checked_files=set())
+        expected = set([TEST_1, TEST_2, TEST_5, TEST_7, TEST_9])
         expected_all_tests = {'presubmit': expected,
-                              'postsubmit': set([TEST_3]),
+                              'postsubmit': set(
+                                  [TEST_3, TEST_6, TEST_8, TEST_10]),
                               'other_group': set([TEST_4])}
         self.assertEqual(expected, tests)
         self.assertEqual(expected_all_tests, all_tests)
 
     def test_find_tests_by_test_mapping_postsubmit(self):
-        """Test _find_tests_by_test_mapping method to locate postsubmit tests."""
-        tests, all_tests = self.ctr._find_tests_by_test_mapping(
-            path=TEST_MAPPING_DIR, test_group=constants.TEST_GROUP_POSTSUBMIT,
-            file_name='test_mapping_sample')
-        expected_presubmit = set([TEST_1, TEST_2])
-        expected = set([TEST_1, TEST_2, TEST_3])
+        """Test _find_tests_by_test_mapping method to locate postsubmit tests.
+        """
+        os_environ_mock = {constants.ANDROID_BUILD_TOP: uc.TEST_DATA_DIR}
+        with mock.patch.dict('os.environ', os_environ_mock, clear=True):
+            tests, all_tests = self.ctr._find_tests_by_test_mapping(
+                path=TEST_MAPPING_DIR,
+                test_group=constants.TEST_GROUP_POSTSUBMIT,
+                file_name='test_mapping_sample', checked_files=set())
+        expected_presubmit = set([TEST_1, TEST_2, TEST_5, TEST_7, TEST_9])
+        expected = set(
+            [TEST_1, TEST_2, TEST_3, TEST_5, TEST_6, TEST_7, TEST_8, TEST_9,
+             TEST_10])
         expected_all_tests = {'presubmit': expected_presubmit,
-                              'postsubmit': set([TEST_3]),
+                              'postsubmit': set(
+                                  [TEST_3, TEST_6, TEST_8, TEST_10]),
                               'other_group': set([TEST_4])}
         self.assertEqual(expected, tests)
         self.assertEqual(expected_all_tests, all_tests)
 
     def test_find_tests_by_test_mapping_all_group(self):
-        """Test _find_tests_by_test_mapping method to locate postsubmit tests."""
-        tests, all_tests = self.ctr._find_tests_by_test_mapping(
-            path=TEST_MAPPING_DIR, test_group=constants.TEST_GROUP_ALL,
-            file_name='test_mapping_sample')
-        expected_presubmit = set([TEST_1, TEST_2])
-        expected = set([TEST_1, TEST_2, TEST_3, TEST_4])
+        """Test _find_tests_by_test_mapping method to locate postsubmit tests.
+        """
+        os_environ_mock = {constants.ANDROID_BUILD_TOP: uc.TEST_DATA_DIR}
+        with mock.patch.dict('os.environ', os_environ_mock, clear=True):
+            tests, all_tests = self.ctr._find_tests_by_test_mapping(
+                path=TEST_MAPPING_DIR, test_group=constants.TEST_GROUP_ALL,
+                file_name='test_mapping_sample', checked_files=set())
+        expected_presubmit = set([TEST_1, TEST_2, TEST_5, TEST_7, TEST_9])
+        expected = set([
+            TEST_1, TEST_2, TEST_3, TEST_4, TEST_5, TEST_6, TEST_7, TEST_8,
+            TEST_9, TEST_10])
         expected_all_tests = {'presubmit': expected_presubmit,
-                              'postsubmit': set([TEST_3]),
+                              'postsubmit': set(
+                                  [TEST_3, TEST_6, TEST_8, TEST_10]),
                               'other_group': set([TEST_4])}
         self.assertEqual(expected, tests)
         self.assertEqual(expected_all_tests, all_tests)
 
     def test_find_tests_by_test_mapping_include_subdir(self):
         """Test _find_tests_by_test_mapping method to include sub directory."""
-        tests, all_tests = self.ctr._find_tests_by_test_mapping(
-            path=TEST_MAPPING_TOP_DIR, file_name='test_mapping_sample',
-            include_subdirs=True)
-        expected = set([TEST_1, TEST_2])
+        os_environ_mock = {constants.ANDROID_BUILD_TOP: uc.TEST_DATA_DIR}
+        with mock.patch.dict('os.environ', os_environ_mock, clear=True):
+            tests, all_tests = self.ctr._find_tests_by_test_mapping(
+                path=TEST_MAPPING_TOP_DIR, file_name='test_mapping_sample',
+                include_subdirs=True, checked_files=set())
+        expected = set([TEST_1, TEST_2, TEST_5, TEST_7, TEST_9])
         expected_all_tests = {'presubmit': expected,
-                              'postsubmit': set([TEST_3]),
+                              'postsubmit': set([
+                                  TEST_3, TEST_6, TEST_8, TEST_10]),
                               'other_group': set([TEST_4])}
         self.assertEqual(expected, tests)
         self.assertEqual(expected_all_tests, all_tests)
diff --git a/atest/constants_default.py b/atest/constants_default.py
index f92f0b7..2fb2fa0 100644
--- a/atest/constants_default.py
+++ b/atest/constants_default.py
@@ -72,6 +72,8 @@
 TEST_GROUP_PRESUBMIT = 'presubmit'
 TEST_GROUP_POSTSUBMIT = 'postsubmit'
 TEST_GROUP_ALL = 'all'
+# Key in TEST_MAPPING file for a list of imported TEST_MAPPING file
+TEST_MAPPING_IMPORTS = 'imports'
 
 # TradeFed command line args
 TF_INCLUDE_FILTER = '--include-filter'
diff --git a/atest/test_mapping.py b/atest/test_mapping.py
index 2a9391f..ef442aa 100644
--- a/atest/test_mapping.py
+++ b/atest/test_mapping.py
@@ -18,6 +18,9 @@
 
 
 import copy
+import os
+
+import constants
 
 
 class TestDetail(object):
@@ -64,3 +67,43 @@
 
     def __eq__(self, other):
         return str(self) == str(other)
+
+
+class Import(object):
+    """Store test mapping import details."""
+
+    def __init__(self, test_mapping_file, details):
+        """Import constructor
+
+        Parse import details from a dictionary, e.g.,
+        {
+            "path": "..\folder1"
+        }
+        in which, project is the name of the project, by default it's the
+        current project of the containing TEST_MAPPING file.
+
+        Args:
+            test_mapping_file: Path to the TEST_MAPPING file that contains the
+                import.
+            details: A dictionary of details about importing another
+                TEST_MAPPING file.
+        """
+        self.test_mapping_file = test_mapping_file
+        self.path = details['path']
+
+    def __str__(self):
+        """String value of the Import object."""
+        return 'Source: %s, path: %s' % (self.test_mapping_file, self.path)
+
+    def get_path(self):
+        """Get the path to TEST_MAPPING import directory."""
+        path = os.path.realpath(os.path.join(
+            os.path.dirname(self.test_mapping_file), self.path))
+        if os.path.exists(path):
+            return path
+        root_dir = os.environ.get(constants.ANDROID_BUILD_TOP, os.sep)
+        path = os.path.realpath(os.path.join(root_dir, self.path))
+        if os.path.exists(path):
+            return path
+        # The import path can't be located.
+        return None
diff --git a/atest/unittest_data/test_mapping/folder1/test_mapping_sample b/atest/unittest_data/test_mapping/folder1/test_mapping_sample
index 8dfa4b2..05cea61 100644
--- a/atest/unittest_data/test_mapping/folder1/test_mapping_sample
+++ b/atest/unittest_data/test_mapping/folder1/test_mapping_sample
@@ -13,5 +13,10 @@
     {
       "name": "test4"
     }
+  ],
+  "imports": [
+    {
+      "path": "../folder2"
+    }
   ]
 }
diff --git a/atest/unittest_data/test_mapping/folder2/test_mapping_sample b/atest/unittest_data/test_mapping/folder2/test_mapping_sample
new file mode 100644
index 0000000..7517cd5
--- /dev/null
+++ b/atest/unittest_data/test_mapping/folder2/test_mapping_sample
@@ -0,0 +1,23 @@
+{
+  "presubmit": [
+    {
+      "name": "test5"
+    }
+  ],
+  "postsubmit": [
+    {
+      "name": "test6"
+    }
+  ],
+  "imports": [
+    {
+      "path": "../folder1"
+    },
+    {
+      "path": "../folder3/folder4"
+    },
+    {
+      "path": "../folder3/non-existing"
+    }
+  ]
+}
diff --git a/atest/unittest_data/test_mapping/folder3/folder4/test_mapping_sample b/atest/unittest_data/test_mapping/folder3/folder4/test_mapping_sample
new file mode 100644
index 0000000..6310055
--- /dev/null
+++ b/atest/unittest_data/test_mapping/folder3/folder4/test_mapping_sample
@@ -0,0 +1,7 @@
+{
+  "imports": [
+    {
+      "path": "../../folder5"
+    }
+  ]
+}
diff --git a/atest/unittest_data/test_mapping/folder3/test_mapping_sample b/atest/unittest_data/test_mapping/folder3/test_mapping_sample
new file mode 100644
index 0000000..ecd5b7d
--- /dev/null
+++ b/atest/unittest_data/test_mapping/folder3/test_mapping_sample
@@ -0,0 +1,17 @@
+{
+  "presubmit": [
+    {
+      "name": "test7"
+    }
+  ],
+  "postsubmit": [
+    {
+      "name": "test8"
+    }
+  ],
+  "imports": [
+    {
+      "path": "../folder1"
+    }
+  ]
+}
diff --git a/atest/unittest_data/test_mapping/folder5/test_mapping_sample b/atest/unittest_data/test_mapping/folder5/test_mapping_sample
new file mode 100644
index 0000000..c449a0a
--- /dev/null
+++ b/atest/unittest_data/test_mapping/folder5/test_mapping_sample
@@ -0,0 +1,12 @@
+{
+  "presubmit": [
+    {
+      "name": "test9"
+    }
+  ],
+  "postsubmit": [
+    {
+      "name": "test10"
+    }
+  ]
+}