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 @Repeatable 61 @Retention(AnnotationRetention.RUNTIME) 62 annotation class FeatureFlag(val name: String, val enabled: Boolean = true) 63 64 /** 65 * This method is the core of the rule, executed by the JUnit framework before each test method. 66 * 67 * It retrieves the test method's metadata. 68 * If any `@FeatureFlag` annotation is found, it passes every feature flag's name 69 * and enabled state into the user-specified lambda to apply custom actions. 70 */ 71 private val parameterizedRegexp = Regex("\\[\\d+\\]$") applynull72 override fun apply(base: Statement, description: Description): Statement { 73 return object : Statement() { 74 override fun evaluate() { 75 // If the same class also uses Parameterized, depending on evaluation order the 76 // method names here may be synthetic method names, where [0] [1] or so are added 77 // at the end of the method name. Find the original method name. 78 val methodName = description.methodName.replace(parameterizedRegexp, "") 79 val testMethod = description.testClass.getMethod(methodName) 80 val featureFlagAnnotations = testMethod.getAnnotationsByType( 81 FeatureFlag::class.java 82 ) 83 84 val valuesToBeRestored = mutableMapOf<String, Boolean?>() 85 for (featureFlagAnnotation in featureFlagAnnotations) { 86 valuesToBeRestored[featureFlagAnnotation.name] = 87 getFlagsMethod(featureFlagAnnotation.name) 88 setFlagsMethod(featureFlagAnnotation.name, featureFlagAnnotation.enabled) 89 } 90 91 // Execute the test method, which includes methods annotated with 92 // @Before, @Test and @After. 93 base.evaluate() 94 95 valuesToBeRestored.forEach { 96 setFlagsMethod(it.key, it.value) 97 } 98 } 99 } 100 } 101 } 102