Use the standard lib to read proto paths

Bug: 365157972
Test: atest FlagJUnitHostTests
Change-Id: I8bff4c94b57b584cf4606c06da714dd2cc95af3d
diff --git a/libraries/flag-helpers/junit/Android.bp b/libraries/flag-helpers/junit/Android.bp
index a5c837c..f842b68 100644
--- a/libraries/flag-helpers/junit/Android.bp
+++ b/libraries/flag-helpers/junit/Android.bp
@@ -59,6 +59,7 @@
     name: "flag-junit-host",
     srcs: ["src_host/**/*.java"],
     static_libs: [
+        "aconfig_host_device_paths_java",
         "flag-junit-base",
         "flag-util",
         "libaconfig_java_proto_full",
diff --git a/libraries/flag-helpers/junit/src_host/android/platform/test/flag/junit/host/DeviceFlags.java b/libraries/flag-helpers/junit/src_host/android/platform/test/flag/junit/host/DeviceFlags.java
index 7149a22..b54ec57 100644
--- a/libraries/flag-helpers/junit/src_host/android/platform/test/flag/junit/host/DeviceFlags.java
+++ b/libraries/flag-helpers/junit/src_host/android/platform/test/flag/junit/host/DeviceFlags.java
@@ -19,6 +19,7 @@
 import android.aconfig.Aconfig.flag_permission;
 import android.aconfig.Aconfig.flag_state;
 import android.aconfig.Aconfig.parsed_flags;
+import android.aconfig.HostDeviceProtos;
 import android.platform.test.flag.util.Flag;
 import android.platform.test.flag.util.FlagReadException;
 
@@ -54,13 +55,6 @@
     private static final String DUMP_DEVICE_CONFIG_CMD = "device_config list";
 
     /**
-     * Partitions that contain the aconfig_flags.pb files. Should be consistent with the partitions
-     * defined in core/packaging/flags.mk.
-     */
-    private static final List<String> FLAG_PARTITIONS =
-            List.of("product", "system", "system_ext", "vendor");
-
-    /**
      * The key is the flag name with namespace ({namespace}/{flagName} for legacy flags,
      * {namespace}/{packageName}.{flagName} for aconfig flags.
      */
@@ -146,10 +140,21 @@
     private parsed_flags getAconfigParsedFlags(ITestDevice testDevice) throws FlagReadException {
         parsed_flags.Builder builder = parsed_flags.newBuilder();
 
-        for (String flagPartition : FLAG_PARTITIONS) {
+        List<String> protoPaths =
+                HostDeviceProtos.parsedFlagsProtoPaths(
+                        command -> {
+                            try {
+                                String adbResult = testDevice.executeAdbCommand(command.split(" "));
+                                LogUtil.CLog.i(
+                                        "Adb command result for '%s': %s", command, adbResult);
+                                return adbResult;
+                            } catch (DeviceNotAvailableException e) {
+                                throw new FlagReadException("ALL_FLAGS", e);
+                            }
+                        });
+
+        for (String aconfigFlagsPbFilePath : protoPaths) {
             try {
-                String aconfigFlagsPbFilePath =
-                        String.format("/%s/etc/aconfig_flags.pb", flagPartition);
                 if (!testDevice.doesFileExist(aconfigFlagsPbFilePath)) {
                     LogUtil.CLog.i("Aconfig flags file %s does not exist", aconfigFlagsPbFilePath);
                     continue;