diff --git a/tools/aconfig/Android.bp b/tools/aconfig/Android.bp
index a4ea7f4..93fd6f7 100644
--- a/tools/aconfig/Android.bp
+++ b/tools/aconfig/Android.bp
@@ -101,7 +101,7 @@
 android_test {
     name: "aconfig.test.java",
     srcs: [
-        "tests/**/*.java",
+        "tests/AconfigTest.java",
     ],
     manifest: "tests/AndroidManifest.xml",
     certificate: "platform",
@@ -113,6 +113,25 @@
     test_suites: ["device-tests"],
 }
 
+java_aconfig_library {
+    name: "aconfig_host_test_java_library",
+    aconfig_declarations: "aconfig.test.flags",
+    host_supported: true,
+    test: true,
+}
+
+java_test_host {
+    name: "AconfigJavaHostTest",
+    srcs: [
+        "tests/AconfigHostTest.java",
+    ],
+    static_libs: [
+        "aconfig_host_test_java_library",
+        "junit",
+    ],
+    test_suites: ["general-tests"],
+}
+
 // integration tests: C++
 
 cc_aconfig_library {
diff --git a/tools/aconfig/src/codegen_java.rs b/tools/aconfig/src/codegen_java.rs
index a117370..2c9dcf5 100644
--- a/tools/aconfig/src/codegen_java.rs
+++ b/tools/aconfig/src/codegen_java.rs
@@ -116,14 +116,14 @@
     use super::*;
     use std::collections::HashMap;
 
-    const EXPECTED_FEATUREFLAGS_CONTENT: &str = r#"
+    const EXPECTED_FEATUREFLAGS_COMMON_CONTENT: &str = r#"
     package com.android.aconfig.test;
     public interface FeatureFlags {
         boolean disabledRo();
         boolean disabledRw();
         boolean enabledRo();
         boolean enabledRw();
-    }"#;
+    "#;
 
     const EXPECTED_FLAG_COMMON_CONTENT: &str = r#"
     package com.android.aconfig.test;
@@ -179,6 +179,9 @@
             CodegenMode::Production,
         )
         .unwrap();
+        let expect_featureflags_content = EXPECTED_FEATUREFLAGS_COMMON_CONTENT.to_string()
+            + r#"
+        }"#;
         let expect_flags_content = EXPECTED_FLAG_COMMON_CONTENT.to_string()
             + r#"
             private static FeatureFlags FEATURE_FLAGS = new FeatureFlagsImpl();
@@ -224,7 +227,7 @@
         let mut file_set = HashMap::from([
             ("com/android/aconfig/test/Flags.java", expect_flags_content.as_str()),
             ("com/android/aconfig/test/FeatureFlagsImpl.java", expect_featureflagsimpl_content),
-            ("com/android/aconfig/test/FeatureFlags.java", EXPECTED_FEATUREFLAGS_CONTENT),
+            ("com/android/aconfig/test/FeatureFlags.java", expect_featureflags_content.as_str()),
             (
                 "com/android/aconfig/test/FakeFeatureFlagsImpl.java",
                 expect_fakefeatureflagsimpl_content.as_str(),
@@ -258,6 +261,11 @@
             CodegenMode::Test,
         )
         .unwrap();
+        let expect_featureflags_content = EXPECTED_FEATUREFLAGS_COMMON_CONTENT.to_string()
+            + r#"
+            public void setFlag(String flagName, boolean value);
+            public void resetAll();
+        }"#;
         let expect_flags_content = EXPECTED_FLAG_COMMON_CONTENT.to_string()
             + r#"
             public static void setFeatureFlags(FeatureFlags featureFlags) {
@@ -275,6 +283,16 @@
             .to_owned()
             + EXPECTED_METHOD_NOT_IMPL_COMMON_CONTENT
             + r#"
+            @Override
+            public void setFlag(String flagName, boolean value) {
+                throw new UnsupportedOperationException(
+                    "Method is not implemented.");
+            }
+            @Override
+            public void resetAll() {
+                throw new UnsupportedOperationException(
+                    "Method is not implemented.");
+            }
         }
         "#;
         let expect_fakefeatureflagsimpl_content = r#"
@@ -300,12 +318,14 @@
             public boolean enabledRw() {
                 return getFlag(Flags.FLAG_ENABLED_RW);
             }
+            @Override
             public void setFlag(String flagName, boolean value) {
                 if (!this.mFlagMap.containsKey(flagName)) {
                     throw new IllegalArgumentException("no such flag" + flagName);
                 }
                 this.mFlagMap.put(flagName, value);
             }
+            @Override
             public void resetAll() {
                 for (Map.Entry entry : mFlagMap.entrySet()) {
                     entry.setValue(null);
@@ -334,7 +354,7 @@
 
         let mut file_set = HashMap::from([
             ("com/android/aconfig/test/Flags.java", expect_flags_content.as_str()),
-            ("com/android/aconfig/test/FeatureFlags.java", EXPECTED_FEATUREFLAGS_CONTENT),
+            ("com/android/aconfig/test/FeatureFlags.java", expect_featureflags_content.as_str()),
             (
                 "com/android/aconfig/test/FeatureFlagsImpl.java",
                 expect_featureflagsimpl_content.as_str(),
diff --git a/tools/aconfig/templates/FakeFeatureFlagsImpl.java.template b/tools/aconfig/templates/FakeFeatureFlagsImpl.java.template
index db35d28..dba82ef 100644
--- a/tools/aconfig/templates/FakeFeatureFlagsImpl.java.template
+++ b/tools/aconfig/templates/FakeFeatureFlagsImpl.java.template
@@ -13,6 +13,7 @@
         return getFlag(Flags.FLAG_{item.flag_name_constant_suffix});
     }
 {{ endfor}}
+    @Override
     public void setFlag(String flagName, boolean value) \{
         if (!this.mFlagMap.containsKey(flagName)) \{
             throw new IllegalArgumentException("no such flag" + flagName);
@@ -20,6 +21,7 @@
         this.mFlagMap.put(flagName, value);
     }
 
+    @Override
     public void resetAll() \{
         for (Map.Entry entry : mFlagMap.entrySet()) \{
             entry.setValue(null);
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 }}
 }
diff --git a/tools/aconfig/templates/FeatureFlagsImpl.java.template b/tools/aconfig/templates/FeatureFlagsImpl.java.template
index ba86ce5..7e1eb15 100644
--- a/tools/aconfig/templates/FeatureFlagsImpl.java.template
+++ b/tools/aconfig/templates/FeatureFlagsImpl.java.template
@@ -29,5 +29,16 @@
             "Method is not implemented.");
     }
 {{ endfor- }}
+    @Override
+    public void setFlag(String flagName, boolean value) \{
+        throw new UnsupportedOperationException(
+            "Method is not implemented.");
+    }
+
+    @Override
+    public void resetAll() \{
+        throw new UnsupportedOperationException(
+            "Method is not implemented.");
+    }
 }
 {{ endif }}
diff --git a/tools/aconfig/tests/AconfigHostTest.java b/tools/aconfig/tests/AconfigHostTest.java
new file mode 100644
index 0000000..29a01e3
--- /dev/null
+++ b/tools/aconfig/tests/AconfigHostTest.java
@@ -0,0 +1,88 @@
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+
+import com.android.aconfig.test.FakeFeatureFlagsImpl;
+import com.android.aconfig.test.FeatureFlags;
+import com.android.aconfig.test.FeatureFlagsImpl;
+import com.android.aconfig.test.Flags;
+
+@RunWith(JUnit4.class)
+public final class AconfigHostTest {
+    @Test
+    public void testThrowsExceptionIfFlagNotSet() {
+        assertThrows(NullPointerException.class, () -> Flags.disabledRo());
+        FeatureFlags featureFlags = new FakeFeatureFlagsImpl();
+        assertThrows(IllegalArgumentException.class, () -> featureFlags.disabledRo());
+    }
+
+    @Test
+    public void testSetFlagInFakeFeatureFlagsImpl() {
+        FeatureFlags featureFlags = new FakeFeatureFlagsImpl();
+        featureFlags.setFlag(Flags.FLAG_ENABLED_RW, true);
+        assertTrue(featureFlags.enabledRw());
+        featureFlags.setFlag(Flags.FLAG_ENABLED_RW, false);
+        assertFalse(featureFlags.enabledRw());
+
+        //Set Flags
+        assertThrows(NullPointerException.class, () -> Flags.enabledRw());
+        Flags.setFeatureFlags(featureFlags);
+        featureFlags.setFlag(Flags.FLAG_ENABLED_RW, true);
+        assertTrue(Flags.enabledRw());
+        Flags.unsetFeatureFlags();
+    }
+
+    @Test
+    public void testSetFlagWithRandomName() {
+        FeatureFlags featureFlags = new FakeFeatureFlagsImpl();
+        assertThrows(IllegalArgumentException.class,
+            () -> featureFlags.setFlag("Randome_name", true));
+    }
+
+    @Test
+    public void testResetFlagsInFakeFeatureFlagsImpl() {
+        FeatureFlags featureFlags = new FakeFeatureFlagsImpl();
+        featureFlags.setFlag(Flags.FLAG_ENABLED_RO, true);
+        assertTrue(featureFlags.enabledRo());
+        featureFlags.resetAll();
+        assertThrows(IllegalArgumentException.class, () -> featureFlags.enabledRo());
+
+        // Set value after reset
+        featureFlags.setFlag(Flags.FLAG_ENABLED_RO, false);
+        assertFalse(featureFlags.enabledRo());
+    }
+
+    @Test
+    public void testFlagsSetFeatureFlags() {
+        FeatureFlags featureFlags = new FakeFeatureFlagsImpl();
+        featureFlags.setFlag(Flags.FLAG_ENABLED_RW, true);
+        assertThrows(NullPointerException.class, () -> Flags.enabledRw());
+        Flags.setFeatureFlags(featureFlags);
+        assertTrue(Flags.enabledRw());
+        Flags.unsetFeatureFlags();
+    }
+
+    @Test
+    public void testFlagsUnsetFeatureFlags() {
+        FeatureFlags featureFlags = new FakeFeatureFlagsImpl();
+        featureFlags.setFlag(Flags.FLAG_ENABLED_RW, true);
+        assertThrows(NullPointerException.class, () -> Flags.enabledRw());
+        Flags.setFeatureFlags(featureFlags);
+        assertTrue(Flags.enabledRw());
+
+        Flags.unsetFeatureFlags();
+        assertThrows(NullPointerException.class, () -> Flags.enabledRw());
+    }
+
+    @Test
+    public void testFeatureFlagsImplNotImpl() {
+        FeatureFlags featureFlags = new FeatureFlagsImpl();
+        assertThrows(UnsupportedOperationException.class,
+            () -> featureFlags.enabledRw());
+    }
+}
diff --git a/tools/aconfig/tests/AconfigTest.java b/tools/aconfig/tests/AconfigTest.java
index 6681f32..317289d 100644
--- a/tools/aconfig/tests/AconfigTest.java
+++ b/tools/aconfig/tests/AconfigTest.java
@@ -8,12 +8,16 @@
 import static com.android.aconfig.test.Flags.enabledRw;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
+import com.android.aconfig.test.FakeFeatureFlagsImpl;
+import com.android.aconfig.test.FeatureFlags;
+
 @RunWith(JUnit4.class)
 public final class AconfigTest {
     @Test
@@ -43,4 +47,10 @@
         // (currently all flags are assigned the default READ_ONLY + DISABLED)
         assertFalse(enabledRw());
     }
+
+    @Test
+    public void testFakeFeatureFlagsImplNotImpl() {
+        FeatureFlags featureFlags = new FakeFeatureFlagsImpl();
+        assertThrows(UnsupportedOperationException.class, () -> featureFlags.enabledRw());
+    }
 }