diff --git a/tools/aconfig/src/codegen_java.rs b/tools/aconfig/src/codegen_java.rs
index 79c6eb4..8ab6ffa 100644
--- a/tools/aconfig/src/codegen_java.rs
+++ b/tools/aconfig/src/codegen_java.rs
@@ -234,6 +234,7 @@
         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 final class FeatureFlagsImpl implements FeatureFlags {
             @Override
@@ -258,6 +259,11 @@
                 }
                 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) {