diff --git a/api/coverage/tools/Android.bp b/api/coverage/tools/Android.bp
index 3e16912..caaca99 100644
--- a/api/coverage/tools/Android.bp
+++ b/api/coverage/tools/Android.bp
@@ -30,3 +30,24 @@
         type: "full",
     },
 }
+
+java_test_host {
+    name: "extract-flagged-apis-test",
+    srcs: ["ExtractFlaggedApisTest.kt"],
+    libs: [
+        "extract_flagged_apis_proto",
+        "junit",
+        "libprotobuf-java-full",
+    ],
+    static_libs: [
+        "truth",
+        "truth-liteproto-extension",
+        "truth-proto-extension",
+    ],
+    data: [
+        ":extract-flagged-apis",
+    ],
+    test_options: {
+        unit_test: true,
+    },
+}
diff --git a/api/coverage/tools/ExtractFlaggedApis.kt b/api/coverage/tools/ExtractFlaggedApis.kt
index d5adfd0..5178f09 100644
--- a/api/coverage/tools/ExtractFlaggedApis.kt
+++ b/api/coverage/tools/ExtractFlaggedApis.kt
@@ -28,12 +28,10 @@
     val builder = FlagApiMap.newBuilder()
     for (pkg in cb.getPackages().packages) {
         val packageName = pkg.qualifiedName()
-        pkg.allClasses()
-            .filter { it.methods().size > 0 }
-            .forEach {
-                extractFlaggedApisFromClass(it, it.methods(), packageName, builder)
-                extractFlaggedApisFromClass(it, it.constructors(), packageName, builder)
-            }
+        pkg.allClasses().forEach {
+            extractFlaggedApisFromClass(it, it.methods(), packageName, builder)
+            extractFlaggedApisFromClass(it, it.constructors(), packageName, builder)
+        }
     }
     val flagApiMap = builder.build()
     FileWriter(args[1]).use { it.write(flagApiMap.toString()) }
@@ -45,6 +43,7 @@
     packageName: String,
     builder: FlagApiMap.Builder
 ) {
+    if (methods.isEmpty()) return
     val classFlag =
         classItem.modifiers
             .findAnnotation("android.annotation.FlaggedApi")
diff --git a/api/coverage/tools/ExtractFlaggedApisTest.kt b/api/coverage/tools/ExtractFlaggedApisTest.kt
new file mode 100644
index 0000000..ee5aaf1
--- /dev/null
+++ b/api/coverage/tools/ExtractFlaggedApisTest.kt
@@ -0,0 +1,160 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.platform.coverage
+
+import com.google.common.truth.extensions.proto.ProtoTruth.assertThat
+import com.google.protobuf.TextFormat
+import java.nio.file.Files
+import java.nio.file.Path
+import java.nio.file.StandardOpenOption
+import org.junit.After
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@RunWith(JUnit4::class)
+class ExtractFlaggedApisTest {
+
+    companion object {
+        const val COMMAND = "java -jar extract-flagged-apis.jar %s %s"
+    }
+
+    private var apiTextFile: Path = Files.createTempFile("current", ".txt")
+    private var flagToApiMap: Path = Files.createTempFile("flag_api_map", ".textproto")
+
+    @Before
+    fun setup() {
+        apiTextFile = Files.createTempFile("current", ".txt")
+        flagToApiMap = Files.createTempFile("flag_api_map", ".textproto")
+    }
+
+    @After
+    fun cleanup() {
+        Files.deleteIfExists(apiTextFile)
+        Files.deleteIfExists(flagToApiMap)
+    }
+
+    @Test
+    fun extractFlaggedApis_onlyMethodFlag_useMethodFlag() {
+        val apiText =
+            """
+            // Signature format: 2.0
+            package android.net.ipsec.ike {
+              public final class IkeSession implements java.lang.AutoCloseable {
+                method @FlaggedApi("com.android.ipsec.flags.dumpsys_api") public void dump(@NonNull java.io.PrintWriter);
+              }
+            }
+        """
+                .trimIndent()
+        Files.write(apiTextFile, apiText.toByteArray(Charsets.UTF_8), StandardOpenOption.APPEND)
+
+        val process = Runtime.getRuntime().exec(createCommand())
+        process.waitFor()
+
+        val content = Files.readAllBytes(flagToApiMap).toString(Charsets.UTF_8)
+        val result = TextFormat.parse(content, FlagApiMap::class.java)
+
+        val expected = FlagApiMap.newBuilder()
+        val api =
+            JavaMethod.newBuilder()
+                .setPackageName("android.net.ipsec.ike")
+                .setClassName("IkeSession")
+                .setMethodName("dump")
+        api.addParameters("java.io.PrintWriter")
+        addFlaggedApi(expected, api, "com.android.ipsec.flags.dumpsys_api")
+        assertThat(result).isEqualTo(expected.build())
+    }
+
+    @Test
+    fun extractFlaggedApis_onlyClassFlag_useClassFlag() {
+        val apiText =
+            """
+            // Signature format: 2.0
+            package android.net.ipsec.ike {
+              @FlaggedApi("com.android.ipsec.flags.dumpsys_api") public final class IkeSession implements java.lang.AutoCloseable {
+                method public void dump(@NonNull java.io.PrintWriter);
+              }
+            }
+        """
+                .trimIndent()
+        Files.write(apiTextFile, apiText.toByteArray(Charsets.UTF_8), StandardOpenOption.APPEND)
+
+        val process = Runtime.getRuntime().exec(createCommand())
+        process.waitFor()
+
+        val content = Files.readAllBytes(flagToApiMap).toString(Charsets.UTF_8)
+        val result = TextFormat.parse(content, FlagApiMap::class.java)
+
+        val expected = FlagApiMap.newBuilder()
+        val api =
+            JavaMethod.newBuilder()
+                .setPackageName("android.net.ipsec.ike")
+                .setClassName("IkeSession")
+                .setMethodName("dump")
+        api.addParameters("java.io.PrintWriter")
+        addFlaggedApi(expected, api, "com.android.ipsec.flags.dumpsys_api")
+        assertThat(result).isEqualTo(expected.build())
+    }
+
+    @Test
+    fun extractFlaggedApis_flaggedConstructorsAreFlaggedApis() {
+        val apiText =
+            """
+            // Signature format: 2.0
+            package android.app.pinner {
+              @FlaggedApi("android.app.pinner_service_client_api") public class PinnerServiceClient {
+                ctor @FlaggedApi("android.app.pinner_service_client_api") public PinnerServiceClient();
+              }
+            }
+        """
+                .trimIndent()
+        Files.write(apiTextFile, apiText.toByteArray(Charsets.UTF_8), StandardOpenOption.APPEND)
+
+        val process = Runtime.getRuntime().exec(createCommand())
+        process.waitFor()
+
+        val content = Files.readAllBytes(flagToApiMap).toString(Charsets.UTF_8)
+        val result = TextFormat.parse(content, FlagApiMap::class.java)
+
+        val expected = FlagApiMap.newBuilder()
+        val api =
+            JavaMethod.newBuilder()
+                .setPackageName("android.app.pinner")
+                .setClassName("PinnerServiceClient")
+                .setMethodName("PinnerServiceClient")
+        addFlaggedApi(expected, api, "android.app.pinner_service_client_api")
+        assertThat(result).isEqualTo(expected.build())
+    }
+
+    private fun addFlaggedApi(builder: FlagApiMap.Builder, api: JavaMethod.Builder, flag: String) {
+        if (builder.containsFlagToApi(flag)) {
+            val updatedApis =
+                builder.getFlagToApiOrThrow(flag).toBuilder().addJavaMethods(api).build()
+            builder.putFlagToApi(flag, updatedApis)
+        } else {
+            val apis = FlaggedApis.newBuilder().addJavaMethods(api).build()
+            builder.putFlagToApi(flag, apis)
+        }
+    }
+
+    private fun createCommand(): Array<String> {
+        val command =
+            String.format(COMMAND, apiTextFile.toAbsolutePath(), flagToApiMap.toAbsolutePath())
+        return command.split(" ").toTypedArray()
+    }
+}