diff --git a/tools/aconfig/templates/FeatureFlagsImpl.java.template b/tools/aconfig/templates/FeatureFlagsImpl.java.template
index 082d476..ba86ce5 100644
--- a/tools/aconfig/templates/FeatureFlagsImpl.java.template
+++ b/tools/aconfig/templates/FeatureFlagsImpl.java.template
@@ -1,65 +1,33 @@
 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;
-{{ else}}
+{{ if not is_test_mode }}
 {{ if is_read_write- }}
 import android.provider.DeviceConfig;
-{{ -endif- }}
 {{ endif }}
 public final class FeatureFlagsImpl implements FeatureFlags \{
 {{ for item in class_elements}}
     @Override
     public boolean {item.method_name}() \{
-        {{ -if not is_test_mode- }}
-        {{ if item.is_read_write }}
+    {{ -if item.is_read_write }}
         return DeviceConfig.getBoolean(
             "{item.device_config_namespace}",
             "{item.device_config_flag}",
             {item.default_value}
         );
-        {{ -else }}
+    {{ else }}
         return {item.default_value};
-        {{ -endif- }}
-        {{ else }}
-        return getFlag(Flags.FLAG_{item.flag_name_constant_suffix});
-        {{ -endif }}
+    {{ endif- }}
     }
 {{ endfor- }}
-{{ if is_test_mode }}
-    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 HashMap<String, Boolean> mFlagMap = Stream.of(
-            {{-for item in class_elements}}
-            Flags.FLAG_{item.flag_name_constant_suffix}{{ if not @last }},{{ endif }}
-            {{ -endfor }}
-        )
-        .collect(
-            HashMap::new,
-            (map, elem) -> map.put(elem, null),
-            HashMap::putAll
-        );
-{{ -endif }}
 }
+{{ else }}
+{#- Generate only stub if in test mode #}
+public final class FeatureFlagsImpl implements FeatureFlags \{
+{{ for item in class_elements}}
+    @Override
+    public boolean {item.method_name}() \{
+        throw new UnsupportedOperationException(
+            "Method is not implemented.");
+    }
+{{ endfor- }}
+}
+{{ endif }}