Add annotation in CSTest to set flag values before instantiating CS

CSTest SetUp sets flag values based on annotations before
constructing ConnectivityService.
For flags that are loaded at ConnectivityService constructor, tests can
use this annotation to set flag values.

Test: TH
Change-Id: I5a995a3de4e433ac39e6ccc116b9d89da468b04b
diff --git a/tests/unit/java/com/android/server/connectivityservice/CSIngressDiscardRuleTests.kt b/tests/unit/java/com/android/server/connectivityservice/CSIngressDiscardRuleTests.kt
index e8664c1..bb7fb51 100644
--- a/tests/unit/java/com/android/server/connectivityservice/CSIngressDiscardRuleTests.kt
+++ b/tests/unit/java/com/android/server/connectivityservice/CSIngressDiscardRuleTests.kt
@@ -30,6 +30,7 @@
 import android.net.VpnTransportInfo
 import android.os.Build
 import androidx.test.filters.SmallTest
+import com.android.server.connectivity.ConnectivityFlags
 import com.android.testutils.DevSdkIgnoreRule
 import com.android.testutils.DevSdkIgnoreRunner
 import com.android.testutils.RecorderCallback.CallbackEntry.LinkPropertiesChanged
@@ -56,7 +57,9 @@
                         TYPE_VPN_SERVICE,
                         "MySession12345",
                         false /* bypassable */,
-                        false /* longLivedTcpConnectionsExpensive */))
+                        false /* longLivedTcpConnectionsExpensive */
+                )
+        )
         .build()
 
 private fun wifiNc() = NetworkCapabilities.Builder()
@@ -286,4 +289,19 @@
         waitForIdle()
         verify(bpfNetMaps).removeIngressDiscardRule(IPV6_ADDRESS)
     }
+
+    @Test @FeatureFlags([Flag(ConnectivityFlags.INGRESS_TO_VPN_ADDRESS_FILTERING, false)])
+    fun testVpnIngressDiscardRule_FeatureDisabled() {
+        val nr = nr(TRANSPORT_VPN)
+        val cb = TestableNetworkCallback()
+        cm.registerNetworkCallback(nr, cb)
+        val nc = vpnNc()
+        val lp = lp(VPN_IFNAME, IPV6_LINK_ADDRESS, LOCAL_IPV6_LINK_ADDRRESS)
+        val agent = Agent(nc = nc, lp = lp)
+        agent.connect()
+        cb.expectAvailableCallbacks(agent.network, validated = false)
+
+        // IngressDiscardRule should not be added since feature is disabled
+        verify(bpfNetMaps, never()).setIngressDiscardRule(any(), any())
+    }
 }
diff --git a/tests/unit/java/com/android/server/connectivityservice/base/CSTest.kt b/tests/unit/java/com/android/server/connectivityservice/base/CSTest.kt
index 5c4617b..bd26c63 100644
--- a/tests/unit/java/com/android/server/connectivityservice/base/CSTest.kt
+++ b/tests/unit/java/com/android/server/connectivityservice/base/CSTest.kt
@@ -79,11 +79,15 @@
 import java.util.concurrent.TimeUnit
 import java.util.function.BiConsumer
 import java.util.function.Consumer
+import kotlin.annotation.AnnotationRetention.RUNTIME
+import kotlin.annotation.AnnotationTarget.FUNCTION
 import kotlin.test.assertNotNull
 import kotlin.test.assertNull
 import kotlin.test.fail
 import org.junit.After
 import org.junit.Before
+import org.junit.Rule
+import org.junit.rules.TestName
 import org.mockito.AdditionalAnswers.delegatesTo
 import org.mockito.Mockito.doAnswer
 import org.mockito.Mockito.doReturn
@@ -126,6 +130,9 @@
 // TODO (b/272685721) : make ConnectivityServiceTest smaller and faster by moving the setup
 // parts into this class and moving the individual tests to multiple separate classes.
 open class CSTest {
+    @get:Rule
+    val testNameRule = TestName()
+
     companion object {
         val CSTestExecutor = Executors.newSingleThreadExecutor()
     }
@@ -155,8 +162,7 @@
         it[ConnectivityService.ALLOW_SATALLITE_NETWORK_FALLBACK] = true
         it[ConnectivityFlags.INGRESS_TO_VPN_ADDRESS_FILTERING] = true
     }
-    fun enableFeature(f: String) = enabledFeatures.set(f, true)
-    fun disableFeature(f: String) = enabledFeatures.set(f, false)
+    fun setFeatureEnabled(flag: String, enabled: Boolean) = enabledFeatures.set(flag, enabled)
 
     // When adding new members, consider if it's not better to build the object in CSTestHelpers
     // to keep this file clean of implementation details. Generally, CSTestHelpers should only
@@ -201,8 +207,32 @@
     lateinit var cm: ConnectivityManager
     lateinit var csHandler: Handler
 
+    // Tests can use this annotation to set flag values before constructing ConnectivityService
+    // e.g. @FeatureFlags([Flag(flagName1, true/false), Flag(flagName2, true/false)])
+    @Retention(RUNTIME)
+    @Target(FUNCTION)
+    annotation class FeatureFlags(val flags: Array<Flag>)
+
+    @Retention(RUNTIME)
+    @Target(FUNCTION)
+    annotation class Flag(val name: String, val enabled: Boolean)
+
     @Before
     fun setUp() {
+        // Set feature flags before constructing ConnectivityService
+        val testMethodName = testNameRule.methodName
+        try {
+            val testMethod = this::class.java.getMethod(testMethodName)
+            val featureFlags = testMethod.getAnnotation(FeatureFlags::class.java)
+            if (featureFlags != null) {
+                for (flag in featureFlags.flags) {
+                    setFeatureEnabled(flag.name, flag.enabled)
+                }
+            }
+        } catch (ignored: NoSuchMethodException) {
+            // This is expected for parameterized tests
+        }
+
         alarmHandlerThread = HandlerThread("TestAlarmManager").also { it.start() }
         alarmManager = makeMockAlarmManager(alarmHandlerThread)
         service = makeConnectivityService(context, netd, deps).also { it.systemReadyInternal() }