diff --git a/staticlibs/testutils/devicetests/com/android/testutils/SetFeatureFlagsRule.kt b/staticlibs/testutils/devicetests/com/android/testutils/SetFeatureFlagsRule.kt
new file mode 100644
index 0000000..4185b05
--- /dev/null
+++ b/staticlibs/testutils/devicetests/com/android/testutils/SetFeatureFlagsRule.kt
@@ -0,0 +1,82 @@
+/*
+ * 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 com.android.testutils.com.android.testutils
+
+import org.junit.rules.TestRule
+import org.junit.runner.Description
+import org.junit.runners.model.Statement
+
+/**
+ * A JUnit Rule that sets feature flags based on `@FeatureFlag` annotations.
+ *
+ * This rule enables dynamic control of feature flag states during testing.
+ *
+ * **Usage:**
+ * ```kotlin
+ * class MyTestClass {
+ *   @get:Rule
+ *   val setFeatureFlagsRule = SetFeatureFlagsRule(setFlagsMethod = (name, enabled) -> {
+ *     // Custom handling code.
+ *   })
+ *
+ *   // ... test methods with @FeatureFlag annotations
+ *   @FeatureFlag("FooBar1", true)
+ *   @FeatureFlag("FooBar2", false)
+ *   @Test
+ *   fun testFooBar() {}
+ * }
+ * ```
+ */
+class SetFeatureFlagsRule(val setFlagsMethod: (name: String, enabled: Boolean) -> Unit) : TestRule {
+    /**
+     * This annotation marks a test method as requiring a specific feature flag to be configured.
+     *
+     * Use this on test methods to dynamically control feature flag states during testing.
+     *
+     * @param name The name of the feature flag.
+     * @param enabled The desired state (true for enabled, false for disabled) of the feature flag.
+     */
+    @Target(AnnotationTarget.FUNCTION)
+    @Retention(AnnotationRetention.RUNTIME)
+    annotation class FeatureFlag(val name: String, val enabled: Boolean = true)
+
+    /**
+     * This method is the core of the rule, executed by the JUnit framework before each test method.
+     *
+     * It retrieves the test method's metadata.
+     * If any `@FeatureFlag` annotation is found, it passes every feature flag's name
+     * and enabled state into the user-specified lambda to apply custom actions.
+     */
+    override fun apply(base: Statement, description: Description): Statement {
+        return object : Statement() {
+            override fun evaluate() {
+                val testMethod = description.testClass.getMethod(description.methodName)
+                val featureFlagAnnotations = testMethod.getAnnotationsByType(
+                    FeatureFlag::class.java
+                )
+
+                for (featureFlagAnnotation in featureFlagAnnotations) {
+                    setFlagsMethod(featureFlagAnnotation.name, featureFlagAnnotation.enabled)
+                }
+
+                // Execute the test method, which includes methods annotated with
+                // @Before, @Test and @After.
+                base.evaluate()
+            }
+        }
+    }
+}
diff --git a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
index 8efab46..91b7bf0 100644
--- a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
+++ b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
@@ -95,9 +95,6 @@
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
-import static java.lang.annotation.ElementType.METHOD;
-import static java.lang.annotation.RetentionPolicy.RUNTIME;
-
 import android.annotation.NonNull;
 import android.app.AlarmManager;
 import android.content.Context;
@@ -165,6 +162,8 @@
 import com.android.testutils.HandlerUtils;
 import com.android.testutils.TestBpfMap;
 import com.android.testutils.TestableNetworkStatsProviderBinder;
+import com.android.testutils.com.android.testutils.SetFeatureFlagsRule;
+import com.android.testutils.com.android.testutils.SetFeatureFlagsRule.FeatureFlag;
 
 import libcore.testing.io.TestIoUtils;
 
@@ -173,7 +172,6 @@
 import org.junit.Ignore;
 import org.junit.Rule;
 import org.junit.Test;
-import org.junit.rules.TestName;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
@@ -183,9 +181,6 @@
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
 import java.io.StringWriter;
-import java.lang.annotation.Retention;
-import java.lang.annotation.Target;
-import java.lang.reflect.Method;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.time.Clock;
@@ -214,8 +209,6 @@
 // NetworkStatsService is not updatable before T, so tests do not need to be backwards compatible
 @DevSdkIgnoreRule.IgnoreUpTo(SC_V2)
 public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
-    @Rule
-    public final TestName mTestNameRule = new TestName();
 
     private static final String TAG = "NetworkStatsServiceTest";
 
@@ -312,6 +305,15 @@
     final TestDependencies mDeps = new TestDependencies();
     final HashMap<String, Boolean> mFeatureFlags = new HashMap<>();
 
+    // This will set feature flags from @FeatureFlag annotations
+    // into the map before setUp() runs.
+    @Rule
+    public final SetFeatureFlagsRule mSetFeatureFlagsRule =
+            new SetFeatureFlagsRule((name, enabled) -> {
+                mFeatureFlags.put(name, enabled);
+                return null;
+            });
+
     private class MockContext extends BroadcastInterceptingContext {
         private final Context mBaseContext;
 
@@ -369,33 +371,6 @@
         return parcel;
     }
 
-    // Tests can use this annotation to set feature flags before constructing
-    // NetworkStatsService, e.g. @FeatureFlag(FeatureName, true/false)).
-    // TODO: Refactor into a Rule, and put in a common place.
-    @Retention(RUNTIME)
-    @Target(METHOD)
-    public @interface FeatureFlag {
-        String name();
-
-        boolean enabled() default true;
-    }
-
-    private void initFeatureFlagsFromAnnotations() {
-        // Setup default and overrides feature flags before creating the service.
-        mFeatureFlags.put(TRAFFICSTATS_RATE_LIMIT_CACHE_ENABLED_FLAG, true);
-
-        final String testMethodName = mTestNameRule.getMethodName();
-        try {
-            final Method method = this.getClass().getMethod(testMethodName);
-            final FeatureFlag[] flags = method.getAnnotationsByType(FeatureFlag.class);
-            for (final FeatureFlag flag : flags) {
-                mFeatureFlags.put(flag.name(), flag.enabled());
-            }
-        } catch (NoSuchMethodException ignored) {
-            // This is expected for parameterized tests
-        }
-    }
-
     @Before
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
@@ -419,9 +394,6 @@
         PowerManager.WakeLock wakeLock =
                 powerManager.newWakeLock(PowerManager.PARTIAL_WAKE_LOCK, TAG);
 
-        // This has to be invoked before initialize the service instance.
-        initFeatureFlagsFromAnnotations();
-
         mHandlerThread = new HandlerThread("NetworkStatsServiceTest-HandlerThread");
         // Create a separate thread for observers to run on. This thread cannot be the same
         // as the handler thread, because the observer callback is fired on this thread, and
@@ -2433,7 +2405,7 @@
         doTestTrafficStatsRateLimitCache(false /* cacheEnabled */);
     }
 
-    @FeatureFlag(name = TRAFFICSTATS_RATE_LIMIT_CACHE_ENABLED_FLAG, enabled = true)
+    @FeatureFlag(name = TRAFFICSTATS_RATE_LIMIT_CACHE_ENABLED_FLAG)
     @Test
     public void testTrafficStatsRateLimitCache_enabled() throws Exception {
         doTestTrafficStatsRateLimitCache(true /* cacheEnabled */);