Add setFlag and resetAll in FeatureFlags test mode

Add methods setFlag and resetAll in FeatureFlags in test mode. For the
injection usecase, user will use the interface FeatureFlags in the code
to control the flags.

Add tests for test mode.

Bug: 280833463
Test: Atest AconfigJavaHostTest --host
Change-Id: Ib59ba35a9011a6400af42fc9c283d37193577997
diff --git a/tools/aconfig/Android.bp b/tools/aconfig/Android.bp
index a4ea7f4..93fd6f7 100644
--- a/tools/aconfig/Android.bp
+++ b/tools/aconfig/Android.bp
@@ -101,7 +101,7 @@
 android_test {
     name: "aconfig.test.java",
     srcs: [
-        "tests/**/*.java",
+        "tests/AconfigTest.java",
     ],
     manifest: "tests/AndroidManifest.xml",
     certificate: "platform",
@@ -113,6 +113,25 @@
     test_suites: ["device-tests"],
 }
 
+java_aconfig_library {
+    name: "aconfig_host_test_java_library",
+    aconfig_declarations: "aconfig.test.flags",
+    host_supported: true,
+    test: true,
+}
+
+java_test_host {
+    name: "AconfigJavaHostTest",
+    srcs: [
+        "tests/AconfigHostTest.java",
+    ],
+    static_libs: [
+        "aconfig_host_test_java_library",
+        "junit",
+    ],
+    test_suites: ["general-tests"],
+}
+
 // integration tests: C++
 
 cc_aconfig_library {
diff --git a/tools/aconfig/src/codegen_java.rs b/tools/aconfig/src/codegen_java.rs
index a117370..2c9dcf5 100644
--- a/tools/aconfig/src/codegen_java.rs
+++ b/tools/aconfig/src/codegen_java.rs
@@ -116,14 +116,14 @@
     use super::*;
     use std::collections::HashMap;
 
-    const EXPECTED_FEATUREFLAGS_CONTENT: &str = r#"
+    const EXPECTED_FEATUREFLAGS_COMMON_CONTENT: &str = r#"
     package com.android.aconfig.test;
     public interface FeatureFlags {
         boolean disabledRo();
         boolean disabledRw();
         boolean enabledRo();
         boolean enabledRw();
-    }"#;
+    "#;
 
     const EXPECTED_FLAG_COMMON_CONTENT: &str = r#"
     package com.android.aconfig.test;
@@ -179,6 +179,9 @@
             CodegenMode::Production,
         )
         .unwrap();
+        let expect_featureflags_content = EXPECTED_FEATUREFLAGS_COMMON_CONTENT.to_string()
+            + r#"
+        }"#;
         let expect_flags_content = EXPECTED_FLAG_COMMON_CONTENT.to_string()
             + r#"
             private static FeatureFlags FEATURE_FLAGS = new FeatureFlagsImpl();
@@ -224,7 +227,7 @@
         let mut file_set = HashMap::from([
             ("com/android/aconfig/test/Flags.java", expect_flags_content.as_str()),
             ("com/android/aconfig/test/FeatureFlagsImpl.java", expect_featureflagsimpl_content),
-            ("com/android/aconfig/test/FeatureFlags.java", EXPECTED_FEATUREFLAGS_CONTENT),
+            ("com/android/aconfig/test/FeatureFlags.java", expect_featureflags_content.as_str()),
             (
                 "com/android/aconfig/test/FakeFeatureFlagsImpl.java",
                 expect_fakefeatureflagsimpl_content.as_str(),
@@ -258,6 +261,11 @@
             CodegenMode::Test,
         )
         .unwrap();
+        let expect_featureflags_content = EXPECTED_FEATUREFLAGS_COMMON_CONTENT.to_string()
+            + r#"
+            public void setFlag(String flagName, boolean value);
+            public void resetAll();
+        }"#;
         let expect_flags_content = EXPECTED_FLAG_COMMON_CONTENT.to_string()
             + r#"
             public static void setFeatureFlags(FeatureFlags featureFlags) {
@@ -275,6 +283,16 @@
             .to_owned()
             + EXPECTED_METHOD_NOT_IMPL_COMMON_CONTENT
             + r#"
+            @Override
+            public void setFlag(String flagName, boolean value) {
+                throw new UnsupportedOperationException(
+                    "Method is not implemented.");
+            }
+            @Override
+            public void resetAll() {
+                throw new UnsupportedOperationException(
+                    "Method is not implemented.");
+            }
         }
         "#;
         let expect_fakefeatureflagsimpl_content = r#"
@@ -300,12 +318,14 @@
             public boolean enabledRw() {
                 return getFlag(Flags.FLAG_ENABLED_RW);
             }
+            @Override
             public void setFlag(String flagName, boolean value) {
                 if (!this.mFlagMap.containsKey(flagName)) {
                     throw new IllegalArgumentException("no such flag" + flagName);
                 }
                 this.mFlagMap.put(flagName, value);
             }
+            @Override
             public void resetAll() {
                 for (Map.Entry entry : mFlagMap.entrySet()) {
                     entry.setValue(null);
@@ -334,7 +354,7 @@
 
         let mut file_set = HashMap::from([
             ("com/android/aconfig/test/Flags.java", expect_flags_content.as_str()),
-            ("com/android/aconfig/test/FeatureFlags.java", EXPECTED_FEATUREFLAGS_CONTENT),
+            ("com/android/aconfig/test/FeatureFlags.java", expect_featureflags_content.as_str()),
             (
                 "com/android/aconfig/test/FeatureFlagsImpl.java",
                 expect_featureflagsimpl_content.as_str(),
diff --git a/tools/aconfig/templates/FakeFeatureFlagsImpl.java.template b/tools/aconfig/templates/FakeFeatureFlagsImpl.java.template
index db35d28..dba82ef 100644
--- a/tools/aconfig/templates/FakeFeatureFlagsImpl.java.template
+++ b/tools/aconfig/templates/FakeFeatureFlagsImpl.java.template
@@ -13,6 +13,7 @@
         return getFlag(Flags.FLAG_{item.flag_name_constant_suffix});
     }
 {{ endfor}}
+    @Override
     public void setFlag(String flagName, boolean value) \{
         if (!this.mFlagMap.containsKey(flagName)) \{
             throw new IllegalArgumentException("no such flag" + flagName);
@@ -20,6 +21,7 @@
         this.mFlagMap.put(flagName, value);
     }
 
+    @Override
     public void resetAll() \{
         for (Map.Entry entry : mFlagMap.entrySet()) \{
             entry.setValue(null);
diff --git a/tools/aconfig/templates/FeatureFlags.java.template b/tools/aconfig/templates/FeatureFlags.java.template
index e0f201f..c99ccbb 100644
--- a/tools/aconfig/templates/FeatureFlags.java.template
+++ b/tools/aconfig/templates/FeatureFlags.java.template
@@ -4,4 +4,10 @@
 {{ for item in class_elements}}
     boolean {item.method_name}();
 {{ endfor }}
+
+{{ -if is_test_mode }}
+    public void setFlag(String flagName, boolean value);
+
+    public void resetAll();
+{{ -endif }}
 }
diff --git a/tools/aconfig/templates/FeatureFlagsImpl.java.template b/tools/aconfig/templates/FeatureFlagsImpl.java.template
index ba86ce5..7e1eb15 100644
--- a/tools/aconfig/templates/FeatureFlagsImpl.java.template
+++ b/tools/aconfig/templates/FeatureFlagsImpl.java.template
@@ -29,5 +29,16 @@
             "Method is not implemented.");
     }
 {{ endfor- }}
+    @Override
+    public void setFlag(String flagName, boolean value) \{
+        throw new UnsupportedOperationException(
+            "Method is not implemented.");
+    }
+
+    @Override
+    public void resetAll() \{
+        throw new UnsupportedOperationException(
+            "Method is not implemented.");
+    }
 }
 {{ endif }}
diff --git a/tools/aconfig/tests/AconfigHostTest.java b/tools/aconfig/tests/AconfigHostTest.java
new file mode 100644
index 0000000..29a01e3
--- /dev/null
+++ b/tools/aconfig/tests/AconfigHostTest.java
@@ -0,0 +1,88 @@
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+
+import com.android.aconfig.test.FakeFeatureFlagsImpl;
+import com.android.aconfig.test.FeatureFlags;
+import com.android.aconfig.test.FeatureFlagsImpl;
+import com.android.aconfig.test.Flags;
+
+@RunWith(JUnit4.class)
+public final class AconfigHostTest {
+    @Test
+    public void testThrowsExceptionIfFlagNotSet() {
+        assertThrows(NullPointerException.class, () -> Flags.disabledRo());
+        FeatureFlags featureFlags = new FakeFeatureFlagsImpl();
+        assertThrows(IllegalArgumentException.class, () -> featureFlags.disabledRo());
+    }
+
+    @Test
+    public void testSetFlagInFakeFeatureFlagsImpl() {
+        FeatureFlags featureFlags = new FakeFeatureFlagsImpl();
+        featureFlags.setFlag(Flags.FLAG_ENABLED_RW, true);
+        assertTrue(featureFlags.enabledRw());
+        featureFlags.setFlag(Flags.FLAG_ENABLED_RW, false);
+        assertFalse(featureFlags.enabledRw());
+
+        //Set Flags
+        assertThrows(NullPointerException.class, () -> Flags.enabledRw());
+        Flags.setFeatureFlags(featureFlags);
+        featureFlags.setFlag(Flags.FLAG_ENABLED_RW, true);
+        assertTrue(Flags.enabledRw());
+        Flags.unsetFeatureFlags();
+    }
+
+    @Test
+    public void testSetFlagWithRandomName() {
+        FeatureFlags featureFlags = new FakeFeatureFlagsImpl();
+        assertThrows(IllegalArgumentException.class,
+            () -> featureFlags.setFlag("Randome_name", true));
+    }
+
+    @Test
+    public void testResetFlagsInFakeFeatureFlagsImpl() {
+        FeatureFlags featureFlags = new FakeFeatureFlagsImpl();
+        featureFlags.setFlag(Flags.FLAG_ENABLED_RO, true);
+        assertTrue(featureFlags.enabledRo());
+        featureFlags.resetAll();
+        assertThrows(IllegalArgumentException.class, () -> featureFlags.enabledRo());
+
+        // Set value after reset
+        featureFlags.setFlag(Flags.FLAG_ENABLED_RO, false);
+        assertFalse(featureFlags.enabledRo());
+    }
+
+    @Test
+    public void testFlagsSetFeatureFlags() {
+        FeatureFlags featureFlags = new FakeFeatureFlagsImpl();
+        featureFlags.setFlag(Flags.FLAG_ENABLED_RW, true);
+        assertThrows(NullPointerException.class, () -> Flags.enabledRw());
+        Flags.setFeatureFlags(featureFlags);
+        assertTrue(Flags.enabledRw());
+        Flags.unsetFeatureFlags();
+    }
+
+    @Test
+    public void testFlagsUnsetFeatureFlags() {
+        FeatureFlags featureFlags = new FakeFeatureFlagsImpl();
+        featureFlags.setFlag(Flags.FLAG_ENABLED_RW, true);
+        assertThrows(NullPointerException.class, () -> Flags.enabledRw());
+        Flags.setFeatureFlags(featureFlags);
+        assertTrue(Flags.enabledRw());
+
+        Flags.unsetFeatureFlags();
+        assertThrows(NullPointerException.class, () -> Flags.enabledRw());
+    }
+
+    @Test
+    public void testFeatureFlagsImplNotImpl() {
+        FeatureFlags featureFlags = new FeatureFlagsImpl();
+        assertThrows(UnsupportedOperationException.class,
+            () -> featureFlags.enabledRw());
+    }
+}
diff --git a/tools/aconfig/tests/AconfigTest.java b/tools/aconfig/tests/AconfigTest.java
index 6681f32..317289d 100644
--- a/tools/aconfig/tests/AconfigTest.java
+++ b/tools/aconfig/tests/AconfigTest.java
@@ -8,12 +8,16 @@
 import static com.android.aconfig.test.Flags.enabledRw;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
+import com.android.aconfig.test.FakeFeatureFlagsImpl;
+import com.android.aconfig.test.FeatureFlags;
+
 @RunWith(JUnit4.class)
 public final class AconfigTest {
     @Test
@@ -43,4 +47,10 @@
         // (currently all flags are assigned the default READ_ONLY + DISABLED)
         assertFalse(enabledRw());
     }
+
+    @Test
+    public void testFakeFeatureFlagsImplNotImpl() {
+        FeatureFlags featureFlags = new FakeFeatureFlagsImpl();
+        assertThrows(UnsupportedOperationException.class, () -> featureFlags.enabledRw());
+    }
 }