kleaf: Test to ensure that the correct modules.load lists are generated

Add a test to ensure that only the expected *modules.load* lists are
generated and no extraneous lists are generated.

Bug: 324304750
Bug: 324356826
Signed-off-by: Isaac J. Manjarres <isaacmanjarres@google.com>
(cherry picked from https://android-review.googlesource.com/q/commit:580b1a61b6a4a4f50c25f3cf3349f49099b1d9cc)
Merged-In: I7aab46022a3fcb5f754d6e517a194d0f553f21e0
Change-Id: I7aab46022a3fcb5f754d6e517a194d0f553f21e0

2 files changed, 50 insertions, 0 deletions

diff --git a/kleaf/artifact_tests/initramfs_modules_lists_test.py b/kleaf/artifact_tests/initramfs_modules_lists_test.py
index f1d655f7..19f14c78 100644
--- a/kleaf/artifact_tests/initramfs_modules_lists_test.py
+++ b/kleaf/artifact_tests/initramfs_modules_lists_test.py
@@ -18,6 +18,7 @@ import subprocess
 import sys
 import tempfile
 import unittest
+import re
 
 from absl.testing import absltest
 
@@ -27,6 +28,8 @@ def load_arguments():
     parser.add_argument("--expected_modules_list")
     parser.add_argument("--expected_modules_recovery_list")
     parser.add_argument("--expected_modules_charger_list")
+    parser.add_argument("--build_vendor_boot")
+    parser.add_argument("--build_vendor_kernel_boot")
     parser.add_argument("files", nargs="*", default=[])
     return parser.parse_known_args()
 
@@ -127,6 +130,44 @@ class InitramfsModulesLists(unittest.TestCase):
                 modules_dir = os.path.join(lib_modules, v)
                 self._diff_modules_lists(modules_lists_map, modules_dir)
 
+    def _verify_modules_load_lists(self, modules_list_name, vendor_boot_name, image_files):
+        """Given a modules.load* name, ensures that no extraneous such files exist.
+
+        This tests to ensure that if a modules.load* file exists, then no other
+        modules.load* file must exist, except for vendor_boot.modules.load* or
+        vendor_kernel_boot.modules.load.
+
+        Args:
+            modules_list_name: The name of the modules.load list.
+            vendor_boot_name: Either vendor_boot or vendor_kernel_boot.
+            image_files: The files associated with the kernel_images() target.
+        """
+        modules_load_lists = [modules_list_name]
+        if vendor_boot_name:
+            modules_load_lists.append(f"{vendor_boot_name}.{modules_list_name}")
+
+        modules_load_list_re = re.compile(f".*{modules_list_name}$")
+        modules_load_list_matches = [os.path.basename(f) for f in image_files if modules_load_list_re.fullmatch(os.path.basename(f))]
+
+        self.assertCountEqual(modules_load_lists, modules_load_list_matches)
+
+    def test_modules_lists_existence(self):
+        vendor_boot_name = None
+
+        if arguments.build_vendor_boot:
+            vendor_boot_name = "vendor_boot"
+        elif arguments.build_vendor_kernel_boot:
+            vendor_boot_name = "vendor_kernel_boot"
+
+        if arguments.expected_modules_list:
+            self._verify_modules_load_lists("modules.load", vendor_boot_name, arguments.files)
+
+        if arguments.expected_modules_recovery_list:
+            self._verify_modules_load_lists("modules.load.recovery", vendor_boot_name, arguments.files)
+
+        if arguments.expected_modules_charger_list:
+            self._verify_modules_load_lists("modules.load.charger", vendor_boot_name, arguments.files)
+
 if __name__ == '__main__':
     arguments, unknown = load_arguments()
     sys.argv[1:] = unknown
diff --git a/kleaf/artifact_tests/kernel_test.bzl b/kleaf/artifact_tests/kernel_test.bzl
index 2631ab61..fe919986 100644
--- a/kleaf/artifact_tests/kernel_test.bzl
+++ b/kleaf/artifact_tests/kernel_test.bzl
@@ -124,6 +124,8 @@ def initramfs_modules_lists_test(
         expected_modules_list = None,
         expected_modules_recovery_list = None,
         expected_modules_charger_list = None,
+        build_vendor_boot = None,
+        build_vendor_kernel_boot = None,
         **kwargs):
     """Tests that the initramfs has modules.load* files with the given content.
 
@@ -133,6 +135,8 @@ def initramfs_modules_lists_test(
         expected_modules_list: file with the expected content for `modules.load`
         expected_modules_recovery_list: file with the expected content for `modules.load.recovery`
         expected_modules_charger_list: file with the expected content for `modules.load.charger`
+        build_vendor_boot: If the `kernel_images` target builds vendor_boot.img
+        build_vendor_kernel_boot: If the `kernel_images` target builds vendor_kernel_boot.img
         **kwargs: Additional attributes to the internal rule, e.g.
           [`visibility`](https://docs.bazel.build/versions/main/visibility.html).
           See complete list
@@ -159,6 +163,11 @@ def initramfs_modules_lists_test(
             "$(rootpath {})".format(expected_modules_charger_list),
         ]
 
+    if build_vendor_boot:
+        args.append("--build_vendor_boot")
+    elif build_vendor_kernel_boot:
+        args.append("--build_vendor_kernel_boot")
+
     args.append("$(rootpaths {})".format(kernel_images))
 
     hermetic_exec_test(