For test mapping config append modified path

Test: presubmit
Bug: 405154519
Change-Id: I4a2bc3806722be9322f0d40070d00c3fafc9faa3

3 files changed, 66 insertions, 5 deletions

diff --git a/ci/optimized_targets.py b/ci/optimized_targets.py
index 0e9723cad8..548e34273d 100644
--- a/ci/optimized_targets.py
+++ b/ci/optimized_targets.py
@@ -207,6 +207,19 @@ class ChangeInfo:
 
     self._change_info_contents = change_info_contents
 
+  def get_changed_paths(self) -> set[str]:
+    changed_paths = set()
+    for change in self._change_info_contents['changes']:
+      project_path = change.get('projectPath') + '/'
+
+      for revision in change.get('revisions'):
+        for file_info in revision.get('fileInfos'):
+          file_path = file_info.get('path')
+          dir_path = os.path.dirname(file_path)
+          changed_paths.add(project_path + dir_path)
+
+    return changed_paths
+
   def find_changed_files(self) -> set[str]:
     changed_files = set()
 
@@ -266,9 +279,11 @@ class GeneralTestsOptimizer(OptimizedBuildTarget):
 
 
   def _get_test_discovery_modules(self) -> set[str]:
+    change_info = ChangeInfo(os.environ.get('CHANGE_INFO'))
+    change_paths = change_info.get_changed_paths()
     test_modules = set()
     for test_info in self.test_infos:
-      tf_command = self._build_tf_command(test_info)
+      tf_command = self._build_tf_command(test_info, change_paths)
       discovery_agent = test_discovery_agent.TestDiscoveryAgent(tradefed_args=tf_command, test_mapping_zip_path=os.environ.get('DIST_DIR')+'/test_mappings.zip')
       modules, dependencies = discovery_agent.discover_test_mapping_test_modules()
       for regex in modules:
@@ -276,7 +291,7 @@ class GeneralTestsOptimizer(OptimizedBuildTarget):
     return test_modules
 
 
-  def _build_tf_command(self, test_info) -> list[str]:
+  def _build_tf_command(self, test_info, change_paths) -> list[str]:
     command = [test_info.command]
     for extra_option in test_info.extra_options:
       if not extra_option.get('key'):
@@ -292,6 +307,10 @@ class GeneralTestsOptimizer(OptimizedBuildTarget):
           command.append(value)
       else:
         command.append(arg_key)
+    if test_info.is_test_mapping:
+      for change_path in change_paths:
+        command.append('--test-mapping-path')
+        command.append(change_path)
 
     return command
 
diff --git a/ci/optimized_targets_test.py b/ci/optimized_targets_test.py
index 8256acd1b4..2935c83cc5 100644
--- a/ci/optimized_targets_test.py
+++ b/ci/optimized_targets_test.py
@@ -44,6 +44,7 @@ class GeneralTestsOptimizerTest(fake_filesystem_unittest.TestCase):
 
   def _setup_working_build_env(self):
     self._write_soong_ui_file()
+    self._write_change_info_file()
     self._host_out_testcases = pathlib.Path('/tmp/top/host_out_testcases')
     self._host_out_testcases.mkdir(parents=True)
     self._target_out_testcases = pathlib.Path('/tmp/top/target_out_testcases')
@@ -62,9 +63,47 @@ class GeneralTestsOptimizerTest(fake_filesystem_unittest.TestCase):
     self.mock_os_environ.update({
         'TOP': '/tmp/top',
         'DIST_DIR': '/tmp/top/out/dist',
-        'TMPDIR': '/tmp/'
+        'TMPDIR': '/tmp/',
+        'CHANGE_INFO': '/tmp/top/change_info'
     })
 
+  def _write_change_info_file(self):
+    change_info_path = pathlib.Path('/tmp/top/')
+    with open(os.path.join(change_info_path, 'change_info'), 'w') as f:
+      f.write("""
+    {
+      "changes": [
+        {
+          "projectPath": "build/ci",
+          "revisions": [
+            {
+              "revisionNumber": 1,
+              "fileInfos": [
+                {
+                  "path": "src/main/java/com/example/MyClass.java",
+                  "action": "MODIFIED"
+                },
+                {
+                  "path": "src/test/java/com/example/MyClassTest.java",
+                  "action": "ADDED"
+                }
+              ]
+            },
+            {
+              "revisionNumber": 2,
+              "fileInfos": [
+                {
+                  "path": "src/main/java/com/example/AnotherClass.java",
+                  "action": "MODIFIED"
+                }
+              ]
+            }
+          ]
+        }
+      ]
+    }
+    """)
+
   def _write_soong_ui_file(self):
     soong_path = pathlib.Path('/tmp/top/build/soong')
     soong_path.mkdir(parents=True)
diff --git a/ci/test_discovery_agent.py b/ci/test_discovery_agent.py
index 6812b071ff..3c1caf45d9 100644
--- a/ci/test_discovery_agent.py
+++ b/ci/test_discovery_agent.py
@@ -118,15 +118,18 @@ class TestDiscoveryAgent:
       )
       java_args.extend(self.tradefed_args)
       env = os.environ.copy()
+      env.update({"SKIP_JAVA_QUERY": "1"})
+      env.update({"ALLOW_EMPTY_TEST_MAPPING": "1"})
       env.update({"TF_TEST_MAPPING_ZIP_FILE": self.test_mapping_zip_path})
       env.update({"DISCOVERY_OUTPUT_FILE": test_discovery_output_file.name})
       logging.info(f"Calling test discovery with args: {java_args}")
       try:
-        result = subprocess.run(args=java_args, env=env, text=True, check=True)
+        result = subprocess.run(args=java_args, env=env, text=True, check=True, stdout = subprocess.PIPE,
+    stderr = subprocess.PIPE)
         logging.info(f"Test discovery agent output: {result.stdout}")
       except subprocess.CalledProcessError as e:
         raise TestDiscoveryError(
-            f"Failed to run test discovery, strout: {e.stdout}, strerr:"
+            f"Failed to run test discovery, stdout: {e.stdout}, stderr:"
             f" {e.stderr}, returncode: {e.returncode}"
         )
       data = json.loads(test_discovery_output_file.read())