diff --git a/tools/aconfig/templates/Flags.java.template b/tools/aconfig/templates/Flags.java.template
index 62116c5..eef98eb 100644
--- a/tools/aconfig/templates/Flags.java.template
+++ b/tools/aconfig/templates/Flags.java.template
@@ -9,6 +9,16 @@
         return FEATURE_FLAGS.{item.method_name}();
     }
     {{ endfor }}
-    private static FeatureFlags FEATURE_FLAGS = new FeatureFlagsImpl();
+    {{ if is_test_mode }}
+    public static void setFeatureFlagsImpl(FeatureFlags featureFlags) \{
+        Flags.FEATURE_FLAGS = featureFlags;
+    }
+
+    public static void unsetFeatureFlagsImpl() \{
+        Flags.FEATURE_FLAGS = null;
+    }
+    {{ -endif}}
+
+    private static FeatureFlags FEATURE_FLAGS{{ -if not is_test_mode }} = new FeatureFlagsImpl(){{ -endif- }};
 
 }