1 /*
2  * Copyright (C) 2024 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.testutils.com.android.testutils
18 
19 import org.junit.rules.TestRule
20 import org.junit.runner.Description
21 import org.junit.runners.model.Statement
22 
23 /**
24  * A JUnit Rule that sets feature flags based on `@FeatureFlag` annotations.
25  *
26  * This rule enables dynamic control of feature flag states during testing.
27  * And restores the original values after performing tests.
28  *
29  * **Usage:**
30  * ```kotlin
31  * class MyTestClass {
32  *   @get:Rule
33  *   val setFeatureFlagsRule = SetFeatureFlagsRule(setFlagsMethod = (name, enabled) -> {
34  *     // Custom handling code.
35  *   }, (name) -> {
36  *     // Custom getter code to retrieve the original values.
37  *   })
38  *
39  *   // ... test methods with @FeatureFlag annotations
40  *   @FeatureFlag("FooBar1", true)
41  *   @FeatureFlag("FooBar2", false)
42  *   @Test
43  *   fun testFooBar() {}
44  * }
45  * ```
46  */
47 class SetFeatureFlagsRule(
48     val setFlagsMethod: (name: String, enabled: Boolean?) -> Unit,
49                           val getFlagsMethod: (name: String) -> Boolean?
50 ) : TestRule {
51     /**
52      * This annotation marks a test method as requiring a specific feature flag to be configured.
53      *
54      * Use this on test methods to dynamically control feature flag states during testing.
55      *
56      * @param name The name of the feature flag.
57      * @param enabled The desired state (true for enabled, false for disabled) of the feature flag.
58      */
59     @Target(AnnotationTarget.FUNCTION)
60     @Retention(AnnotationRetention.RUNTIME)
61     annotation class FeatureFlag(val name: String, val enabled: Boolean = true)
62 
63     /**
64      * This method is the core of the rule, executed by the JUnit framework before each test method.
65      *
66      * It retrieves the test method's metadata.
67      * If any `@FeatureFlag` annotation is found, it passes every feature flag's name
68      * and enabled state into the user-specified lambda to apply custom actions.
69      */
applynull70     override fun apply(base: Statement, description: Description): Statement {
71         return object : Statement() {
72             override fun evaluate() {
73                 val testMethod = description.testClass.getMethod(description.methodName)
74                 val featureFlagAnnotations = testMethod.getAnnotationsByType(
75                     FeatureFlag::class.java
76                 )
77 
78                 val valuesToBeRestored = mutableMapOf<String, Boolean?>()
79                 for (featureFlagAnnotation in featureFlagAnnotations) {
80                     valuesToBeRestored[featureFlagAnnotation.name] =
81                             getFlagsMethod(featureFlagAnnotation.name)
82                     setFlagsMethod(featureFlagAnnotation.name, featureFlagAnnotation.enabled)
83                 }
84 
85                 // Execute the test method, which includes methods annotated with
86                 // @Before, @Test and @After.
87                 base.evaluate()
88 
89                 valuesToBeRestored.forEach {
90                     setFlagsMethod(it.key, it.value)
91                 }
92             }
93         }
94     }
95 }
96