diff --git a/tools/aconfig/src/commands.rs b/tools/aconfig/src/commands.rs
index 58831cc..dd2087b 100644
--- a/tools/aconfig/src/commands.rs
+++ b/tools/aconfig/src/commands.rs
@@ -129,12 +129,18 @@
     Ok(output)
 }
 
-pub fn create_java_lib(mut input: Input) -> Result<Vec<OutputFile>> {
+#[derive(Copy, Clone, Debug, PartialEq, Eq, ValueEnum)]
+pub enum CodegenMode {
+    Production,
+    Test,
+}
+
+pub fn create_java_lib(mut input: Input, codegen_mode: CodegenMode) -> Result<Vec<OutputFile>> {
     let parsed_flags = input.try_parse_flags()?;
     let Some(package) = find_unique_package(&parsed_flags) else {
         bail!("no parsed flags, or the parsed flags use different packages");
     };
-    generate_java_code(package, parsed_flags.parsed_flag.iter())
+    generate_java_code(package, parsed_flags.parsed_flag.iter(), codegen_mode)
 }
 
 pub fn create_cpp_lib(mut input: Input) -> Result<OutputFile> {