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 }}
 }