| /* |
| * Copyright (C) 2023 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.test.flag.junit; |
| |
| import static org.junit.Assume.assumeFalse; |
| |
| import android.platform.test.flag.util.FlagReadException; |
| import android.platform.test.flag.util.FlagSetException; |
| |
| import com.google.common.base.CaseFormat; |
| import com.google.common.collect.Sets; |
| |
| import org.junit.rules.TestRule; |
| import org.junit.runner.Description; |
| import org.junit.runners.model.Statement; |
| |
| import java.lang.reflect.Field; |
| import java.util.HashMap; |
| import java.util.HashSet; |
| import java.util.Map; |
| import java.util.Objects; |
| import java.util.Set; |
| import java.util.concurrent.ConcurrentHashMap; |
| import java.util.function.BiPredicate; |
| import java.util.function.Predicate; |
| |
| import javax.annotation.Nonnull; |
| import javax.annotation.Nullable; |
| |
| /** A {@link TestRule} that helps to set flag values in unit test. */ |
| public final class SetFlagsRule implements TestRule { |
| private static final String FAKE_FEATURE_FLAGS_IMPL_CLASS_NAME = "FakeFeatureFlagsImpl"; |
| private static final String REAL_FEATURE_FLAGS_IMPL_CLASS_NAME = "FeatureFlagsImpl"; |
| private static final String CUSTOM_FEATURE_FLAGS_CLASS_NAME = "CustomFeatureFlags"; |
| private static final String FEATURE_FLAGS_CLASS_NAME = "FeatureFlags"; |
| private static final String FEATURE_FLAGS_FIELD_NAME = "FEATURE_FLAGS"; |
| private static final String FLAGS_CLASS_NAME = "Flags"; |
| private static final String FLAG_CONSTANT_PREFIX = "FLAG_"; |
| private static final String SET_FLAG_METHOD_NAME = "setFlag"; |
| private static final String RESET_ALL_METHOD_NAME = "resetAll"; |
| private static final String IS_FLAG_READ_ONLY_OPTIMIZED_METHOD_NAME = "isFlagReadOnlyOptimized"; |
| |
| // Store instances for entire life of a SetFlagsRule instance |
| private final Map<Class<?>, Object> mFlagsClassToFakeFlagsImpl = new HashMap<>(); |
| private final Map<Class<?>, Object> mFlagsClassToRealFlagsImpl = new HashMap<>(); |
| |
| // Store classes that are currently mutated by this rule |
| private final Set<Class<?>> mMutatedFlagsClasses = new HashSet<>(); |
| |
| // Any flags added to this list cannot be set imperatively (i.e. with enableFlags/disableFlags) |
| private final Set<String> mLockedFlagNames = new HashSet<>(); |
| |
| // listener to be called before setting a flag |
| private final Listener mListener; |
| |
| // TODO(322377082): remove repackage prefix list |
| private static final String[] REPACKAGE_PREFIX_LIST = |
| new String[] { |
| "", "com.android.internal.hidden_from_bootclasspath.", |
| }; |
| private final Map<String, Set<String>> mPackageToRepackage = new HashMap<>(); |
| |
| private final boolean mIsInitWithDefault; |
| private FlagsParameterization mFlagsParameterization; |
| private boolean mIsRuleEvaluating = false; |
| |
| public enum DefaultInitValueType { |
| /** |
| * Initialize flag value as null |
| * |
| * <p>Flag value need to be set before using |
| */ |
| NULL_DEFAULT, |
| |
| /** |
| * Initialize flag value with the default value from the device |
| * |
| * <p>If flag value is not overridden by adb, then the default value is from the release |
| * configuration when the test is built. |
| */ |
| DEVICE_DEFAULT, |
| } |
| |
| public SetFlagsRule() { |
| this(DefaultInitValueType.DEVICE_DEFAULT); |
| } |
| |
| public SetFlagsRule(DefaultInitValueType defaultType) { |
| this(defaultType, null); |
| } |
| |
| public SetFlagsRule(@Nullable FlagsParameterization flagsParameterization) { |
| this(DefaultInitValueType.DEVICE_DEFAULT, flagsParameterization); |
| } |
| |
| public SetFlagsRule( |
| DefaultInitValueType defaultType, |
| @Nullable FlagsParameterization flagsParameterization) { |
| this(defaultType, flagsParameterization, null); |
| } |
| |
| private SetFlagsRule( |
| DefaultInitValueType defaultType, |
| @Nullable FlagsParameterization flagsParameterization, |
| @Nullable Listener listener) { |
| mIsInitWithDefault = defaultType == DefaultInitValueType.DEVICE_DEFAULT; |
| mFlagsParameterization = flagsParameterization; |
| if (flagsParameterization != null) { |
| mLockedFlagNames.addAll(flagsParameterization.mOverrides.keySet()); |
| } |
| mListener = listener; |
| } |
| |
| /** |
| * Set the FlagsParameterization to be used during this test. This cannot be used to override a |
| * previous call, and cannot be called once the rule has been evaluated. |
| */ |
| public void setFlagsParameterization(@Nonnull FlagsParameterization flagsParameterization) { |
| Objects.requireNonNull(flagsParameterization, "FlagsParameterization cannot be cleared"); |
| if (mFlagsParameterization != null) { |
| throw new AssertionError("FlagsParameterization cannot be overridden"); |
| } |
| if (mIsRuleEvaluating) { |
| throw new AssertionError("Cannot set FlagsParameterization once the rule is running"); |
| } |
| ensureFlagsAreUnset(); |
| mFlagsParameterization = flagsParameterization; |
| mLockedFlagNames.addAll(flagsParameterization.mOverrides.keySet()); |
| } |
| |
| /** |
| * Enables the given flags. |
| * |
| * @param fullFlagNames The name of the flags in the flag class with the format |
| * {packageName}.{flagName} |
| * |
| * @deprecated Annotate your test or class with <code>@EnableFlags(String...)</code> instead |
| */ |
| @Deprecated |
| public void enableFlags(String... fullFlagNames) { |
| if (!mIsRuleEvaluating) { |
| throw new IllegalStateException("Not allowed to set flags outside test and setup code"); |
| } |
| for (String fullFlagName : fullFlagNames) { |
| if (mLockedFlagNames.contains(fullFlagName)) { |
| throw new FlagSetException(fullFlagName, "Not allowed to change locked flags"); |
| } |
| setFlagValue(fullFlagName, true); |
| } |
| } |
| |
| /** |
| * Disables the given flags. |
| * |
| * @param fullFlagNames The name of the flags in the flag class with the format |
| * {packageName}.{flagName} |
| * |
| * @deprecated Annotate your test or class with <code>@DisableFlags(String...)</code> instead |
| */ |
| @Deprecated |
| public void disableFlags(String... fullFlagNames) { |
| if (!mIsRuleEvaluating) { |
| throw new IllegalStateException("Not allowed to set flags outside test and setup code"); |
| } |
| for (String fullFlagName : fullFlagNames) { |
| if (mLockedFlagNames.contains(fullFlagName)) { |
| throw new FlagSetException(fullFlagName, "Not allowed to change locked flags"); |
| } |
| setFlagValue(fullFlagName, false); |
| } |
| } |
| |
| private void ensureFlagsAreUnset() { |
| if (!mFlagsClassToFakeFlagsImpl.isEmpty()) { |
| throw new IllegalStateException("Some flags were set before the rule was initialized"); |
| } |
| } |
| |
| @Override |
| public Statement apply(Statement base, Description description) { |
| return new Statement() { |
| @Override |
| public void evaluate() throws Throwable { |
| Throwable throwable = null; |
| try { |
| if (mListener != null) { |
| mListener.onStartedEvaluating(); |
| } |
| AnnotationsRetriever.FlagAnnotations flagAnnotations = |
| AnnotationsRetriever.getFlagAnnotations(description); |
| assertAnnotationsMatchParameterization(flagAnnotations, mFlagsParameterization); |
| flagAnnotations.assumeAllSetFlagsMatchParameterization(mFlagsParameterization); |
| if (mFlagsParameterization != null) { |
| ensureFlagsAreUnset(); |
| for (Map.Entry<String, Boolean> pair : |
| mFlagsParameterization.mOverrides.entrySet()) { |
| setFlagValue(pair.getKey(), pair.getValue()); |
| } |
| } |
| for (Map.Entry<String, Boolean> pair : |
| flagAnnotations.mSetFlagValues.entrySet()) { |
| setFlagValue(pair.getKey(), pair.getValue()); |
| } |
| mLockedFlagNames.addAll(flagAnnotations.mRequiredFlagValues.keySet()); |
| mLockedFlagNames.addAll(flagAnnotations.mSetFlagValues.keySet()); |
| mIsRuleEvaluating = true; |
| base.evaluate(); |
| } catch (Throwable t) { |
| throwable = t; |
| } finally { |
| mIsRuleEvaluating = false; |
| try { |
| resetFlags(); |
| } catch (Throwable t) { |
| if (throwable != null) { |
| t.addSuppressed(throwable); |
| } |
| throwable = t; |
| } |
| try { |
| if (mListener != null) { |
| mListener.onFinishedEvaluating(); |
| } |
| } catch (Throwable t) { |
| if (throwable != null) { |
| t.addSuppressed(throwable); |
| } |
| throwable = t; |
| } |
| } |
| if (throwable != null) throw throwable; |
| } |
| }; |
| } |
| |
| private static void assertAnnotationsMatchParameterization( |
| AnnotationsRetriever.FlagAnnotations flagAnnotations, |
| FlagsParameterization parameterization) { |
| if (parameterization == null) return; |
| Set<String> parameterizedFlags = parameterization.mOverrides.keySet(); |
| Set<String> requiredFlags = flagAnnotations.mRequiredFlagValues.keySet(); |
| // Assert that NO Annotation-Required flag is in the parameterization |
| Set<String> parameterizedAndRequiredFlags = |
| Sets.intersection(parameterizedFlags, requiredFlags); |
| if (!parameterizedAndRequiredFlags.isEmpty()) { |
| throw new AssertionError( |
| "The following flags have required values (per @RequiresFlagsEnabled or" |
| + " @RequiresFlagsDisabled) but they are part of the" |
| + " FlagParameterization: " |
| + parameterizedAndRequiredFlags); |
| } |
| } |
| |
| private void setFlagValue(String fullFlagName, boolean value) { |
| if (!fullFlagName.contains(".")) { |
| throw new FlagSetException( |
| fullFlagName, "Flag name is not the expected format {packgeName}.{flagName}."); |
| } |
| // Get all packages containing Flags referencing the same fullFlagName. |
| Set<String> packageSet = getPackagesContainsFlag(fullFlagName); |
| |
| for (String packageName : packageSet) { |
| setFlagValue(Flag.createFlag(fullFlagName, packageName), value); |
| } |
| } |
| |
| private Set<String> getPackagesContainsFlag(String fullFlagName) { |
| return getAllPackagesForFlag(fullFlagName, mPackageToRepackage); |
| } |
| |
| private static Set<String> getAllPackagesForFlag( |
| String fullFlagName, Map<String, Set<String>> packageToRepackage) { |
| |
| String packageName = Flag.getFlagPackageName(fullFlagName); |
| Set<String> packageSet = packageToRepackage.getOrDefault(packageName, new HashSet<>()); |
| |
| if (!packageSet.isEmpty()) { |
| return packageSet; |
| } |
| |
| for (String prefix : REPACKAGE_PREFIX_LIST) { |
| String repackagedName = String.format("%s%s", prefix, packageName); |
| String flagClassName = String.format("%s.%s", repackagedName, FLAGS_CLASS_NAME); |
| try { |
| Class.forName(flagClassName, false, SetFlagsRule.class.getClassLoader()); |
| packageSet.add(repackagedName); |
| } catch (ClassNotFoundException e) { |
| // Skip if the class is not found |
| // An error will be thrown if no package containing flags referencing |
| // the passed in flag |
| } |
| } |
| packageToRepackage.put(packageName, packageSet); |
| if (packageSet.isEmpty()) { |
| throw new FlagSetException( |
| fullFlagName, |
| "Cannot find package containing Flags class referencing to this flag."); |
| } |
| return packageSet; |
| } |
| |
| private void setFlagValue(Flag flag, boolean value) { |
| if (mListener != null) { |
| mListener.onBeforeSetFlag(flag, value); |
| } |
| |
| Object fakeFlagsImplInstance = null; |
| |
| Class<?> flagsClass = getFlagClassFromFlag(flag); |
| fakeFlagsImplInstance = getOrCreateFakeFlagsImp(flagsClass); |
| |
| if (!mMutatedFlagsClasses.contains(flagsClass)) { |
| // Replace FeatureFlags in Flags class with FakeFeatureFlagsImpl |
| replaceFlagsImpl(flagsClass, fakeFlagsImplInstance); |
| mMutatedFlagsClasses.add(flagsClass); |
| } |
| |
| // If the test is trying to set the flag value on a read_only flag in an optimized build |
| // skip this test, since it is not a valid testing case |
| // The reason for skipping instead of throwning error here is all read_write flag will be |
| // change to read_only in the final release configuration. Thus the test could be executed |
| // in other release configuration cases |
| // TODO(b/337449119): SetFlagsRule should still run tests that are consistent with the |
| // read-only values of flags. But be careful, if a ClassRule exists, the value returned by |
| // the original FeatureFlags instance may be overridden, and reading it may not be allowed. |
| boolean isOptimized = verifyFlagReadOnlyAndOptimized(fakeFlagsImplInstance, flag); |
| assumeFalse( |
| String.format( |
| "Flag %s is read_only, and the code is optimized. " |
| + " The flag value should not be modified on this build" |
| + " Skip this test.", |
| flag.fullFlagName()), |
| isOptimized); |
| |
| // Set desired flag value in the FakeFeatureFlagsImpl |
| setFlagValueInFakeFeatureFlagsImpl(fakeFlagsImplInstance, flag, value); |
| } |
| |
| private static Class<?> getFlagClassFromFlag(Flag flag) { |
| String className = flag.flagsClassName(); |
| Class<?> flagsClass = null; |
| try { |
| flagsClass = Class.forName(className); |
| } catch (ClassNotFoundException e) { |
| throw new FlagSetException( |
| flag.fullFlagName(), |
| String.format( |
| "Can not load the Flags class %s to set its values. Please check the " |
| + "flag name and ensure that the aconfig auto generated " |
| + "library is in the dependency.", |
| className), |
| e); |
| } |
| return flagsClass; |
| } |
| |
| private static Class<?> getFlagClassFromFlagsClassName(String className) { |
| if (!className.endsWith("." + FLAGS_CLASS_NAME)) { |
| throw new FlagSetException( |
| className, |
| "Can not watch this Flags class because it is not named 'Flags'. Please ensure" |
| + " your @UsesFlags() annotations only reference the Flags classes."); |
| } |
| try { |
| return Class.forName(className); |
| } catch (ClassNotFoundException e) { |
| throw new FlagSetException( |
| className, |
| "Cannot load this Flags class to set its values. Please check the flag name and" |
| + " ensure that the aconfig auto generated library is in the dependency.", |
| e); |
| } |
| } |
| |
| private boolean getFlagValue(Object featureFlagsImpl, Flag flag) { |
| // Must be consistent with method name in aconfig auto generated code. |
| String methodName = getFlagMethodName(flag); |
| String fullFlagName = flag.fullFlagName(); |
| |
| try { |
| Object result = |
| featureFlagsImpl.getClass().getMethod(methodName).invoke(featureFlagsImpl); |
| if (result instanceof Boolean) { |
| return (Boolean) result; |
| } |
| throw new FlagReadException( |
| fullFlagName, |
| String.format( |
| "Flag type is %s, not boolean", result.getClass().getSimpleName())); |
| } catch (NoSuchMethodException e) { |
| throw new FlagReadException( |
| fullFlagName, |
| String.format( |
| "No method %s in the Flags class %s to read the flag value. Please" |
| + " check the flag name.", |
| methodName, featureFlagsImpl.getClass().getName()), |
| e); |
| } catch (ReflectiveOperationException e) { |
| throw new FlagReadException( |
| fullFlagName, |
| String.format( |
| "Fail to get value of flag %s from instance %s", |
| fullFlagName, featureFlagsImpl.getClass().getName()), |
| e); |
| } |
| } |
| |
| private String getFlagMethodName(Flag flag) { |
| return CaseFormat.LOWER_UNDERSCORE.to(CaseFormat.LOWER_CAMEL, flag.simpleFlagName()); |
| } |
| |
| private void setFlagValueInFakeFeatureFlagsImpl( |
| Object fakeFeatureFlagsImpl, Flag flag, boolean value) { |
| String fullFlagName = flag.fullFlagName(); |
| try { |
| fakeFeatureFlagsImpl |
| .getClass() |
| .getMethod(SET_FLAG_METHOD_NAME, String.class, boolean.class) |
| .invoke(fakeFeatureFlagsImpl, fullFlagName, value); |
| } catch (NoSuchMethodException e) { |
| throw new FlagSetException( |
| fullFlagName, |
| String.format( |
| "Flag implementation %s is not fake implementation", |
| fakeFeatureFlagsImpl.getClass().getName()), |
| e); |
| } catch (ReflectiveOperationException e) { |
| throw new FlagSetException(fullFlagName, e); |
| } |
| } |
| |
| private static boolean verifyFlagReadOnlyAndOptimized(Object fakeFeatureFlagsImpl, Flag flag) { |
| String fullFlagName = flag.fullFlagName(); |
| try { |
| boolean result = |
| (Boolean) |
| fakeFeatureFlagsImpl |
| .getClass() |
| .getMethod( |
| IS_FLAG_READ_ONLY_OPTIMIZED_METHOD_NAME, String.class) |
| .invoke(fakeFeatureFlagsImpl, fullFlagName); |
| return result; |
| } catch (NoSuchMethodException e) { |
| // If the flag is generated under exported mode, then it doesn't have this method |
| String simpleClassName = fakeFeatureFlagsImpl.getClass().getSimpleName(); |
| if (simpleClassName.equals(FAKE_FEATURE_FLAGS_IMPL_CLASS_NAME)) { |
| return false; |
| } |
| if (simpleClassName.equals(CUSTOM_FEATURE_FLAGS_CLASS_NAME)) { |
| return false; |
| } |
| throw new FlagSetException( |
| fullFlagName, |
| String.format( |
| "Cannot check whether flag is optimized. " |
| + "Flag implementation %s is not fake implementation", |
| fakeFeatureFlagsImpl.getClass().getName()), |
| e); |
| } catch (ReflectiveOperationException e) { |
| throw new FlagSetException(fullFlagName, e); |
| } |
| } |
| |
| @Nonnull |
| private Object getOrCreateFakeFlagsImp(Class<?> flagsClass) { |
| Object fakeFlagsImplInstance = mFlagsClassToFakeFlagsImpl.get(flagsClass); |
| if (fakeFlagsImplInstance != null) { |
| return fakeFlagsImplInstance; |
| } |
| |
| String packageName = flagsClass.getPackageName(); |
| String fakeClassName = |
| String.format("%s.%s", packageName, FAKE_FEATURE_FLAGS_IMPL_CLASS_NAME); |
| String interfaceName = String.format("%s.%s", packageName, FEATURE_FLAGS_CLASS_NAME); |
| |
| Object realFlagsImplInstance = readFlagsImpl(flagsClass); |
| mFlagsClassToRealFlagsImpl.put(flagsClass, realFlagsImplInstance); |
| |
| try { |
| Class<?> flagImplClass = Class.forName(fakeClassName); |
| Class<?> flagInterface = Class.forName(interfaceName); |
| fakeFlagsImplInstance = |
| flagImplClass |
| .getConstructor(flagInterface) |
| .newInstance(mIsInitWithDefault ? realFlagsImplInstance : null); |
| } catch (ReflectiveOperationException e) { |
| throw new UnsupportedOperationException( |
| String.format( |
| "Cannot create FakeFeatureFlagsImpl in Flags class %s.", |
| flagsClass.getName()), |
| e); |
| } |
| |
| mFlagsClassToFakeFlagsImpl.put(flagsClass, fakeFlagsImplInstance); |
| |
| return fakeFlagsImplInstance; |
| } |
| |
| private static void replaceFlagsImpl(Class<?> flagsClass, Object flagsImplInstance) { |
| Field featureFlagsField = getFeatureFlagsField(flagsClass); |
| try { |
| featureFlagsField.set(null, flagsImplInstance); |
| } catch (IllegalAccessException e) { |
| throw new UnsupportedOperationException( |
| String.format( |
| "Cannot replace FeatureFlagsImpl to %s.", |
| flagsImplInstance.getClass().getName()), |
| e); |
| } |
| } |
| |
| private static Object readFlagsImpl(Class<?> flagsClass) { |
| Field featureFlagsField = getFeatureFlagsField(flagsClass); |
| try { |
| return featureFlagsField.get(null); |
| } catch (IllegalAccessException e) { |
| throw new UnsupportedOperationException( |
| String.format( |
| "Cannot get FeatureFlags from Flags class %s.", flagsClass.getName()), |
| e); |
| } |
| } |
| |
| private static Field getFeatureFlagsField(Class<?> flagsClass) { |
| Field featureFlagsField = null; |
| try { |
| featureFlagsField = flagsClass.getDeclaredField(FEATURE_FLAGS_FIELD_NAME); |
| } catch (ReflectiveOperationException e) { |
| throw new UnsupportedOperationException( |
| String.format( |
| "Cannot store FeatureFlagsImpl in Flag %s.", flagsClass.getName()), |
| e); |
| } |
| featureFlagsField.setAccessible(true); |
| return featureFlagsField; |
| } |
| |
| private void resetFlags() { |
| String flagsClassName = null; |
| try { |
| for (Class<?> flagsClass : mMutatedFlagsClasses) { |
| flagsClassName = flagsClass.getName(); |
| Object fakeFlagsImplInstance = mFlagsClassToFakeFlagsImpl.get(flagsClass); |
| Object flagsImplInstance = mFlagsClassToRealFlagsImpl.get(flagsClass); |
| // Replace FeatureFlags in Flags class with real FeatureFlagsImpl |
| replaceFlagsImpl(flagsClass, flagsImplInstance); |
| fakeFlagsImplInstance |
| .getClass() |
| .getMethod(RESET_ALL_METHOD_NAME) |
| .invoke(fakeFlagsImplInstance); |
| } |
| mMutatedFlagsClasses.clear(); |
| } catch (Exception e) { |
| throw new FlagSetException(flagsClassName, e); |
| } |
| } |
| |
| /** An interface that provides hooks to the ClassRule. */ |
| private interface Listener { |
| /** Called before a flag is set. */ |
| void onBeforeSetFlag(SetFlagsRule.Flag flag, boolean value); |
| |
| /** Called after the rule has started evaluating for a test. */ |
| void onStartedEvaluating(); |
| |
| /** Called after the rule has finished evaluating for a test. */ |
| void onFinishedEvaluating(); |
| } |
| |
| /** |
| * A @ClassRule which adds extra consistency checks for SetFlagsRule. |
| * <li>Requires that tests monitor the Flags class of any flag that is set. |
| * <li>Fails a test if a flag that is set was read before the test started. |
| */ |
| public static class ClassRule implements TestRule { |
| /** The flags classes that are requested to be watched during construction. */ |
| private final Set<Class<?>> mGlobalFlagsClassesToWatch = new HashSet<>(); |
| |
| /** The flags packages that are allowed to be set, for quick per-flag lookup */ |
| private final Set<String> mSettableFlagsPackages = new HashSet<>(); |
| |
| /** The mapping from the Flags classes to the real implementations */ |
| private final Map<Class<?>, Object> mFlagsClassToRealFlagsImpl = new HashMap<>(); |
| |
| /** The mapping from the Flags classes to the watcher implementations */ |
| private final Map<Class<?>, Object> mFlagsClassToWatcherImpl = new HashMap<>(); |
| |
| /** The flags classes that have actually been mutated */ |
| private final Set<Class<?>> mMutatedFlagsClasses = new HashSet<>(); |
| |
| /** The flag values set by class annotations */ |
| private final Map<String, Boolean> mClassLevelSetFlagValues = new ConcurrentHashMap<>(); |
| |
| /** |
| * The individual flags which have been read from prior to tests starting, mapped to the |
| * stack trace of the first read. |
| */ |
| private final Map<String, FirstFlagRead> mFirstReadOutsideTestsByFlag = |
| new ConcurrentHashMap<>(); |
| |
| /** |
| * The individual flags which have been read from within a test, mapped to the stack trace |
| * of the first read. |
| */ |
| private final Map<String, FirstFlagRead> mFirstReadWithinTestByFlag = |
| new ConcurrentHashMap<>(); |
| |
| /** repackage cache */ |
| private final Map<String, Set<String>> mPackageToRepackage = new HashMap<>(); |
| |
| /** The depth of the ClassRule evaluating on potentially nested suites */ |
| private int mSuiteRunDepth = 0; |
| |
| /** Whether the SetFlagsRule is evaluating for a test */ |
| private boolean mIsTestRunning = false; |
| |
| /** Typical constructor takes an initial list flags classes to watch */ |
| public ClassRule(Class<?>... flagsClasses) { |
| for (Class<?> flagsClass : flagsClasses) { |
| mGlobalFlagsClassesToWatch.add(flagsClass); |
| } |
| } |
| |
| /** Listener to be notified of events in any created SetFlagsRule */ |
| private SetFlagsRule.Listener mListener = |
| new SetFlagsRule.Listener() { |
| @Override |
| public void onBeforeSetFlag(SetFlagsRule.Flag flag, boolean value) { |
| if (!mIsTestRunning) { |
| throw new IllegalStateException("Inner rule should be running!"); |
| } |
| assertFlagCanBeSet(flag, value); |
| } |
| |
| @Override |
| public void onStartedEvaluating() { |
| if (mSuiteRunDepth == 0) { |
| throw new IllegalStateException("Outer rule should be running!"); |
| } |
| if (mIsTestRunning) { |
| throw new IllegalStateException("Inner rule is still running!"); |
| } |
| mIsTestRunning = true; |
| } |
| |
| @Override |
| public void onFinishedEvaluating() { |
| if (!mIsTestRunning) { |
| throw new IllegalStateException("Inner rule did not start!"); |
| } |
| mIsTestRunning = false; |
| checkAllFlagsWatchersRestored(); |
| mFirstReadWithinTestByFlag.clear(); |
| } |
| }; |
| |
| /** |
| * Creates a SetFlagsRule which will work as normal, but additionally enforce the guarantees |
| * about not setting flags that were read within the ClassRule |
| */ |
| public SetFlagsRule createSetFlagsRule() { |
| return createSetFlagsRule(null); |
| } |
| |
| /** |
| * Creates a SetFlagsRule with parameterization which will work as normal, but additionally |
| * enforce the guarantees about not setting flags that were read within the ClassRule |
| */ |
| public SetFlagsRule createSetFlagsRule( |
| @Nullable FlagsParameterization flagsParameterization) { |
| return new SetFlagsRule( |
| DefaultInitValueType.DEVICE_DEFAULT, flagsParameterization, mListener); |
| } |
| |
| private boolean isFlagsClassMonitored(SetFlagsRule.Flag flag) { |
| return mSettableFlagsPackages.contains(flag.flagPackageName()); |
| } |
| |
| private void assertFlagCanBeSet(SetFlagsRule.Flag flag, boolean value) { |
| Exception firstReadWithinTest = mFirstReadWithinTestByFlag.get(flag.fullFlagName()); |
| if (firstReadWithinTest != null) { |
| throw new FlagSetException( |
| flag.fullFlagName(), |
| "This flag was locked when it was read earlier in this test. To fix this" |
| + " error, always use @EnableFlags() and @DisableFlags() to set" |
| + " flags, which ensures flags are set before even any" |
| + " @Before-annotated setup methods.", |
| firstReadWithinTest); |
| } |
| Exception firstReadOutsideTest = mFirstReadOutsideTestsByFlag.get(flag.fullFlagName()); |
| if (firstReadOutsideTest != null) { |
| throw new FlagSetException( |
| flag.fullFlagName(), |
| "This flag was locked when it was read outside of the test code; likely" |
| + " during initialization of the test class. To fix this error," |
| + " move test fixture initialization code into your" |
| + " @Before-annotated setup method, and ensure you are using" |
| + " @EnableFlags() and @DisableFlags() to set flags.", |
| firstReadOutsideTest); |
| } |
| if (!isFlagsClassMonitored(flag)) { |
| throw new FlagSetException( |
| flag.fullFlagName(), |
| "This flag's class is not monitored. Always use @EnableFlags() and" |
| + " @DisableFlags() on the class or method instead of" |
| + " .enableFlags() or .disableFlags() to prevent this error. When" |
| + " using FlagsParameterization, add `@UsesFlags(" |
| + flag.flagPackageName() |
| + ".Flags.class)` to the test class. As a last resort, pass the" |
| + " Flags class to the constructor of your" |
| + " SetFlagsRule.ClassRule."); |
| } |
| // Detect errors where the rule messed up and set the wrong flag value. |
| Boolean classLevelValue = mClassLevelSetFlagValues.get(flag.fullFlagName()); |
| if (classLevelValue != null && classLevelValue != value) { |
| throw new FlagSetException( |
| flag.fullFlagName(), |
| "This flag's value was set at the class level to a different value."); |
| } |
| } |
| |
| private void checkInstanceOfRealFlagsImpl(Object actual) { |
| if (!actual.getClass().getSimpleName().equals(REAL_FEATURE_FLAGS_IMPL_CLASS_NAME)) { |
| throw new IllegalStateException( |
| String.format( |
| "Wrong impl type during setup: '%s' is not a %s", |
| actual, REAL_FEATURE_FLAGS_IMPL_CLASS_NAME)); |
| } |
| } |
| |
| private void checkSameAs(Object expected, Object actual) { |
| if (expected != actual) { |
| throw new IllegalStateException( |
| String.format( |
| "Wrong impl instance found during teardown: expected %s but was %s", |
| expected, actual)); |
| } |
| } |
| |
| private Object getOrCreateFlagReadWatcher(Class<?> flagsClass) { |
| Object watcher = mFlagsClassToWatcherImpl.get(flagsClass); |
| if (watcher != null) { |
| return watcher; |
| } |
| Object flagsImplInstance = readFlagsImpl(flagsClass); |
| // strict mode: ensure that the current impl is the real impl |
| checkInstanceOfRealFlagsImpl(flagsImplInstance); |
| // save the real impl for restoration later |
| mFlagsClassToRealFlagsImpl.put(flagsClass, flagsImplInstance); |
| watcher = newFlagReadWatcher(flagsClass, flagsImplInstance); |
| mFlagsClassToWatcherImpl.put(flagsClass, watcher); |
| return watcher; |
| } |
| |
| private void recordFlagRead(String flagName) { |
| if (mIsTestRunning) { |
| mFirstReadWithinTestByFlag.computeIfAbsent(flagName, FirstFlagRead::new); |
| } else { |
| mFirstReadOutsideTestsByFlag.computeIfAbsent(flagName, FirstFlagRead::new); |
| } |
| } |
| |
| private Object newFlagReadWatcher(Class<?> flagsClass, Object flagsImplInstance) { |
| String packageName = flagsClass.getPackageName(); |
| String customClassName = |
| String.format("%s.%s", packageName, CUSTOM_FEATURE_FLAGS_CLASS_NAME); |
| BiPredicate<String, Predicate<Object>> getValueImpl = |
| (flagName, predicate) -> { |
| // Flags set at the class level pose no consistency risk |
| Boolean value = mClassLevelSetFlagValues.get(flagName); |
| if (value != null) { |
| return value; |
| } |
| recordFlagRead(flagName); |
| return predicate.test(flagsImplInstance); |
| }; |
| try { |
| Class<?> customFlagsClass = Class.forName(customClassName); |
| return customFlagsClass.getConstructor(BiPredicate.class).newInstance(getValueImpl); |
| } catch (ReflectiveOperationException e) { |
| throw new UnsupportedOperationException( |
| String.format( |
| "Cannot create CustomFeatureFlags in Flags class %s.", |
| flagsClass.getName()), |
| e); |
| } |
| } |
| |
| /** Get the package name of the flags in this class. This is the non-repackaged name. */ |
| private String getFlagPackageName(Class<?> flagsClass) { |
| String classPackageName = flagsClass.getPackageName(); |
| String shortestPackageName = classPackageName; |
| for (String prefix : REPACKAGE_PREFIX_LIST) { |
| if (prefix.isEmpty()) continue; |
| if (classPackageName.startsWith(prefix)) { |
| String unprefixedPackage = classPackageName.substring(prefix.length()); |
| if (unprefixedPackage.length() < shortestPackageName.length()) { |
| shortestPackageName = unprefixedPackage; |
| } |
| } |
| } |
| return shortestPackageName; |
| } |
| |
| private void setupClassLevelFlagValues(Description description) { |
| mClassLevelSetFlagValues.putAll( |
| AnnotationsRetriever.getFlagAnnotations(description).mSetFlagValues); |
| } |
| |
| private void setupFlagsWatchers(Description description) { |
| // Start with the static list of Flags classes to watch |
| Set<Class<?>> flagsClassesToWatch = new HashSet<>(mGlobalFlagsClassesToWatch); |
| // Collect the Flags classes from @UsedFlags annotations on the Descriptor |
| Set<String> usedFlagsClasses = AnnotationsRetriever.getAllUsedFlagsClasses(description); |
| for (String flagsClassName : usedFlagsClasses) { |
| flagsClassesToWatch.add(getFlagClassFromFlagsClassName(flagsClassName)); |
| } |
| // Now setup watchers on the provided Flags classes |
| for (Class<?> flagsClass : flagsClassesToWatch) { |
| setupFlagsWatcher(flagsClass, getFlagPackageName(flagsClass)); |
| } |
| // Get all annotated flags and then the distinct packages for each flag |
| Set<String> setFlags = AnnotationsRetriever.getAllAnnotationSetFlags(description); |
| Set<String> extraFlagPackages = new HashSet<>(); |
| for (String setFlag : setFlags) { |
| extraFlagPackages.add(Flag.getFlagPackageName(setFlag)); |
| } |
| // Do not bother with flags that are already monitored |
| extraFlagPackages.removeAll(mSettableFlagsPackages); |
| // Expand packages to all repackaged versions, stored as Flag objects |
| Set<Flag> extraWildcardFlags = new HashSet<>(); |
| for (String extraFlagPackage : extraFlagPackages) { |
| String fullFlagName = extraFlagPackage + ".*"; |
| Set<String> packages = getAllPackagesForFlag(fullFlagName, mPackageToRepackage); |
| for (String packageName : packages) { |
| Flag flag = Flag.createFlag(fullFlagName, packageName); |
| extraWildcardFlags.add(flag); |
| } |
| } |
| // Set up watchers for each wildcard flag |
| for (Flag flag : extraWildcardFlags) { |
| Class<?> flagsClass = getFlagClassFromFlag(flag); |
| setupFlagsWatcher(flagsClass, flag.flagPackageName()); |
| } |
| } |
| |
| private void setupFlagsWatcher(Class<?> flagsClass, String flagPackageName) { |
| if (mMutatedFlagsClasses.contains(flagsClass)) { |
| throw new IllegalStateException( |
| String.format("Flags class %s is already mutated", flagsClass.getName())); |
| } |
| Object watcher = getOrCreateFlagReadWatcher(flagsClass); |
| replaceFlagsImpl(flagsClass, watcher); |
| mMutatedFlagsClasses.add(flagsClass); |
| mSettableFlagsPackages.add(flagPackageName); |
| } |
| |
| private void teardownFlagsWatchers() { |
| try { |
| for (Class<?> flagsClass : mMutatedFlagsClasses) { |
| Object flagsImplInstance = mFlagsClassToRealFlagsImpl.get(flagsClass); |
| // strict mode: ensure that the watcher is still in place |
| Object watcher = readFlagsImpl(flagsClass); |
| checkSameAs(mFlagsClassToWatcherImpl.get(flagsClass), watcher); |
| // Replace FeatureFlags in Flags class with real FeatureFlagsImpl |
| replaceFlagsImpl(flagsClass, flagsImplInstance); |
| } |
| mMutatedFlagsClasses.clear(); |
| mSettableFlagsPackages.clear(); |
| mFirstReadOutsideTestsByFlag.clear(); |
| } catch (IllegalStateException e) { |
| throw e; |
| } catch (Exception e) { |
| throw new IllegalStateException("Failed to teardown Flags watchers", e); |
| } |
| if (mIsTestRunning) { |
| throw new IllegalStateException("An inner SetFlagsRule is still running"); |
| } |
| if (!mFirstReadWithinTestByFlag.isEmpty()) { |
| throw new IllegalStateException("An inner SetFlagsRule did not fully clean up"); |
| } |
| } |
| |
| private void checkAllFlagsWatchersRestored() { |
| for (Class<?> flagsClass : mMutatedFlagsClasses) { |
| Object watcher = readFlagsImpl(flagsClass); |
| checkSameAs(mFlagsClassToWatcherImpl.get(flagsClass), watcher); |
| } |
| } |
| |
| @Override |
| public Statement apply(Statement base, Description description) { |
| return new Statement() { |
| @Override |
| public void evaluate() throws Throwable { |
| Throwable throwable = null; |
| final int initialDepth = mSuiteRunDepth; |
| try { |
| mSuiteRunDepth++; |
| if (initialDepth == 0) { |
| setupFlagsWatchers(description); |
| setupClassLevelFlagValues(description); |
| } |
| base.evaluate(); |
| } catch (Throwable t) { |
| throwable = t; |
| } finally { |
| mSuiteRunDepth--; |
| try { |
| if (initialDepth == 0) { |
| mClassLevelSetFlagValues.clear(); |
| teardownFlagsWatchers(); |
| } |
| if (mSuiteRunDepth != initialDepth) { |
| throw new IllegalStateException( |
| String.format( |
| "Evaluations were not correctly nested: initial" |
| + " depth was %d but final depth was %d", |
| initialDepth, mSuiteRunDepth)); |
| } |
| } catch (Throwable t) { |
| if (throwable != null) { |
| t.addSuppressed(throwable); |
| } |
| throwable = t; |
| } |
| } |
| if (throwable != null) throw throwable; |
| } |
| }; |
| } |
| } |
| |
| private static class FirstFlagRead extends Exception { |
| FirstFlagRead(String flagName) { |
| super(String.format("Flag '%s' was first read at this location:", flagName)); |
| } |
| } |
| |
| private static class Flag { |
| private static final String PACKAGE_NAME_SIMPLE_NAME_SEPARATOR = "."; |
| private final String mFullFlagName; |
| private final String mFlagPackageName; |
| private final String mClassPackageName; |
| private final String mSimpleFlagName; |
| |
| public static String getFlagPackageName(String fullFlagName) { |
| int index = fullFlagName.lastIndexOf(PACKAGE_NAME_SIMPLE_NAME_SEPARATOR); |
| return fullFlagName.substring(0, index); |
| } |
| |
| public static Flag createFlag(String fullFlagName) { |
| int index = fullFlagName.lastIndexOf(PACKAGE_NAME_SIMPLE_NAME_SEPARATOR); |
| String packageName = fullFlagName.substring(0, index); |
| return createFlag(fullFlagName, packageName); |
| } |
| |
| public static Flag createFlag(String fullFlagName, String classPackageName) { |
| if (!fullFlagName.contains(PACKAGE_NAME_SIMPLE_NAME_SEPARATOR) |
| || !classPackageName.contains(PACKAGE_NAME_SIMPLE_NAME_SEPARATOR)) { |
| throw new IllegalArgumentException( |
| String.format( |
| "Flag %s is invalid. The format should be {packageName}" |
| + ".{simpleFlagName}", |
| fullFlagName)); |
| } |
| int index = fullFlagName.lastIndexOf(PACKAGE_NAME_SIMPLE_NAME_SEPARATOR); |
| String flagPackageName = fullFlagName.substring(0, index); |
| String simpleFlagName = fullFlagName.substring(index + 1); |
| |
| return new Flag(fullFlagName, flagPackageName, classPackageName, simpleFlagName); |
| } |
| |
| private Flag( |
| String fullFlagName, |
| String flagPackageName, |
| String classPackageName, |
| String simpleFlagName) { |
| this.mFullFlagName = fullFlagName; |
| this.mFlagPackageName = flagPackageName; |
| this.mClassPackageName = classPackageName; |
| this.mSimpleFlagName = simpleFlagName; |
| } |
| |
| public String fullFlagName() { |
| return mFullFlagName; |
| } |
| |
| public String flagPackageName() { |
| return mFlagPackageName; |
| } |
| |
| public String classPackageName() { |
| return mClassPackageName; |
| } |
| |
| public String simpleFlagName() { |
| return mSimpleFlagName; |
| } |
| |
| public String flagsClassName() { |
| return String.format("%s.%s", classPackageName(), FLAGS_CLASS_NAME); |
| } |
| } |
| } |
