From 23fbd1e2f6e1b5d7a8d2de272e8264bcd77ee7de Mon Sep 17 00:00:00 2001
From: Jeff DeCew <jeffdq@google.com>
Date: Sat, 27 Apr 2024 18:06:24 +0000
Subject: [PATCH] FakeFeatureFlagsImpl optionally takes defaults via a
 FeatureFlags.

Bug: 337449122
Flag: test_only
Test: presubmit
Change-Id: I45e2a523b36a6b14627c89ee7deffda711a32f5e
---
 tools/aconfig/aconfig/src/codegen/java.rs       | 17 +++++++++++++----
 .../FakeFeatureFlagsImpl.java.template          | 17 +++++++++++++----
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/tools/aconfig/aconfig/src/codegen/java.rs b/tools/aconfig/aconfig/src/codegen/java.rs
index 9abc892908..3360ddd68b 100644
--- a/tools/aconfig/aconfig/src/codegen/java.rs
+++ b/tools/aconfig/aconfig/src/codegen/java.rs
@@ -428,10 +428,16 @@ mod tests {
 
     /** @hide */
     public class FakeFeatureFlagsImpl extends CustomFeatureFlags {
-        private Map<String, Boolean> mFlagMap = new HashMap<>();
+        private final Map<String, Boolean> mFlagMap = new HashMap<>();
+        private final FeatureFlags mDefaults;
 
         public FakeFeatureFlagsImpl() {
+            this(null);
+        }
+
+        public FakeFeatureFlagsImpl(FeatureFlags defaults) {
             super(null);
+            mDefaults = defaults;
             // Initialize the map with null values
             for (String flagName : getFlagNames()) {
                 mFlagMap.put(flagName, null);
@@ -441,10 +447,13 @@ mod tests {
         @Override
         protected boolean getValue(String flagName, Predicate<FeatureFlags> getter) {
             Boolean value = this.mFlagMap.get(flagName);
-            if (value == null) {
-                throw new IllegalArgumentException(flagName + " is not set");
+            if (value != null) {
+                return value;
             }
-            return value;
+            if (mDefaults != null) {
+                return getter.test(mDefaults);
+            }
+            throw new IllegalArgumentException(flagName + " is not set");
         }
 
         public void setFlag(String flagName, boolean value) {
diff --git a/tools/aconfig/aconfig/templates/FakeFeatureFlagsImpl.java.template b/tools/aconfig/aconfig/templates/FakeFeatureFlagsImpl.java.template
index c20d3c5061..290d2c4b24 100644
--- a/tools/aconfig/aconfig/templates/FakeFeatureFlagsImpl.java.template
+++ b/tools/aconfig/aconfig/templates/FakeFeatureFlagsImpl.java.template
@@ -6,10 +6,16 @@ import java.util.function.Predicate;
 
 /** @hide */
 public class FakeFeatureFlagsImpl extends CustomFeatureFlags \{
-    private Map<String, Boolean> mFlagMap = new HashMap<>();
+    private final Map<String, Boolean> mFlagMap = new HashMap<>();
+    private final FeatureFlags mDefaults;
 
     public FakeFeatureFlagsImpl() \{
+        this(null);
+    }
+
+    public FakeFeatureFlagsImpl(FeatureFlags defaults) \{
         super(null);
+        mDefaults = defaults;
         // Initialize the map with null values
         for (String flagName : getFlagNames()) \{
             mFlagMap.put(flagName, null);
@@ -19,10 +25,13 @@ public class FakeFeatureFlagsImpl extends CustomFeatureFlags \{
     @Override
     protected boolean getValue(String flagName, Predicate<FeatureFlags> getter) \{
         Boolean value = this.mFlagMap.get(flagName);
-        if (value == null) \{
-            throw new IllegalArgumentException(flagName + " is not set");
+        if (value != null) \{
+            return value;
         }
-        return value;
+        if (mDefaults != null) \{
+            return getter.test(mDefaults);
+        }
+        throw new IllegalArgumentException(flagName + " is not set");
     }
 
     public void setFlag(String flagName, boolean value) \{