1 // Copyright 2015 The Chromium Authors 2 // Use of this source code is governed by a BSD-style license that can be 3 // found in the LICENSE file. 4 5 package org.chromium.base.test.util; 6 7 import android.app.Activity; 8 import android.text.TextUtils; 9 10 import org.junit.Assert; 11 import org.junit.Rule; 12 import org.junit.rules.TestRule; 13 import org.junit.runner.Description; 14 import org.junit.runners.model.Statement; 15 16 import org.chromium.base.ActivityState; 17 import org.chromium.base.ApplicationStatus; 18 import org.chromium.base.CommandLine; 19 import org.chromium.base.CommandLineInitUtil; 20 import org.chromium.base.Log; 21 import org.chromium.base.test.BaseJUnit4ClassRunner.ClassHook; 22 import org.chromium.base.test.BaseJUnit4ClassRunner.TestHook; 23 24 import java.lang.annotation.ElementType; 25 import java.lang.annotation.Inherited; 26 import java.lang.annotation.Retention; 27 import java.lang.annotation.RetentionPolicy; 28 import java.lang.annotation.Target; 29 import java.lang.reflect.Field; 30 import java.lang.reflect.Method; 31 import java.util.ArrayList; 32 import java.util.Arrays; 33 import java.util.Collections; 34 import java.util.HashMap; 35 import java.util.HashSet; 36 import java.util.List; 37 import java.util.Map; 38 import java.util.Map.Entry; 39 import java.util.Set; 40 41 /** 42 * Provides annotations related to command-line flag handling. 43 * 44 * <p>This can be used in either an on-device instrumentation test or a junit (robolectric) test 45 * running on the host. To use in an instrumentation test, just {@code RunWith} {@link 46 * BaseJUnit4ClassRunner} (or a runner which extends that class). To use from a robolectric test, 47 * add the following test rule to your class: 48 * 49 * <pre> 50 * @Rule 51 * TestRule mRule = CommandLineFlags.getTestRule(); 52 * </pre> 53 * 54 * <p>Then you can annotate the test class, test methods, or test rules with {@code 55 * CommandLineFlags.Add} or {@code CommandLineFlags.Remove}. Uses of these annotations on a derived 56 * class will take precedence over uses on its base classes, so a derived class can add a 57 * command-line flag that a base class has removed (or vice versa). Similarly, uses of these 58 * annotations on a test method will take precedence over uses on the containing class. 59 * 60 * <p> 61 * These annonations may also be used on Junit4 Rule classes and on their base classes. Note, 62 * however that the annotation processor only looks at the declared type of the Rule, not its actual 63 * type, so in, for example: 64 * 65 * <pre> 66 * @Rule 67 * TestRule mRule = new ChromeActivityTestRule(); 68 * </pre> 69 * 70 * will only look for CommandLineFlags annotations on TestRule, not for CommandLineFlags annotations 71 * on ChromeActivityTestRule. 72 * <p> 73 * In addition a rule may not remove flags added by an independently invoked rule, although it may 74 * remove flags added by its base classes. 75 * <p> 76 * Uses of these annotations on the test class or methods take precedence over uses on Rule classes. 77 * <p> 78 * Note that this class should never be instantiated. 79 */ 80 public final class CommandLineFlags { 81 private static final String TAG = "CommandLineFlags"; 82 private static final String DISABLE_FEATURES = "disable-features"; 83 private static final String ENABLE_FEATURES = "enable-features"; 84 85 // These members are used to track CommandLine state modifications made by the class/test method 86 // currently being run, to be undone when the class/test method finishes. 87 private static Set<String> sClassFlagsToRemove; 88 private static Map<String, String> sClassFlagsToAdd; 89 private static Set<String> sMethodFlagsToRemove; 90 private static Map<String, String> sMethodFlagsToAdd; 91 92 /** Adds command-line flags to the {@link org.chromium.base.CommandLine} for this test. */ 93 @Inherited 94 @Retention(RetentionPolicy.RUNTIME) 95 @Target({ElementType.METHOD, ElementType.TYPE}) 96 public @interface Add { value()97 String[] value(); 98 } 99 100 /** 101 * Removes command-line flags from the {@link org.chromium.base.CommandLine} from this test. 102 * 103 * Note that this can only be applied to test methods. This restriction is due to complexities 104 * in resolving the order that annotations are applied, and given how rare it is to need to 105 * remove command line flags, this annotation must be applied directly to each test method 106 * wishing to remove a flag. 107 */ 108 @Inherited 109 @Retention(RetentionPolicy.RUNTIME) 110 @Target({ElementType.METHOD}) 111 public @interface Remove { value()112 String[] value(); 113 } 114 115 /** 116 * Sets up the CommandLine with the appropriate flags. 117 * 118 * This will add the difference of the sets of flags specified by {@link CommandLineFlags.Add} 119 * and {@link CommandLineFlags.Remove} to the {@link org.chromium.base.CommandLine}. Note that 120 * trying to remove a flag set externally, i.e. by the command-line flags file, will not work. 121 */ setUpClass(Class<?> clazz)122 public static void setUpClass(Class<?> clazz) { 123 if (!CommandLine.isInitialized()) { 124 CommandLineInitUtil.initCommandLine(getTestCmdLineFile()); 125 } 126 127 Set<String> flags = new HashSet<>(); 128 updateFlagsForClass(clazz, flags); 129 sClassFlagsToRemove = new HashSet<>(); 130 sClassFlagsToAdd = new HashMap<>(); 131 applyFlags(flags, null, sClassFlagsToRemove, sClassFlagsToAdd); 132 } 133 tearDownClass()134 public static void tearDownClass() { 135 if (ApplicationStatus.isInitialized()) { 136 for (Activity a : ApplicationStatus.getRunningActivities()) { 137 if (ApplicationStatus.getStateForActivity(a) < ActivityState.RESUMED) { 138 Log.w( 139 TAG, 140 "Activity " 141 + a 142 + ", is still starting up while the Command Line flags " 143 + "are being reset. This is a known source of flakiness."); 144 } 145 } 146 } 147 restoreFlags(sClassFlagsToRemove, sClassFlagsToAdd); 148 sClassFlagsToRemove = null; 149 sClassFlagsToAdd = null; 150 } 151 setUpMethod(Method method)152 public static void setUpMethod(Method method) { 153 Set<String> flagsToAdd = new HashSet<>(); 154 Set<String> flagsToRemove = new HashSet<>(); 155 updateFlagsForMethod(method, flagsToAdd, flagsToRemove); 156 sMethodFlagsToRemove = new HashSet<>(); 157 sMethodFlagsToAdd = new HashMap<>(); 158 applyFlags(flagsToAdd, flagsToRemove, sMethodFlagsToRemove, sMethodFlagsToAdd); 159 } 160 tearDownMethod()161 public static void tearDownMethod() { 162 restoreFlags(sMethodFlagsToRemove, sMethodFlagsToAdd); 163 sMethodFlagsToRemove = null; 164 sMethodFlagsToAdd = null; 165 } 166 restoreFlags(Set<String> flagsToRemove, Map<String, String> flagsToAdd)167 private static void restoreFlags(Set<String> flagsToRemove, Map<String, String> flagsToAdd) { 168 for (String flag : flagsToRemove) { 169 CommandLine.getInstance().removeSwitch(flag); 170 } 171 for (Entry<String, String> flag : flagsToAdd.entrySet()) { 172 if (flag.getValue() == null) { 173 CommandLine.getInstance().appendSwitch(flag.getKey()); 174 } else { 175 CommandLine.getInstance().appendSwitchWithValue(flag.getKey(), flag.getValue()); 176 } 177 } 178 } 179 applyFlags( Set<String> flagsToAdd, Set<String> flagsToRemove, Set<String> flagsToRemoveForRestore, Map<String, String> flagsToAddForRestore)180 private static void applyFlags( 181 Set<String> flagsToAdd, 182 Set<String> flagsToRemove, 183 Set<String> flagsToRemoveForRestore, 184 Map<String, String> flagsToAddForRestore) { 185 if (flagsToRemove != null) { 186 for (String flag : flagsToRemove) { 187 if (CommandLine.getInstance().hasSwitch(flag)) { 188 String existingValue = CommandLine.getInstance().getSwitchValue(flag); 189 CommandLine.getInstance().removeSwitch(flag); 190 flagsToAddForRestore.put(flag, existingValue); 191 } 192 } 193 } 194 195 Set<String> enableFeatures = new HashSet<String>(getFeatureValues(ENABLE_FEATURES)); 196 Set<String> disableFeatures = new HashSet<String>(getFeatureValues(DISABLE_FEATURES)); 197 for (String flag : flagsToAdd) { 198 String[] parsedFlags = flag.split("=", 2); 199 if (parsedFlags.length == 1) { 200 if (!CommandLine.getInstance().hasSwitch(flag)) { 201 CommandLine.getInstance().appendSwitch(flag); 202 flagsToRemoveForRestore.add(flag); 203 } 204 } else if (ENABLE_FEATURES.equals(parsedFlags[0])) { 205 // We collect enable/disable features flags separately and aggregate them because 206 // they may be specified multiple times, in which case the values will trample each 207 // other. 208 Collections.addAll(enableFeatures, parsedFlags[1].split(",")); 209 } else if (DISABLE_FEATURES.equals(parsedFlags[0])) { 210 Collections.addAll(disableFeatures, parsedFlags[1].split(",")); 211 } else { 212 String existingValue = CommandLine.getInstance().getSwitchValue(parsedFlags[0]); 213 if (parsedFlags[1].equals(existingValue)) continue; 214 if (existingValue != null) { 215 flagsToAddForRestore.put(parsedFlags[0], existingValue); 216 CommandLine.getInstance().removeSwitch(parsedFlags[0]); 217 } 218 CommandLine.getInstance().appendSwitchWithValue(parsedFlags[0], parsedFlags[1]); 219 flagsToRemoveForRestore.add(parsedFlags[0]); 220 } 221 } 222 223 if (enableFeatures.size() > 0) { 224 String existingValue = CommandLine.getInstance().getSwitchValue(ENABLE_FEATURES); 225 if (existingValue != null) { 226 flagsToAddForRestore.put(ENABLE_FEATURES, existingValue); 227 CommandLine.getInstance().removeSwitch(ENABLE_FEATURES); 228 } 229 CommandLine.getInstance() 230 .appendSwitchWithValue(ENABLE_FEATURES, TextUtils.join(",", enableFeatures)); 231 flagsToRemoveForRestore.add(ENABLE_FEATURES); 232 } 233 if (disableFeatures.size() > 0) { 234 String existingValue = CommandLine.getInstance().getSwitchValue(DISABLE_FEATURES); 235 if (existingValue != null) { 236 flagsToAddForRestore.put(DISABLE_FEATURES, existingValue); 237 CommandLine.getInstance().removeSwitch(DISABLE_FEATURES); 238 } 239 CommandLine.getInstance() 240 .appendSwitchWithValue(DISABLE_FEATURES, TextUtils.join(",", disableFeatures)); 241 flagsToRemoveForRestore.add(DISABLE_FEATURES); 242 } 243 } 244 updateFlagsForClass(Class<?> clazz, Set<String> flags)245 private static void updateFlagsForClass(Class<?> clazz, Set<String> flags) { 246 // Get flags from rules within the class. 247 for (Field field : clazz.getFields()) { 248 if (field.isAnnotationPresent(Rule.class)) { 249 // The order in which fields are returned is undefined, so, for consistency, 250 // a rule must only ever add flags. 251 updateFlagsForClass(field.getType(), flags); 252 } 253 } 254 for (Method method : clazz.getMethods()) { 255 Assert.assertFalse( 256 "@Rule annotations on methods are unsupported. Cause: " 257 + method.toGenericString(), 258 method.isAnnotationPresent(Rule.class)); 259 } 260 261 // Add the flags from the parent. Override any flags defined by the rules. 262 Class<?> parent = clazz.getSuperclass(); 263 if (parent != null) updateFlagsForClass(parent, flags); 264 265 // Flags on the element itself override all other flag sources. 266 if (clazz.isAnnotationPresent(CommandLineFlags.Add.class)) { 267 flags.addAll(Arrays.asList(clazz.getAnnotation(CommandLineFlags.Add.class).value())); 268 } 269 } 270 updateFlagsForMethod( Method method, Set<String> flagsToAdd, Set<String> flagsToRemove)271 private static void updateFlagsForMethod( 272 Method method, Set<String> flagsToAdd, Set<String> flagsToRemove) { 273 if (method.isAnnotationPresent(CommandLineFlags.Add.class)) { 274 flagsToAdd.addAll( 275 Arrays.asList(method.getAnnotation(CommandLineFlags.Add.class).value())); 276 } 277 if (method.isAnnotationPresent(CommandLineFlags.Remove.class)) { 278 flagsToRemove.addAll( 279 Arrays.asList(method.getAnnotation(CommandLineFlags.Remove.class).value())); 280 } 281 } 282 getFeatureValues(String flag)283 private static List<String> getFeatureValues(String flag) { 284 String value = CommandLine.getInstance().getSwitchValue(flag); 285 if (value == null) return new ArrayList<>(); 286 return Arrays.asList(value.split(",")); 287 } 288 CommandLineFlags()289 private CommandLineFlags() { 290 throw new AssertionError("CommandLineFlags is a non-instantiable class"); 291 } 292 293 private static class CommandLineFlagsTestRule implements TestRule { 294 @Override apply(final Statement base, Description description)295 public Statement apply(final Statement base, Description description) { 296 return new Statement() { 297 @Override 298 public void evaluate() throws Throwable { 299 try { 300 Class clazz = description.getTestClass(); 301 CommandLineFlags.setUpClass(clazz); 302 CommandLineFlags.setUpMethod(clazz.getMethod(description.getMethodName())); 303 304 base.evaluate(); 305 } finally { 306 CommandLineFlags.tearDownMethod(); 307 CommandLineFlags.tearDownClass(); 308 } 309 } 310 }; 311 } 312 } 313 314 public static TestRule getTestRule() { 315 return new CommandLineFlagsTestRule(); 316 } 317 318 public static TestHook getPreTestHook() { 319 return (targetContext, testMethod) -> CommandLineFlags.setUpMethod(testMethod.getMethod()); 320 } 321 322 public static ClassHook getPreClassHook() { 323 return (targetContext, testClass) -> CommandLineFlags.setUpClass(testClass); 324 } 325 326 public static TestHook getPostTestHook() { 327 return (targetContext, testMethod) -> CommandLineFlags.tearDownMethod(); 328 } 329 330 public static ClassHook getPostClassHook() { 331 return (targetContext, testClass) -> CommandLineFlags.tearDownClass(); 332 } 333 334 public static String getTestCmdLineFile() { 335 return "test-cmdline-file"; 336 } 337 } 338