Simplify main function for mainline_modules_sdks.py

Bug: 205100626
Test: presubmit
Change-Id: I2fdd371aad6a9c37703a190754604da2d86975b2
Merged-In: I2fdd371aad6a9c37703a190754604da2d86975b2
(cherry picked from commit a584ab5dc0)

build/mainline_modules_sdks.py

@@ -317,7 +317,8 @@ class SdkDistProducer:
         return os.path.join(self.out_dir, "soong/mainline-sdks",
                             f"{sdk_name}-{sdk_version}.zip")
 
-    def produce_dist(self, sdk_versions, modules):
+    def produce_dist(self, modules):
+        sdk_versions = SDK_VERSIONS
         self.build_sdks(sdk_versions, modules)
         self.populate_dist(sdk_versions, modules)
 
@@ -503,26 +504,30 @@ def apply_transformations(producer, tmp_dir, transformations):
         os.utime(path, (modified, modified))
 
 
-def main():
-    """Program entry point."""
-    if not os.path.exists("build/make/core/Makefile"):
-        sys.exit("This script must be run from the top of the tree.")
-
-    producer = SdkDistProducer(
+def create_producer():
+    return SdkDistProducer(
         # Variables initialized from environment variables that are set by the
         # calling mainline_modules_sdks.sh.
         out_dir=os.environ["OUT_DIR"],
         dist_dir=os.environ["DIST_DIR"],
     )
 
+def filter_modules(modules):
     target_build_apps = os.environ.get("TARGET_BUILD_APPS")
     if target_build_apps:
-        build_mainline_modules = [m for m in MAINLINE_MODULES
-                                  if m.apex in target_build_apps.split()]
+        target_build_apps = target_build_apps.split()
+        return [m for m in modules if m.apex in target_build_apps]
     else:
-        build_mainline_modules = MAINLINE_MODULES
+        return modules
 
-    producer.produce_dist(SDK_VERSIONS, build_mainline_modules)
+def main():
+    """Program entry point."""
+    if not os.path.exists("build/make/core/Makefile"):
+        sys.exit("This script must be run from the top of the tree.")
+
+    producer = create_producer()
+    modules = filter_modules(MAINLINE_MODULES)
+    producer.produce_dist(modules)
 
 
 if __name__ == "__main__":

build/mainline_modules_sdks_test.py

@@ -162,5 +162,18 @@ class TestSoongConfigBoilerplateInserter(unittest.TestCase):
         self.apply_transformations(src, transformations, expected)
 
 
+class TestFilterModules(unittest.TestCase):
+
+    def test_no_filter(self):
+        modules = mm.filter_modules(mm.MAINLINE_MODULES)
+        self.assertEqual(modules, mm.MAINLINE_MODULES)
+
+    def test_with_filter(self):
+        os.environ["TARGET_BUILD_APPS"] = "com.android.art"
+        modules = mm.filter_modules(mm.MAINLINE_MODULES)
+        expected = mm.MAINLINE_MODULES_BY_APEX["com.android.art"]
+        self.assertEqual(modules, [expected])
+
+
 if __name__ == "__main__":
     unittest.main(verbosity=2)