diff --git a/tools/aconfig/src/codegen_java.rs b/tools/aconfig/src/codegen_java.rs
index be0ec97..7cdf486 100644
--- a/tools/aconfig/src/codegen_java.rs
+++ b/tools/aconfig/src/codegen_java.rs
@@ -124,6 +124,7 @@
         boolean enabledFixedRo();
         boolean enabledRo();
         boolean enabledRw();
+    }
     "#;
 
     const EXPECTED_FLAG_COMMON_CONTENT: &str = r#"
@@ -152,32 +153,62 @@
         }
     "#;
 
-    const EXPECTED_METHOD_NOT_IMPL_COMMON_CONTENT: &str = r#"
+    const EXPECTED_FAKEFEATUREFLAGSIMPL_CONTENT: &str = r#"
+    package com.android.aconfig.test;
+    import java.util.HashMap;
+    import java.util.Map;
+    public class FakeFeatureFlagsImpl implements FeatureFlags {
+        public FakeFeatureFlagsImpl() {
+            resetAll();
+        }
         @Override
         public boolean disabledRo() {
-            throw new UnsupportedOperationException(
-                "Method is not implemented.");
+            return getFlag(Flags.FLAG_DISABLED_RO);
         }
         @Override
         public boolean disabledRw() {
-            throw new UnsupportedOperationException(
-                "Method is not implemented.");
+            return getFlag(Flags.FLAG_DISABLED_RW);
         }
         @Override
         public boolean enabledFixedRo() {
-            throw new UnsupportedOperationException(
-                "Method is not implemented.");
+            return getFlag(Flags.FLAG_ENABLED_FIXED_RO);
         }
         @Override
         public boolean enabledRo() {
-            throw new UnsupportedOperationException(
-                "Method is not implemented.");
+            return getFlag(Flags.FLAG_ENABLED_RO);
         }
         @Override
         public boolean enabledRw() {
-            throw new UnsupportedOperationException(
-                "Method is not implemented.");
+            return getFlag(Flags.FLAG_ENABLED_RW);
         }
+        public void setFlag(String flagName, boolean value) {
+            if (!this.mFlagMap.containsKey(flagName)) {
+                throw new IllegalArgumentException("no such flag " + flagName);
+            }
+            this.mFlagMap.put(flagName, value);
+        }
+        public void resetAll() {
+            for (Map.Entry entry : mFlagMap.entrySet()) {
+                entry.setValue(null);
+            }
+        }
+        private boolean getFlag(String flagName) {
+            Boolean value = this.mFlagMap.get(flagName);
+            if (value == null) {
+                throw new IllegalArgumentException(flagName + " is not set");
+            }
+            return value;
+        }
+        private Map<String, Boolean> mFlagMap = new HashMap<>(
+            Map.of(
+                Flags.FLAG_DISABLED_RO, false,
+                Flags.FLAG_DISABLED_RW, false,
+                Flags.FLAG_ENABLED_FIXED_RO, false,
+                Flags.FLAG_ENABLED_RO, false,
+                Flags.FLAG_ENABLED_RW, false
+            )
+        );
+    }
     "#;
 
     #[test]
@@ -189,21 +220,11 @@
             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();
         }"#;
-        let expect_fakefeatureflagsimpl_content = r#"
-        package com.android.aconfig.test;
-        public class FakeFeatureFlagsImpl implements FeatureFlags {"#
-            .to_owned()
-            + EXPECTED_METHOD_NOT_IMPL_COMMON_CONTENT
-            + r#"
-        }
-        "#;
+
         let expect_featureflagsimpl_content = r#"
         package com.android.aconfig.test;
         import android.provider.DeviceConfig;
@@ -241,10 +262,10 @@
         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", expect_featureflags_content.as_str()),
+            ("com/android/aconfig/test/FeatureFlags.java", EXPECTED_FEATUREFLAGS_COMMON_CONTENT),
             (
                 "com/android/aconfig/test/FakeFeatureFlagsImpl.java",
-                expect_fakefeatureflagsimpl_content.as_str(),
+                EXPECTED_FAKEFEATUREFLAGSIMPL_CONTENT,
             ),
         ]);
 
@@ -275,11 +296,7 @@
             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) {
@@ -293,94 +310,42 @@
         "#;
         let expect_featureflagsimpl_content = r#"
         package com.android.aconfig.test;
-        public final class FeatureFlagsImpl implements FeatureFlags {"#
-            .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#"
-        package com.android.aconfig.test;
-        import static java.util.stream.Collectors.toMap;
-        import java.util.HashMap;
-        import java.util.Map;
-        import java.util.stream.Stream;
-        public class FakeFeatureFlagsImpl implements FeatureFlags {
+        public final class FeatureFlagsImpl implements FeatureFlags {
             @Override
             public boolean disabledRo() {
-                return getFlag(Flags.FLAG_DISABLED_RO);
+                throw new UnsupportedOperationException(
+                    "Method is not implemented.");
             }
             @Override
             public boolean disabledRw() {
-                return getFlag(Flags.FLAG_DISABLED_RW);
+                throw new UnsupportedOperationException(
+                    "Method is not implemented.");
             }
             @Override
             public boolean enabledFixedRo() {
-                return getFlag(Flags.FLAG_ENABLED_FIXED_RO);
+                throw new UnsupportedOperationException(
+                    "Method is not implemented.");
             }
             @Override
             public boolean enabledRo() {
-                return getFlag(Flags.FLAG_ENABLED_RO);
+                throw new UnsupportedOperationException(
+                    "Method is not implemented.");
             }
             @Override
             public boolean enabledRw() {
-                return getFlag(Flags.FLAG_ENABLED_RW);
+                throw new UnsupportedOperationException(
+                    "Method is not implemented.");
             }
-            @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);
-                }
-            }
-            private boolean getFlag(String flagName) {
-                Boolean value = this.mFlagMap.get(flagName);
-                if (value == null) {
-                    throw new IllegalArgumentException(flagName + " is not set");
-                }
-                return value;
-            }
-            private HashMap<String, Boolean> mFlagMap = Stream.of(
-                    Flags.FLAG_DISABLED_RO,
-                    Flags.FLAG_DISABLED_RW,
-                    Flags.FLAG_ENABLED_FIXED_RO,
-                    Flags.FLAG_ENABLED_RO,
-                    Flags.FLAG_ENABLED_RW
-                )
-                .collect(
-                    HashMap::new,
-                    (map, elem) -> map.put(elem, null),
-                    HashMap::putAll
-                );
         }
         "#;
 
         let mut file_set = HashMap::from([
             ("com/android/aconfig/test/Flags.java", expect_flags_content.as_str()),
-            ("com/android/aconfig/test/FeatureFlags.java", expect_featureflags_content.as_str()),
-            (
-                "com/android/aconfig/test/FeatureFlagsImpl.java",
-                expect_featureflagsimpl_content.as_str(),
-            ),
+            ("com/android/aconfig/test/FeatureFlags.java", EXPECTED_FEATUREFLAGS_COMMON_CONTENT),
+            ("com/android/aconfig/test/FeatureFlagsImpl.java", expect_featureflagsimpl_content),
             (
                 "com/android/aconfig/test/FakeFeatureFlagsImpl.java",
-                expect_fakefeatureflagsimpl_content,
+                EXPECTED_FAKEFEATUREFLAGSIMPL_CONTENT,
             ),
         ]);
 
diff --git a/tools/aconfig/templates/FakeFeatureFlagsImpl.java.template b/tools/aconfig/templates/FakeFeatureFlagsImpl.java.template
index dba82ef..82bea81 100644
--- a/tools/aconfig/templates/FakeFeatureFlagsImpl.java.template
+++ b/tools/aconfig/templates/FakeFeatureFlagsImpl.java.template
@@ -1,27 +1,26 @@
 package {package_name};
-{{ if is_test_mode }}
-import static java.util.stream.Collectors.toMap;
 
 import java.util.HashMap;
 import java.util.Map;
-import java.util.stream.Stream;
 
 public class FakeFeatureFlagsImpl implements FeatureFlags \{
+    public FakeFeatureFlagsImpl() \{
+        resetAll();
+    }
+
 {{ for item in class_elements}}
     @Override
     public boolean {item.method_name}() \{
         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);
+            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);
@@ -36,26 +35,11 @@
         return value;
     }
 
-    private HashMap<String, Boolean> mFlagMap = Stream.of(
+    private Map<String, Boolean> mFlagMap = new HashMap<>(
+        Map.of(
             {{-for item in class_elements}}
-            Flags.FLAG_{item.flag_name_constant_suffix}{{ if not @last }},{{ endif }}
+            Flags.FLAG_{item.flag_name_constant_suffix}, false{{ if not @last }},{{ endif }}
             {{ -endfor }}
         )
-        .collect(
-            HashMap::new,
-            (map, elem) -> map.put(elem, null),
-            HashMap::putAll
-        );
+    );
 }
-{{ else }}
-{#- Generate only stub if in prod mode #}
-public class FakeFeatureFlagsImpl implements FeatureFlags \{
-{{ for item in class_elements}}
-    @Override
-    public boolean {item.method_name}() \{
-        throw new UnsupportedOperationException(
-            "Method is not implemented.");
-    }
-{{ endfor}}
-}
-{{ endif }}
diff --git a/tools/aconfig/templates/FeatureFlags.java.template b/tools/aconfig/templates/FeatureFlags.java.template
index c99ccbb..e0f201f 100644
--- a/tools/aconfig/templates/FeatureFlags.java.template
+++ b/tools/aconfig/templates/FeatureFlags.java.template
@@ -4,10 +4,4 @@
 {{ 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 7e1eb15..96de06c 100644
--- a/tools/aconfig/templates/FeatureFlagsImpl.java.template
+++ b/tools/aconfig/templates/FeatureFlagsImpl.java.template
@@ -17,7 +17,7 @@
         return {item.default_value};
     {{ endif- }}
     }
-{{ endfor- }}
+{{ endfor }}
 }
 {{ else }}
 {#- Generate only stub if in test mode #}
@@ -28,17 +28,6 @@
         throw new UnsupportedOperationException(
             "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.");
-    }
+{{ endfor }}
 }
 {{ endif }}
diff --git a/tools/aconfig/tests/AconfigHostTest.java b/tools/aconfig/tests/AconfigHostTest.java
index 29a01e3..ea71b7e 100644
--- a/tools/aconfig/tests/AconfigHostTest.java
+++ b/tools/aconfig/tests/AconfigHostTest.java
@@ -17,13 +17,13 @@
     @Test
     public void testThrowsExceptionIfFlagNotSet() {
         assertThrows(NullPointerException.class, () -> Flags.disabledRo());
-        FeatureFlags featureFlags = new FakeFeatureFlagsImpl();
+        FakeFeatureFlagsImpl featureFlags = new FakeFeatureFlagsImpl();
         assertThrows(IllegalArgumentException.class, () -> featureFlags.disabledRo());
     }
 
     @Test
     public void testSetFlagInFakeFeatureFlagsImpl() {
-        FeatureFlags featureFlags = new FakeFeatureFlagsImpl();
+        FakeFeatureFlagsImpl featureFlags = new FakeFeatureFlagsImpl();
         featureFlags.setFlag(Flags.FLAG_ENABLED_RW, true);
         assertTrue(featureFlags.enabledRw());
         featureFlags.setFlag(Flags.FLAG_ENABLED_RW, false);
@@ -39,14 +39,14 @@
 
     @Test
     public void testSetFlagWithRandomName() {
-        FeatureFlags featureFlags = new FakeFeatureFlagsImpl();
+        FakeFeatureFlagsImpl featureFlags = new FakeFeatureFlagsImpl();
         assertThrows(IllegalArgumentException.class,
             () -> featureFlags.setFlag("Randome_name", true));
     }
 
     @Test
     public void testResetFlagsInFakeFeatureFlagsImpl() {
-        FeatureFlags featureFlags = new FakeFeatureFlagsImpl();
+        FakeFeatureFlagsImpl featureFlags = new FakeFeatureFlagsImpl();
         featureFlags.setFlag(Flags.FLAG_ENABLED_RO, true);
         assertTrue(featureFlags.enabledRo());
         featureFlags.resetAll();
@@ -59,7 +59,7 @@
 
     @Test
     public void testFlagsSetFeatureFlags() {
-        FeatureFlags featureFlags = new FakeFeatureFlagsImpl();
+        FakeFeatureFlagsImpl featureFlags = new FakeFeatureFlagsImpl();
         featureFlags.setFlag(Flags.FLAG_ENABLED_RW, true);
         assertThrows(NullPointerException.class, () -> Flags.enabledRw());
         Flags.setFeatureFlags(featureFlags);
@@ -69,7 +69,7 @@
 
     @Test
     public void testFlagsUnsetFeatureFlags() {
-        FeatureFlags featureFlags = new FakeFeatureFlagsImpl();
+        FakeFeatureFlagsImpl featureFlags = new FakeFeatureFlagsImpl();
         featureFlags.setFlag(Flags.FLAG_ENABLED_RW, true);
         assertThrows(NullPointerException.class, () -> Flags.enabledRw());
         Flags.setFeatureFlags(featureFlags);
diff --git a/tools/aconfig/tests/AconfigTest.java b/tools/aconfig/tests/AconfigTest.java
index 79b5a55..958b02e 100644
--- a/tools/aconfig/tests/AconfigTest.java
+++ b/tools/aconfig/tests/AconfigTest.java
@@ -59,8 +59,9 @@
     }
 
     @Test
-    public void testFakeFeatureFlagsImplNotImpl() {
-        FeatureFlags featureFlags = new FakeFeatureFlagsImpl();
-        assertThrows(UnsupportedOperationException.class, () -> featureFlags.enabledRw());
+    public void testFakeFeatureFlagsImplImpled() {
+        FakeFeatureFlagsImpl fakeFeatureFlags = new FakeFeatureFlagsImpl();
+        fakeFeatureFlags.setFlag(FLAG_ENABLED_RW, false);
+        assertFalse(fakeFeatureFlags.enabledRw());
     }
 }