diff --git a/ci/Android.bp b/ci/Android.bp
index 22c4851..6d4ac35 100644
--- a/ci/Android.bp
+++ b/ci/Android.bp
@@ -71,11 +71,36 @@
     },
 }
 
+python_test_host {
+    name: "optimized_targets_test",
+    main: "optimized_targets_test.py",
+    pkg_path: "testdata",
+    srcs: [
+        "optimized_targets_test.py",
+    ],
+    libs: [
+        "build_test_suites",
+        "pyfakefs",
+    ],
+    test_options: {
+        unit_test: true,
+    },
+    data: [
+        ":py3-cmd",
+    ],
+    version: {
+        py3: {
+            embedded_launcher: true,
+        },
+    },
+}
+
 python_library_host {
     name: "build_test_suites",
     srcs: [
         "build_test_suites.py",
         "optimized_targets.py",
+        "test_mapping_module_retriever.py",
         "build_context.py",
     ],
 }
diff --git a/ci/optimized_targets.py b/ci/optimized_targets.py
index 116d6f8..fddde17 100644
--- a/ci/optimized_targets.py
+++ b/ci/optimized_targets.py
@@ -16,8 +16,13 @@
 from abc import ABC
 import argparse
 import functools
-from typing import Self
 from build_context import BuildContext
+import json
+import logging
+import os
+from typing import Self
+
+import test_mapping_module_retriever
 
 
 class OptimizedBuildTarget(ABC):
@@ -41,7 +46,10 @@
   def get_build_targets(self) -> set[str]:
     features = self.build_context.enabled_build_features
     if self.get_enabled_flag() in features:
-      return self.get_build_targets_impl()
+      self.modules_to_build = self.get_build_targets_impl()
+      return self.modules_to_build
+
+    self.modules_to_build = {self.target}
     return {self.target}
 
   def package_outputs(self):
@@ -82,6 +90,30 @@
     pass
 
 
+class ChangeInfo:
+
+  def __init__(self, change_info_file_path):
+    try:
+      with open(change_info_file_path) as change_info_file:
+        change_info_contents = json.load(change_info_file)
+    except json.decoder.JSONDecodeError:
+      logging.error(f'Failed to load CHANGE_INFO: {change_info_file_path}')
+      raise
+
+    self._change_info_contents = change_info_contents
+
+  def find_changed_files(self) -> set[str]:
+    changed_files = 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'):
+          changed_files.add(project_path + file_info.get('path'))
+
+    return changed_files
+
 class GeneralTestsOptimizer(OptimizedBuildTarget):
   """general-tests optimizer
 
@@ -94,8 +126,55 @@
   normally built.
   """
 
+  # List of modules that are always required to be in general-tests.zip.
+  _REQUIRED_MODULES = frozenset(
+      ['cts-tradefed', 'vts-tradefed', 'compatibility-host-util']
+  )
+
+  def get_build_targets_impl(self) -> set[str]:
+    change_info_file_path = os.environ.get('CHANGE_INFO')
+    if not change_info_file_path:
+      logging.info(
+          'No CHANGE_INFO env var found, general-tests optimization disabled.'
+      )
+      return {'general-tests'}
+
+    test_infos = self.build_context.test_infos
+    test_mapping_test_groups = set()
+    for test_info in test_infos:
+      is_test_mapping = test_info.is_test_mapping
+      current_test_mapping_test_groups = test_info.test_mapping_test_groups
+      uses_general_tests = test_info.build_target_used('general-tests')
+
+      if uses_general_tests and not is_test_mapping:
+        logging.info(
+            'Test uses general-tests.zip but is not test-mapping, general-tests'
+            ' optimization disabled.'
+        )
+        return {'general-tests'}
+
+      if is_test_mapping:
+        test_mapping_test_groups.update(current_test_mapping_test_groups)
+
+    change_info = ChangeInfo(change_info_file_path)
+    changed_files = change_info.find_changed_files()
+
+    test_mappings = test_mapping_module_retriever.GetTestMappings(
+        changed_files, set()
+    )
+
+    modules_to_build = set(self._REQUIRED_MODULES)
+
+    modules_to_build.update(
+        test_mapping_module_retriever.FindAffectedModules(
+            test_mappings, changed_files, test_mapping_test_groups
+        )
+    )
+
+    return modules_to_build
+
   def get_enabled_flag(self):
-    return 'general-tests-optimized'
+    return 'general_tests_optimized'
 
   @classmethod
   def get_optimized_targets(cls) -> dict[str, OptimizedBuildTarget]:
diff --git a/ci/optimized_targets_test.py b/ci/optimized_targets_test.py
new file mode 100644
index 0000000..919c193
--- /dev/null
+++ b/ci/optimized_targets_test.py
@@ -0,0 +1,206 @@
+# Copyright 2024, 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.
+
+"""Tests for optimized_targets.py"""
+
+import json
+import logging
+import os
+import pathlib
+import re
+import unittest
+from unittest import mock
+import optimized_targets
+from build_context import BuildContext
+from pyfakefs import fake_filesystem_unittest
+
+
+class GeneralTestsOptimizerTest(fake_filesystem_unittest.TestCase):
+
+  def setUp(self):
+    self.setUpPyfakefs()
+
+    os_environ_patcher = mock.patch.dict('os.environ', {})
+    self.addCleanup(os_environ_patcher.stop)
+    self.mock_os_environ = os_environ_patcher.start()
+
+    self._setup_working_build_env()
+    self._write_change_info_file()
+    test_mapping_dir = pathlib.Path('/project/path/file/path')
+    test_mapping_dir.mkdir(parents=True)
+    self._write_test_mapping_file()
+
+  def _setup_working_build_env(self):
+    self.change_info_file = pathlib.Path('/tmp/change_info')
+
+    self.mock_os_environ.update({
+        'CHANGE_INFO': str(self.change_info_file),
+    })
+
+  def test_general_tests_optimized(self):
+    optimizer = self._create_general_tests_optimizer()
+
+    build_targets = optimizer.get_build_targets()
+
+    expected_build_targets = set(
+        optimized_targets.GeneralTestsOptimizer._REQUIRED_MODULES
+    )
+    expected_build_targets.add('test_mapping_module')
+
+    self.assertSetEqual(build_targets, expected_build_targets)
+
+  def test_no_change_info_no_optimization(self):
+    del os.environ['CHANGE_INFO']
+
+    optimizer = self._create_general_tests_optimizer()
+
+    build_targets = optimizer.get_build_targets()
+
+    self.assertSetEqual(build_targets, {'general-tests'})
+
+  def test_mapping_groups_unused_module_not_built(self):
+    test_context = self._create_test_context()
+    test_context['testInfos'][0]['extraOptions'] = [
+        {
+            'key': 'additional-files-filter',
+            'values': ['general-tests.zip'],
+        },
+        {
+            'key': 'test-mapping-test-group',
+            'values': ['unused-test-mapping-group'],
+        },
+    ]
+    optimizer = self._create_general_tests_optimizer(
+        build_context=self._create_build_context(test_context=test_context)
+    )
+
+    build_targets = optimizer.get_build_targets()
+
+    expected_build_targets = set(
+        optimized_targets.GeneralTestsOptimizer._REQUIRED_MODULES
+    )
+    self.assertSetEqual(build_targets, expected_build_targets)
+
+  def test_general_tests_used_by_non_test_mapping_test_no_optimization(self):
+    test_context = self._create_test_context()
+    test_context['testInfos'][0]['extraOptions'] = [{
+        'key': 'additional-files-filter',
+        'values': ['general-tests.zip'],
+    }]
+    optimizer = self._create_general_tests_optimizer(
+        build_context=self._create_build_context(test_context=test_context)
+    )
+
+    build_targets = optimizer.get_build_targets()
+
+    self.assertSetEqual(build_targets, {'general-tests'})
+
+  def test_malformed_change_info_raises(self):
+    with open(self.change_info_file, 'w') as f:
+      f.write('not change info')
+
+    optimizer = self._create_general_tests_optimizer()
+
+    with self.assertRaises(json.decoder.JSONDecodeError):
+      build_targets = optimizer.get_build_targets()
+
+  def test_malformed_test_mapping_raises(self):
+    with open('/project/path/file/path/TEST_MAPPING', 'w') as f:
+      f.write('not test mapping')
+
+    optimizer = self._create_general_tests_optimizer()
+
+    with self.assertRaises(json.decoder.JSONDecodeError):
+      build_targets = optimizer.get_build_targets()
+
+  def _write_change_info_file(self):
+    change_info_contents = {
+        'changes': [{
+            'projectPath': '/project/path',
+            'revisions': [{
+                'fileInfos': [{
+                    'path': 'file/path/file_name',
+                }],
+            }],
+        }]
+    }
+
+    with open(self.change_info_file, 'w') as f:
+      json.dump(change_info_contents, f)
+
+  def _write_test_mapping_file(self):
+    test_mapping_contents = {
+        'test-mapping-group': [
+            {
+                'name': 'test_mapping_module',
+            },
+        ],
+    }
+
+    with open('/project/path/file/path/TEST_MAPPING', 'w') as f:
+      json.dump(test_mapping_contents, f)
+
+  def _create_general_tests_optimizer(
+      self, build_context: BuildContext = None
+  ):
+    if not build_context:
+      build_context = self._create_build_context()
+    return optimized_targets.GeneralTestsOptimizer(
+        'general-tests', build_context, None
+    )
+
+  def _create_build_context(
+      self,
+      general_tests_optimized: bool = True,
+      test_context: dict[str, any] = None,
+  ) -> BuildContext:
+    if not test_context:
+      test_context = self._create_test_context()
+    build_context_dict = {}
+    build_context_dict['enabledBuildFeatures'] = [{'name': 'optimized_build'}]
+    if general_tests_optimized:
+      build_context_dict['enabledBuildFeatures'].append({'name': 'general_tests_optimized'})
+    build_context_dict['testContext'] = test_context
+    return BuildContext(build_context_dict)
+
+  def _create_test_context(self):
+    return {
+        'testInfos': [
+            {
+                'name': 'atp_test',
+                'target': 'test_target',
+                'branch': 'branch',
+                'extraOptions': [
+                    {
+                        'key': 'additional-files-filter',
+                        'values': ['general-tests.zip'],
+                    },
+                    {
+                        'key': 'test-mapping-test-group',
+                        'values': ['test-mapping-group'],
+                    },
+                ],
+                'command': '/tf/command',
+                'extraBuildTargets': [
+                    'extra_build_target',
+                ],
+            },
+        ],
+    }
+
+
+if __name__ == '__main__':
+  # Setup logging to be silent so unit tests can pass through TF.
+  logging.disable(logging.ERROR)
+  unittest.main()
diff --git a/ci/test_mapping_module_retriever.py b/ci/test_mapping_module_retriever.py
index d2c13c0..c93cdd5 100644
--- a/ci/test_mapping_module_retriever.py
+++ b/ci/test_mapping_module_retriever.py
@@ -17,11 +17,13 @@
 modules are needed to build for the given list of changed files.
 TODO(lucafarsi): Deduplicate from artifact_helper.py
 """
+# TODO(lucafarsi): Share this logic with the original logic in
+# test_mapping_test_retriever.py
 
-from typing import Any, Dict, Set, Text
 import json
 import os
 import re
+from typing import Any
 
 # Regex to extra test name from the path of test config file.
 TEST_NAME_REGEX = r'(?:^|.*/)([^/]+)\.config'
@@ -39,7 +41,7 @@
 _COMMENTS_RE = re.compile(r'(\"(?:[^\"\\]|\\.)*\"|(?=//))(?://.*)?')
 
 
-def FilterComments(test_mapping_file: Text) -> Text:
+def FilterComments(test_mapping_file: str) -> str:
   """Remove comments in TEST_MAPPING file to valid format.
 
   Only '//' is regarded as comments.
@@ -52,8 +54,8 @@
   """
   return re.sub(_COMMENTS_RE, r'\1', test_mapping_file)
 
-def GetTestMappings(paths: Set[Text],
-                    checked_paths: Set[Text]) -> Dict[Text, Dict[Text, Any]]:
+def GetTestMappings(paths: set[str],
+                    checked_paths: set[str]) -> dict[str, dict[str, Any]]:
   """Get the affected TEST_MAPPING files.
 
   TEST_MAPPING files in source code are packaged into a build artifact
@@ -123,3 +125,68 @@
       pass
 
   return test_mappings
+
+
+def FindAffectedModules(
+    test_mappings: dict[str, Any],
+    changed_files: set[str],
+    test_mapping_test_groups: set[str],
+) -> set[str]:
+  """Find affected test modules.
+
+  Find the affected set of test modules that would run in a test mapping run based on the given test mappings, changed files, and test mapping test group.
+
+  Args:
+    test_mappings: A set of test mappings returned by GetTestMappings in the following format:
+      {
+        'test_mapping_file_path': {
+          'group_name' : [
+            'name': 'module_name',
+          ],
+        }
+      }
+    changed_files: A set of files changed for the given run.
+    test_mapping_test_groups: A set of test mapping test groups that are being considered for the given run.
+
+  Returns:
+    A set of test module names which would run for a test mapping test run with the given parameters.
+  """
+
+  modules = set()
+
+  for test_mapping in test_mappings.values():
+    for group_name, group in test_mapping.items():
+      # If a module is not in any of the test mapping groups being tested skip
+      # it.
+      if group_name not in test_mapping_test_groups:
+        continue
+
+      for entry in group:
+        module_name = entry.get('name')
+
+        if not module_name:
+          continue
+
+        file_patterns = entry.get('file_patterns')
+        if not file_patterns:
+          modules.add(module_name)
+          continue
+
+        if matches_file_patterns(file_patterns, changed_files):
+          modules.add(module_name)
+
+  return modules
+
+def MatchesFilePatterns(
+    file_patterns: list[set], changed_files: set[str]
+) -> bool:
+  """Checks if any of the changed files match any of the file patterns.
+
+  Args:
+    file_patterns: A list of file patterns to match against.
+    changed_files: A set of files to check against the file patterns.
+
+  Returns:
+    True if any of the changed files match any of the file patterns.
+  """
+  return any(re.search(pattern, "|".join(changed_files)) for pattern in file_patterns)