1 /*
2  * Copyright (C) 2023 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 package com.android.adservices.shared.testing;
17 
18 import com.android.adservices.shared.testing.Logger.LogLevel;
19 import com.android.adservices.shared.testing.Logger.RealLogger;
20 import com.android.adservices.shared.testing.NameValuePair.Matcher;
21 import com.android.adservices.shared.testing.annotations.DisableDebugFlag;
22 import com.android.adservices.shared.testing.annotations.DisableDebugFlags;
23 import com.android.adservices.shared.testing.annotations.EnableDebugFlag;
24 import com.android.adservices.shared.testing.annotations.EnableDebugFlags;
25 import com.android.adservices.shared.testing.annotations.SetDoubleFlag;
26 import com.android.adservices.shared.testing.annotations.SetDoubleFlags;
27 import com.android.adservices.shared.testing.annotations.SetFlagDisabled;
28 import com.android.adservices.shared.testing.annotations.SetFlagEnabled;
29 import com.android.adservices.shared.testing.annotations.SetFlagFalse;
30 import com.android.adservices.shared.testing.annotations.SetFlagTrue;
31 import com.android.adservices.shared.testing.annotations.SetFlagsDisabled;
32 import com.android.adservices.shared.testing.annotations.SetFlagsEnabled;
33 import com.android.adservices.shared.testing.annotations.SetFlagsFalse;
34 import com.android.adservices.shared.testing.annotations.SetFlagsTrue;
35 import com.android.adservices.shared.testing.annotations.SetFloatFlag;
36 import com.android.adservices.shared.testing.annotations.SetFloatFlags;
37 import com.android.adservices.shared.testing.annotations.SetIntegerFlag;
38 import com.android.adservices.shared.testing.annotations.SetIntegerFlags;
39 import com.android.adservices.shared.testing.annotations.SetLogcatTag;
40 import com.android.adservices.shared.testing.annotations.SetLogcatTags;
41 import com.android.adservices.shared.testing.annotations.SetLongDebugFlag;
42 import com.android.adservices.shared.testing.annotations.SetLongDebugFlags;
43 import com.android.adservices.shared.testing.annotations.SetLongFlag;
44 import com.android.adservices.shared.testing.annotations.SetLongFlags;
45 import com.android.adservices.shared.testing.annotations.SetStringArrayFlag;
46 import com.android.adservices.shared.testing.annotations.SetStringArrayFlags;
47 import com.android.adservices.shared.testing.annotations.SetStringFlag;
48 import com.android.adservices.shared.testing.annotations.SetStringFlags;
49 import com.android.adservices.shared.testing.concurrency.SyncCallback;
50 import com.android.adservices.shared.testing.device.DeviceConfig;
51 
52 import com.google.errorprone.annotations.FormatMethod;
53 import com.google.errorprone.annotations.FormatString;
54 
55 import org.junit.runner.Description;
56 import org.junit.runners.model.Statement;
57 
58 import java.lang.annotation.Annotation;
59 import java.util.ArrayList;
60 import java.util.Arrays;
61 import java.util.LinkedHashSet;
62 import java.util.List;
63 import java.util.Objects;
64 import java.util.Set;
65 import java.util.function.Consumer;
66 
67 // TODO(b/294423183): add unit tests for the most relevant / less repetitive stuff (don't need to
68 // test all setters / getters, for example)
69 /**
70  * Rule used to properly set "Android flags"- it will take care of permissions, restoring values at
71  * the end, setting {@link android.provider.DeviceConfig} or {@link android.os.SystemProperties},
72  * etc...
73  *
74  * @param <T> type of the concrete rule
75  */
76 public abstract class AbstractFlagsSetterRule<T extends AbstractFlagsSetterRule<T>>
77         extends AbstractRethrowerRule {
78 
79     protected static final String SYSTEM_PROPERTY_FOR_LOGCAT_TAGS_PREFIX = "log.tag.";
80 
81     // TODO(b/331781012): add all of them
82     // TODO(b/331781012): include itself (mLog.getTag()), although it would require a new method
83     // to make sure the tag is set before logging (as the initial command to set the logcat tag is
84     // cached)
85     private static final String[] INFRA_TAGS = {SyncCallback.LOG_TAG};
86     private static final Matcher INFRA_LOGTAG_MATCHER =
87             (prop) ->
88                     Arrays.stream(INFRA_TAGS)
89                             .anyMatch(
90                                     tag ->
91                                             prop.name.equals(
92                                                     SYSTEM_PROPERTY_FOR_LOGCAT_TAGS_PREFIX + tag));
93 
94     private final String mDeviceConfigNamespace;
95     private final DeviceConfigHelper mDeviceConfig;
96     // TODO(b/338067482): move system properties to its own rule?
97     // Prefix used on SystemProperties used for DebugFlags
98     private final String mDebugFlagPrefix;
99     private final SystemPropertiesHelper mSystemProperties;
100 
101     // Cache methods that were called before the test started, so the rule can be
102     // instantiated using a builder-like approach.
103     // NOTE: they MUST be cached and only executed after the test starts, because there's no
104     // guarantee that the rule will be executed at all (for example, a rule executed earlier might
105     // throw an AssumptionViolatedException)
106     private final List<Command> mInitialCommands = new ArrayList<>();
107 
108     // Name of flags that were changed by the test
109     private final Set<String> mChangedFlags = new LinkedHashSet<>();
110     // Name of system properties that were changed by the test
111     private final Set<String> mChangedSystemProperties = new LinkedHashSet<>();
112     private final Matcher mSystemPropertiesMatcher;
113 
114     private final List<NameValuePair> mPreTestFlags = new ArrayList<>();
115     private final List<NameValuePair> mPreTestSystemProperties = new ArrayList<>();
116     private final List<NameValuePair> mOnTestFailureFlags = new ArrayList<>();
117     private final List<NameValuePair> mOnTestFailureSystemProperties = new ArrayList<>();
118 
119     private DeviceConfig.SyncDisabledModeForTest mPreviousSyncDisabledModeForTest;
120 
121     private boolean mIsRunning;
122     private boolean mFlagsClearedByTest;
123 
124     private final Consumer<NameValuePair> mFlagsSetter;
125 
126     private final boolean mSkipStuffWhenObjectsAreNullOnUnitTests;
127 
AbstractFlagsSetterRule( RealLogger logger, String deviceConfigNamespace, String debugFlagPrefix, DeviceConfigHelper.InterfaceFactory deviceConfigInterfaceFactory, SystemPropertiesHelper.Interface systemPropertiesInterface)128     protected AbstractFlagsSetterRule(
129             RealLogger logger,
130             String deviceConfigNamespace,
131             String debugFlagPrefix,
132             DeviceConfigHelper.InterfaceFactory deviceConfigInterfaceFactory,
133             SystemPropertiesHelper.Interface systemPropertiesInterface) {
134         this(
135                 logger,
136                 deviceConfigNamespace,
137                 debugFlagPrefix,
138                 // TODO(b/294423183, 328682831): should not be necessary, but integrated with
139                 // setLogcatTag()
140                 (prop) ->
141                         prop.name.startsWith(debugFlagPrefix)
142                                 || prop.name.startsWith(SYSTEM_PROPERTY_FOR_LOGCAT_TAGS_PREFIX),
143                 deviceConfigInterfaceFactory,
144                 systemPropertiesInterface);
145     }
146 
AbstractFlagsSetterRule( RealLogger logger, String deviceConfigNamespace, String debugFlagPrefix, Matcher systemPropertiesMatcher, DeviceConfigHelper.InterfaceFactory deviceConfigInterfaceFactory, SystemPropertiesHelper.Interface systemPropertiesInterface)147     protected AbstractFlagsSetterRule(
148             RealLogger logger,
149             String deviceConfigNamespace,
150             String debugFlagPrefix,
151             Matcher systemPropertiesMatcher,
152             DeviceConfigHelper.InterfaceFactory deviceConfigInterfaceFactory,
153             SystemPropertiesHelper.Interface systemPropertiesInterface) {
154 
155         super(logger);
156 
157         mDeviceConfigNamespace =
158                 Objects.requireNonNull(
159                         deviceConfigNamespace, "deviceConfigNamespace cannot be null");
160         mDebugFlagPrefix =
161                 Objects.requireNonNull(debugFlagPrefix, "debugFlagPrefix cannot be null");
162         Objects.requireNonNull(systemPropertiesMatcher, "systemPropertiesMatcher cannot be null");
163         mSystemPropertiesMatcher =
164                 (prop) ->
165                         systemPropertiesMatcher.matches(prop) || INFRA_LOGTAG_MATCHER.matches(prop);
166         mDeviceConfig =
167                 new DeviceConfigHelper(deviceConfigInterfaceFactory, deviceConfigNamespace, logger);
168         mSystemProperties = new SystemPropertiesHelper(systemPropertiesInterface, logger);
169         storeSyncDisabledMode();
170         // Must set right away to avoid race conditions (for example, backend setting flags before
171         // apply() is called)
172         setSyncDisabledMode(DeviceConfig.SyncDisabledModeForTest.PERSISTENT);
173 
174         mFlagsSetter = flag -> defaultFlagsSetterImplementation(flag);
175         mSkipStuffWhenObjectsAreNullOnUnitTests = false;
176 
177         mLog.v(
178                 "Constructor: mDeviceConfigNamespace=%s,"
179                         + " mDebugFlagPrefix=%s,mDeviceConfig=%s, mSystemProperties=%s",
180                 mDeviceConfigNamespace, mDebugFlagPrefix, mDeviceConfig, mSystemProperties);
181     }
182 
183     // TODO(b/340882758): this constructor is only used by AbstractFlagsSetterRuleTestCase, which
184     // for now is only testing that the flags are set (it's not testing other stuff like checking
185     // they're reset, system properties, etc...), hence it sets some non-null fields as null. This
186     // is temporary, as this class should be refactored to use the new DeviceConfig class and be
187     // split into multiple rules (for example, to set DebugFlags and Logcat tags) - as more features
188     // are tested and/or refactored, these references should be properly set (and eventually the
189     // constructors merged);
AbstractFlagsSetterRule(RealLogger logger, Consumer<NameValuePair> flagsSetter)190     protected AbstractFlagsSetterRule(RealLogger logger, Consumer<NameValuePair> flagsSetter) {
191         super(logger);
192         mFlagsSetter = flagsSetter;
193 
194         mSkipStuffWhenObjectsAreNullOnUnitTests = true;
195         mSystemPropertiesMatcher = null;
196         mDeviceConfigNamespace = null;
197         mDeviceConfig = null;
198         mDebugFlagPrefix = null;
199         mSystemProperties = null;
200     }
201 
202     @Override
preTest(Statement base, Description description, List<Throwable> cleanUpErrors)203     protected void preTest(Statement base, Description description, List<Throwable> cleanUpErrors) {
204         String testName = TestHelper.getTestName(description);
205         mIsRunning = true;
206 
207         if (!mSkipStuffWhenObjectsAreNullOnUnitTests) {
208             // TODO(b/294423183): ideally should be "setupErrors", but it's not used yet (other
209             // than logging), so it doesn't matter
210             runSafely(cleanUpErrors, () -> mPreTestFlags.addAll(mDeviceConfig.getAll()));
211         }
212         // Log flags set on the device prior to test execution. Useful for verifying if flag state
213         // is correct for flag-ramp / AOAO testing.
214         log(mPreTestFlags, "pre-test flags");
215         if (!mSkipStuffWhenObjectsAreNullOnUnitTests) {
216         runSafely(
217                 cleanUpErrors,
218                 () ->
219                         mPreTestSystemProperties.addAll(
220                                 mSystemProperties.getAll(mSystemPropertiesMatcher)));
221         }
222 
223         runInitialCommands(testName);
224         setAnnotatedFlags(description);
225     }
226 
227     @Override
onTestFailure( Statement base, Description description, List<Throwable> cleanUpErrors, Throwable testFailure)228     protected void onTestFailure(
229             Statement base,
230             Description description,
231             List<Throwable> cleanUpErrors,
232             Throwable testFailure) {
233         runSafely(cleanUpErrors, () -> mOnTestFailureFlags.addAll(mDeviceConfig.getAll()));
234         runSafely(
235                 cleanUpErrors,
236                 () ->
237                         mOnTestFailureSystemProperties.addAll(
238                                 mSystemProperties.getAll(mSystemPropertiesMatcher)));
239     }
240 
241     @Override
postTest( Statement base, Description description, List<Throwable> cleanUpErrors)242     protected void postTest(
243             Statement base, Description description, List<Throwable> cleanUpErrors) {
244         String testName = TestHelper.getTestName(description);
245         runSafely(cleanUpErrors, () -> resetFlags(testName));
246         runSafely(cleanUpErrors, () -> resetSystemProperties(testName));
247         restoreSyncDisabledMode(cleanUpErrors);
248         mIsRunning = false;
249     }
250 
restoreSyncDisabledMode(List<Throwable> cleanUpErrors)251     private void restoreSyncDisabledMode(List<Throwable> cleanUpErrors) {
252         if (mPreviousSyncDisabledModeForTest != null) {
253             mLog.v(
254                     "mPreviousSyncDisabledModeForTest=%s; restoring flag sync mode",
255                     mPreviousSyncDisabledModeForTest);
256             runSafely(cleanUpErrors, () -> setSyncDisabledMode(mPreviousSyncDisabledModeForTest));
257         } else {
258             mLog.v("mPreviousSyncDisabledModeForTest=null; not restoring flag sync mode");
259         }
260     }
261 
setSyncDisabledMode(DeviceConfig.SyncDisabledModeForTest mode)262     private void setSyncDisabledMode(DeviceConfig.SyncDisabledModeForTest mode) {
263         runOrCache(
264                 "setSyncDisabledMode(" + mode + ")", () -> mDeviceConfig.setSyncDisabledMode(mode));
265     }
266 
storeSyncDisabledMode()267     private void storeSyncDisabledMode() {
268         runOrCache(
269                 "storeSyncDisabledMode()",
270                 () -> mPreviousSyncDisabledModeForTest = mDeviceConfig.getSyncDisabledMode());
271     }
272 
273     @Override
decorateTestFailureMessage(StringBuilder dump, List<Throwable> cleanUpErrors)274     protected String decorateTestFailureMessage(StringBuilder dump, List<Throwable> cleanUpErrors) {
275         if (mFlagsClearedByTest) {
276             dump.append("NOTE: test explicitly cleared all flags.\n");
277         }
278 
279         logAllAndDumpDiff("flags", dump, mChangedFlags, mPreTestFlags, mOnTestFailureFlags);
280         logAllAndDumpDiff(
281                 "system properties",
282                 dump,
283                 mChangedSystemProperties,
284                 mPreTestSystemProperties,
285                 mOnTestFailureSystemProperties);
286         return "flags / system properties state";
287     }
288 
logAllAndDumpDiff( String what, StringBuilder dump, Set<String> changedNames, List<NameValuePair> preTest, List<NameValuePair> postTest)289     private void logAllAndDumpDiff(
290             String what,
291             StringBuilder dump,
292             Set<String> changedNames,
293             List<NameValuePair> preTest,
294             List<NameValuePair> postTest) {
295         // Log all values
296         log(preTest, "%s before the test", what);
297         log(postTest, "%s after the test", what);
298 
299         // Dump only what was changed
300         appendChanges(dump, what, changedNames, preTest, postTest);
301     }
302 
appendChanges( StringBuilder dump, String what, Set<String> changedNames, List<NameValuePair> preTest, List<NameValuePair> postTest)303     private void appendChanges(
304             StringBuilder dump,
305             String what,
306             Set<String> changedNames,
307             List<NameValuePair> preTest,
308             List<NameValuePair> postTest) {
309         if (changedNames.isEmpty()) {
310             dump.append("Test didn't change any ").append(what).append('\n');
311             return;
312         }
313         dump.append("Test changed ")
314                 .append(changedNames.size())
315                 .append(' ')
316                 .append(what)
317                 .append(" (see log for all changes): \n");
318 
319         for (String name : changedNames) {
320             String before = getValue(preTest, name);
321             String after = getValue(postTest, name);
322             dump.append('\t')
323                     .append(name)
324                     .append(": ")
325                     .append("before=")
326                     .append(before)
327                     .append(", after=")
328                     .append(Objects.equals(before, after) ? "<<unchanged>>" : after)
329                     .append('\n');
330         }
331     }
332 
getValue(List<NameValuePair> list, String name)333     private String getValue(List<NameValuePair> list, String name) {
334         for (NameValuePair candidate : list) {
335             if (candidate.name.equals(name)) {
336                 return candidate.value;
337             }
338         }
339         return null;
340     }
341 
342     /**
343      * Dumps all flags using the {@value #TAG} tag.
344      *
345      * <p>Typically use for temporary debugging purposes like {@code dumpFlags("getFoo(%s)", bar)}.
346      */
347     @FormatMethod
dumpFlags(@ormatString String reasonFmt, @Nullable Object... reasonArgs)348     public final void dumpFlags(@FormatString String reasonFmt, @Nullable Object... reasonArgs) {
349         log(mDeviceConfig.getAll(), "flags (Reason: %s)", String.format(reasonFmt, reasonArgs));
350     }
351 
352     /**
353      * Dumps all system properties using the {@value #TAG} tag.
354      *
355      * <p>Typically use for temporary debugging purposes like {@code
356      * dumpSystemProperties("getFoo(%s)", bar)}.
357      */
358     @FormatMethod
dumpSystemProperties( @ormatString String reasonFmt, @Nullable Object... reasonArgs)359     public final void dumpSystemProperties(
360             @FormatString String reasonFmt, @Nullable Object... reasonArgs) {
361         log(
362                 mSystemProperties.getAll(mSystemPropertiesMatcher),
363                 "system properties (Reason: %s)",
364                 String.format(reasonFmt, reasonArgs));
365     }
366 
367     @FormatMethod
log( List<NameValuePair> values, @FormatString String whatFmt, @Nullable Object... whatArgs)368     private void log(
369             List<NameValuePair> values,
370             @FormatString String whatFmt,
371             @Nullable Object... whatArgs) {
372         String what = String.format(whatFmt, whatArgs);
373         if (values.isEmpty()) {
374             mLog.d("%s: empty", what);
375             return;
376         }
377         mLog.d("Logging (on VERBOSE) name/value of %d %s", values.size(), what);
378         values.forEach(value -> mLog.v("\t%s", value));
379     }
380 
381     /** Clears all flags from the namespace */
clearFlags()382     public final T clearFlags() {
383         return runOrCache(
384                 "clearFlags()",
385                 () -> {
386                     mLog.i("Clearing all flags. mIsRunning=%b", mIsRunning);
387                     mDeviceConfig.clearFlags();
388                     // TODO(b/294423183): ideally we should save the flags and restore - possibly
389                     // using DeviceConfig properties - but for now let's just clear it.
390                     mFlagsClearedByTest = true;
391                 });
392     }
393 
394     /** Sets the flag with the given value. */
setFlag(String name, boolean value)395     public final T setFlag(String name, boolean value) {
396         return setOrCacheFlag(name, Boolean.toString(value));
397     }
398 
399     /** Sets the flag with the given value. */
setFlag(String name, int value)400     public final T setFlag(String name, int value) {
401         return setOrCacheFlag(name, Integer.toString(value));
402     }
403 
404     /** Sets the flag with the given value. */
setFlag(String name, long value)405     public final T setFlag(String name, long value) {
406         return setOrCacheFlag(name, Long.toString(value));
407     }
408 
409     /** Sets the flag with the given value. */
setFlag(String name, float value)410     public final T setFlag(String name, float value) {
411         return setOrCacheFlag(name, Float.toString(value));
412     }
413 
414     /** Sets the flag with the given value. */
setFlag(String name, double value)415     public final T setFlag(String name, double value) {
416         return setOrCacheFlag(name, Double.toString(value));
417     }
418 
419     /**
420      * Sets the flag with the given values and the {@link #ARRAY_SPLITTER_COMMA} separator.
421      *
422      * <p>This method could also be used to set a simple (i.e., no array) String flag, as the
423      * separator is not added after the last element.
424      */
setFlag(String name, String... values)425     public final T setFlag(String name, String... values) {
426         return setArrayFlagWithExplicitSeparator(name, ARRAY_SPLITTER_COMMA, values);
427     }
428 
429     // TODO(b/303901926): static import from shared code instead
430     private static final String ARRAY_SPLITTER_COMMA = ",";
431 
432     /**
433      * Sets a string array flag with the given elements, separated by {@code separator}.
434      *
435      * <p>Use the method when you need to pass a explicitly {@code separator} - otherwise, just use
436      * {@link #setFlag(String, String...)}, it's simpler.
437      */
setArrayFlagWithExplicitSeparator( String name, String separator, String... values)438     public final T setArrayFlagWithExplicitSeparator(
439             String name, String separator, String... values) {
440         Objects.requireNonNull(separator, "separator cannot be null");
441         Objects.requireNonNull(values, "values cannot be null");
442         if (values.length == 0) {
443             throw new IllegalArgumentException("no values (name=" + name + ")");
444         }
445         if (values.length == 1) {
446             return setOrCacheFlag(name, values[0]);
447         }
448 
449         // TODO(b/303901926): use some existing helper / utility to flatten it - or a stream like
450         // list.stream().map(Object::toString).collect(Collectors.joining(delimiter) - once it's
451         // unit tested
452         StringBuilder flattenedValue = new StringBuilder().append(values[0]);
453         for (int i = 1; i < values.length; i++) {
454             String nextValue = values[i];
455             if (i < values.length) {
456                 flattenedValue.append(separator);
457             }
458             flattenedValue.append(nextValue);
459         }
460         return setOrCacheFlag(name, flattenedValue.toString(), separator);
461     }
462 
463     /**
464      * Sets a {@code logcat} tag.
465      *
466      * <p><b>Note: </b> it's clearer to use the {@link SetLogcatTag} annotation instead.
467      */
setLogcatTag(String tag, LogLevel level)468     public final T setLogcatTag(String tag, LogLevel level) {
469         setOrCacheLogtagSystemProperty(tag, level.name());
470         return getThis();
471     }
472 
473     // TODO(b/331781012): create @SetInfraLogcatTags as well
474     /** Sets the {@code logcat} tags for the (shared) infra classes. */
setInfraLogcatTags()475     public final T setInfraLogcatTags() {
476         for (String tag : INFRA_TAGS) {
477             setLogcatTag(tag, LogLevel.VERBOSE);
478         }
479         return getThis();
480     }
481 
482     /** Gets the value of the given flag. */
483     @Nullable
getFlag(String flag)484     public final String getFlag(String flag) {
485         return mDeviceConfig.get(flag);
486     }
487 
488     // TODO(295007931): abstract SDK-related methods in a new SdkLevelHelper and reuse them on
489     // SdkLevelSupportRule
490     /** Gets the device's SDK level. */
getDeviceSdk()491     protected abstract int getDeviceSdk();
492 
isAtLeastR()493     protected boolean isAtLeastR() {
494         return getDeviceSdk() >= 30;
495     }
496 
isAtLeastS()497     protected boolean isAtLeastS() {
498         return getDeviceSdk() >= 31;
499     }
500 
isAtLeastT()501     protected boolean isAtLeastT() {
502         return getDeviceSdk() > 32;
503     }
504 
505     // Helper to get a reference to this object, taking care of the generic casting.
506     @SuppressWarnings("unchecked")
getThis()507     protected final T getThis() {
508         return (T) this;
509     }
510 
511     // Set the annotated flags with the specified value for a particular test method.
512     // NOTE: when adding an annotation here, you also need to add it on isFlagAnnotationPresent()
setAnnotatedFlags(Description description)513     private void setAnnotatedFlags(Description description) {
514         List<Annotation> annotations = getAllFlagAnnotations(description);
515 
516         // Apply the annotations in the reverse order. First apply from the super classes, test
517         // class and then test method. If same annotated flag is present in class and test
518         // method, test method takes higher priority.
519         // NOTE: add annotations sorted by "most likely usage" and "groups"
520         for (int i = annotations.size() - 1; i >= 0; i--) {
521             Annotation annotation = annotations.get(i);
522 
523             // Boolean
524             if (annotation instanceof SetFlagEnabled) {
525                 setAnnotatedFlag((SetFlagEnabled) annotation);
526             } else if (annotation instanceof SetFlagsEnabled) {
527                 setAnnotatedFlag((SetFlagsEnabled) annotation);
528             } else if (annotation instanceof SetFlagDisabled) {
529                 setAnnotatedFlag((SetFlagDisabled) annotation);
530             } else if (annotation instanceof SetFlagsDisabled) {
531                 setAnnotatedFlag((SetFlagsDisabled) annotation);
532             } else if (annotation instanceof SetFlagTrue) {
533                 setAnnotatedFlag((SetFlagTrue) annotation);
534             } else if (annotation instanceof SetFlagsTrue) {
535                 setAnnotatedFlag((SetFlagsTrue) annotation);
536             } else if (annotation instanceof SetFlagFalse) {
537                 setAnnotatedFlag((SetFlagFalse) annotation);
538             } else if (annotation instanceof SetFlagsFalse) {
539                 setAnnotatedFlag((SetFlagsFalse) annotation);
540 
541                 // Numbers
542             } else if (annotation instanceof SetIntegerFlag) {
543                 setAnnotatedFlag((SetIntegerFlag) annotation);
544             } else if (annotation instanceof SetIntegerFlags) {
545                 setAnnotatedFlag((SetIntegerFlags) annotation);
546             } else if (annotation instanceof SetLongFlag) {
547                 setAnnotatedFlag((SetLongFlag) annotation);
548             } else if (annotation instanceof SetLongFlags) {
549                 setAnnotatedFlag((SetLongFlags) annotation);
550             } else if (annotation instanceof SetFloatFlag) {
551                 setAnnotatedFlag((SetFloatFlag) annotation);
552             } else if (annotation instanceof SetFloatFlags) {
553                 setAnnotatedFlag((SetFloatFlags) annotation);
554             } else if (annotation instanceof SetDoubleFlag) {
555                 setAnnotatedFlag((SetDoubleFlag) annotation);
556             } else if (annotation instanceof SetDoubleFlags) {
557                 setAnnotatedFlag((SetDoubleFlags) annotation);
558 
559                 // String
560             } else if (annotation instanceof SetStringFlag) {
561                 setAnnotatedFlag((SetStringFlag) annotation);
562             } else if (annotation instanceof SetStringFlags) {
563                 setAnnotatedFlag((SetStringFlags) annotation);
564             } else if (annotation instanceof SetStringArrayFlag) {
565                 setAnnotatedFlag((SetStringArrayFlag) annotation);
566             } else if (annotation instanceof SetStringArrayFlags) {
567                 setAnnotatedFlag((SetStringArrayFlags) annotation);
568 
569                 // Debug flags
570             } else if (annotation instanceof EnableDebugFlag) {
571                 setAnnotatedFlag((EnableDebugFlag) annotation);
572             } else if (annotation instanceof EnableDebugFlags) {
573                 setAnnotatedFlag((EnableDebugFlags) annotation);
574             } else if (annotation instanceof DisableDebugFlag) {
575                 setAnnotatedFlag((DisableDebugFlag) annotation);
576             } else if (annotation instanceof DisableDebugFlags) {
577                 setAnnotatedFlag((DisableDebugFlags) annotation);
578             } else if (annotation instanceof SetLongDebugFlag) {
579                 setAnnotatedFlag((SetLongDebugFlag) annotation);
580             } else if (annotation instanceof SetLongDebugFlags) {
581                 setAnnotatedFlag((SetLongDebugFlags) annotation);
582 
583                 // Logcat flags
584             } else if (annotation instanceof SetLogcatTag) {
585                 setAnnotatedFlag((SetLogcatTag) annotation);
586             } else if (annotation instanceof SetLogcatTags) {
587                 setAnnotatedFlag((SetLogcatTags) annotation);
588             } else {
589                 processAnnotation(description, annotation);
590             }
591         }
592     }
593 
setOrCacheFlag(String name, String value)594     private T setOrCacheFlag(String name, String value) {
595         return setOrCacheFlag(name, value, /* separator= */ null);
596     }
597 
598     // TODO(b/294423183): need to add unit test for setters that call this
setOrCacheFlag(String name, String value, @Nullable String separator)599     protected final T setOrCacheFlag(String name, String value, @Nullable String separator) {
600         Objects.requireNonNull(name, "name cannot be null");
601         NameValuePair flag = new NameValuePair(name, value, separator);
602         if (!mIsRunning) {
603             if (isFlagManagedByRunner(name)) {
604                 return getThis();
605             }
606             cacheCommand(new SetFlagCommand(flag));
607             return getThis();
608         }
609         return setFlag(flag);
610     }
611 
612     // TODO(b/295321663): need to provide a more elegant way to integrate it with the custom runners
isFlagManagedByRunner(String flag)613     protected boolean isFlagManagedByRunner(String flag) {
614         return false;
615     }
616 
setFlag(NameValuePair flag)617     private T setFlag(NameValuePair flag) {
618         mFlagsSetter.accept(flag);
619         return getThis();
620     }
621 
622     // Only used by the default mFlagsSetter - other methods should call setFlag()
defaultFlagsSetterImplementation(NameValuePair flag)623     private void defaultFlagsSetterImplementation(NameValuePair flag) {
624         mLog.d("Setting flag: %s", flag);
625         if (flag.separator == null) {
626             mDeviceConfig.set(flag.name, flag.value);
627         } else {
628             mDeviceConfig.setWithSeparator(flag.name, flag.value, flag.separator);
629         }
630         mChangedFlags.add(flag.name);
631     }
632 
resetFlags(String testName)633     private void resetFlags(String testName) {
634         if (mSkipStuffWhenObjectsAreNullOnUnitTests) {
635             mLog.w("resetFlags(%s): skipping (should only happen on rule test itself)", testName);
636             return;
637         }
638         mLog.d("Resetting flags after %s", testName);
639         mDeviceConfig.reset();
640     }
641 
642     /** Sets the value of the given {@link com.android.adservices.service.DebugFlag}. */
setDebugFlag(String name, boolean value)643     public final T setDebugFlag(String name, boolean value) {
644         return setDebugFlag(name, Boolean.toString(value));
645     }
646 
647     /** Sets the value of the given {@link com.android.adservices.service.DebugFlag}. */
setDebugFlag(String name, int value)648     public final T setDebugFlag(String name, int value) {
649         return setDebugFlag(name, Integer.toString(value));
650     }
651 
setOrCacheLogtagSystemProperty(String name, String value)652     private T setOrCacheLogtagSystemProperty(String name, String value) {
653         return setOrCacheSystemProperty(SYSTEM_PROPERTY_FOR_LOGCAT_TAGS_PREFIX + name, value);
654     }
655 
656     /** Sets the value of the given {@link com.android.adservices.service.DebugFlag}. */
setDebugFlag(String name, String value)657     public final T setDebugFlag(String name, String value) {
658         return setOrCacheSystemProperty(mDebugFlagPrefix + name, value);
659     }
660 
setOrCacheSystemProperty(String name, String value)661     private T setOrCacheSystemProperty(String name, String value) {
662         NameValuePair systemProperty = new NameValuePair(name, value);
663         if (!mIsRunning) {
664             cacheCommand(new SetSystemPropertyCommand(systemProperty));
665             return getThis();
666         }
667         return setSystemProperty(systemProperty);
668     }
669 
setSystemProperty(NameValuePair systemProperty)670     private T setSystemProperty(NameValuePair systemProperty) {
671         mLog.d("Setting system property: %s", systemProperty);
672         mSystemProperties.set(systemProperty.name, systemProperty.value);
673         mChangedSystemProperties.add(systemProperty.name);
674         return getThis();
675     }
676 
resetSystemProperties(String testName)677     private void resetSystemProperties(String testName) {
678         if (mSkipStuffWhenObjectsAreNullOnUnitTests) {
679             mLog.w(
680                     "resetSystemProperties(%s): skipping (should only happen on rule test itself)",
681                     testName);
682             return;
683         }
684         mLog.d("Resetting SystemProperties after %s", testName);
685         mSystemProperties.reset();
686     }
687 
runOrCache(String description, Runnable r)688     protected T runOrCache(String description, Runnable r) {
689         RunnableCommand command = new RunnableCommand(description, r);
690         if (!mIsRunning) {
691             cacheCommand(command);
692             return getThis();
693         }
694         command.execute();
695         return getThis();
696     }
697 
cacheCommand(Command command)698     private void cacheCommand(Command command) {
699         if (mIsRunning) {
700             throw new IllegalStateException(
701                     "Cannot cache " + command + " as test is already running");
702         }
703         mLog.v("Caching %s as test is not running yet", command);
704         mInitialCommands.add(command);
705     }
706 
runCommand(String description, Runnable runnable)707     private void runCommand(String description, Runnable runnable) {
708         mLog.v("Running runnable for %s", description);
709         runnable.run();
710     }
711 
712     // TODO(b/294423183): make private once not used by subclass for legacy methods
runInitialCommands(String testName)713     protected final void runInitialCommands(String testName) {
714         if (mInitialCommands.isEmpty()) {
715             mLog.d("Not executing any command before %s", testName);
716         } else {
717             int size = mInitialCommands.size();
718             mLog.d("Executing %d commands before %s", size, testName);
719             for (int i = 0; i < mInitialCommands.size(); i++) {
720                 Command command = mInitialCommands.get(i);
721                 mLog.v("\t%d: %s", i, command);
722                 command.execute();
723             }
724         }
725     }
726 
727     // TODO(b/294423183): improve logic used here and on setAnnotatedFlags()
728     // NOTE: when adding an annotation here, you also need to add it on setAnnotatedFlags()
isFlagAnnotationPresent(Annotation annotation)729     private boolean isFlagAnnotationPresent(Annotation annotation) {
730         // NOTE: add annotations sorted by "most likely usage" and "groups"
731         boolean processedHere =
732                 // Boolean
733                 (annotation instanceof SetFlagEnabled)
734                         || (annotation instanceof SetFlagsEnabled)
735                         || (annotation instanceof SetFlagDisabled)
736                         || (annotation instanceof SetFlagsDisabled)
737                         || (annotation instanceof SetFlagTrue)
738                         || (annotation instanceof SetFlagsTrue)
739                         || (annotation instanceof SetFlagFalse)
740                         || (annotation instanceof SetFlagsFalse)
741                         // Numbers
742                         || (annotation instanceof SetIntegerFlag)
743                         || (annotation instanceof SetIntegerFlags)
744                         || (annotation instanceof SetLongFlag)
745                         || (annotation instanceof SetLongFlags)
746                         || (annotation instanceof SetFloatFlag)
747                         || (annotation instanceof SetFloatFlags)
748                         || (annotation instanceof SetDoubleFlag)
749                         || (annotation instanceof SetDoubleFlags)
750                         // Strings
751                         || (annotation instanceof SetStringFlag)
752                         || (annotation instanceof SetStringFlags)
753                         || (annotation instanceof SetStringArrayFlag)
754                         || (annotation instanceof SetStringArrayFlags)
755                         // Debug flags
756                         || (annotation instanceof DisableDebugFlag)
757                         || (annotation instanceof DisableDebugFlags)
758                         || (annotation instanceof EnableDebugFlag)
759                         || (annotation instanceof EnableDebugFlags)
760                         || (annotation instanceof SetLongDebugFlag)
761                         || (annotation instanceof SetLongDebugFlags)
762                         // Logcat flags
763                         || (annotation instanceof SetLogcatTag)
764                         || (annotation instanceof SetLogcatTags);
765         return processedHere || isAnnotationSupported(annotation);
766     }
767 
768     /**
769      * By default returns {@code false}, but subclasses can override to support custom annotations.
770      *
771      * <p>Note: when overridden, {@link #processAnnotation(Description, Annotation)} should be
772      * overridden as well.
773      */
isAnnotationSupported(Annotation annotation)774     protected boolean isAnnotationSupported(Annotation annotation) {
775         return false;
776     }
777 
778     /**
779      * Called to process custom annotations present in the test (when {@link
780      * #isAnnotationSupported(Annotation)} returns {@code true} for that annotation type).
781      */
processAnnotation(Description description, Annotation annotation)782     protected void processAnnotation(Description description, Annotation annotation) {
783         throw new IllegalStateException(
784                 "Rule subclass ("
785                         + this.getClass().getName()
786                         + ") supports annotation "
787                         + annotation.annotationType().getName()
788                         + ", but doesn't override processAnnotation(), which was called with "
789                         + annotation);
790     }
791 
792     // TODO(b/377592216, 373477535): use TestHelper.getAnnotations() instead (after it has unit or
793     // integration tests)
getAllFlagAnnotations(Description description)794     private List<Annotation> getAllFlagAnnotations(Description description) {
795         List<Annotation> result = new ArrayList<>();
796         for (Annotation testMethodAnnotation : description.getAnnotations()) {
797             if (isFlagAnnotationPresent(testMethodAnnotation)) {
798                 result.add(testMethodAnnotation);
799             } else {
800                 mLog.v("Ignoring annotation %s", testMethodAnnotation);
801             }
802         }
803 
804         // Get all the flag based annotations from test class and super classes
805         Class<?> clazz = description.getTestClass();
806         do {
807             addFlagAnnotations(result, clazz);
808             for (Class<?> classInterface : clazz.getInterfaces()) {
809                 // TODO(b/340882758): add unit test for this as well. Also, unit test need to make
810                 // sure class prevails - for example, if interface has SetFlag(x, true) and test
811                 // have SetFlag(x, false), the interface annotation should be applied before the
812                 // class one.
813                 addFlagAnnotations(result, classInterface);
814             }
815             clazz = clazz.getSuperclass();
816         } while (clazz != null);
817 
818         return result;
819     }
820 
addFlagAnnotations(List<Annotation> annotations, Class<?> clazz)821     private void addFlagAnnotations(List<Annotation> annotations, Class<?> clazz) {
822         Annotation[] classAnnotations = clazz.getAnnotations();
823         if (classAnnotations == null) {
824             return;
825         }
826         for (Annotation annotation : classAnnotations) {
827             if (isFlagAnnotationPresent(annotation)) {
828                 annotations.add(annotation);
829             }
830         }
831     }
832 
833     // Single SetFlagEnabled annotations present
setAnnotatedFlag(SetFlagEnabled annotation)834     private void setAnnotatedFlag(SetFlagEnabled annotation) {
835         setFlag(annotation.value(), true);
836     }
837 
838     // Multiple SetFlagEnabled annotations present
setAnnotatedFlag(SetFlagsEnabled repeatedAnnotation)839     private void setAnnotatedFlag(SetFlagsEnabled repeatedAnnotation) {
840         for (SetFlagEnabled annotation : repeatedAnnotation.value()) {
841             setAnnotatedFlag(annotation);
842         }
843     }
844 
845     // Single SetFlagDisabled annotations present
setAnnotatedFlag(SetFlagDisabled annotation)846     private void setAnnotatedFlag(SetFlagDisabled annotation) {
847         setFlag(annotation.value(), false);
848     }
849 
850     // Multiple SetFlagDisabled annotations present
setAnnotatedFlag(SetFlagsDisabled repeatedAnnotation)851     private void setAnnotatedFlag(SetFlagsDisabled repeatedAnnotation) {
852         for (SetFlagDisabled annotation : repeatedAnnotation.value()) {
853             setAnnotatedFlag(annotation);
854         }
855     }
856 
857     // Single SetFlagTrue annotations present
setAnnotatedFlag(SetFlagTrue annotation)858     private void setAnnotatedFlag(SetFlagTrue annotation) {
859         setFlag(annotation.value(), true);
860     }
861 
862     // Multiple SetFlagTrue annotations present
setAnnotatedFlag(SetFlagsTrue repeatedAnnotation)863     private void setAnnotatedFlag(SetFlagsTrue repeatedAnnotation) {
864         for (SetFlagTrue annotation : repeatedAnnotation.value()) {
865             setAnnotatedFlag(annotation);
866         }
867     }
868 
869     // Single SetFlagFalse annotations present
setAnnotatedFlag(SetFlagFalse annotation)870     private void setAnnotatedFlag(SetFlagFalse annotation) {
871         setFlag(annotation.value(), false);
872     }
873 
874     // Multiple SetFlagFalse annotations present
setAnnotatedFlag(SetFlagsFalse repeatedAnnotation)875     private void setAnnotatedFlag(SetFlagsFalse repeatedAnnotation) {
876         for (SetFlagFalse annotation : repeatedAnnotation.value()) {
877             setAnnotatedFlag(annotation);
878         }
879     }
880 
881     // Single SetIntegerFlag annotations present
setAnnotatedFlag(SetIntegerFlag annotation)882     private void setAnnotatedFlag(SetIntegerFlag annotation) {
883         setFlag(annotation.name(), annotation.value());
884     }
885 
886     // Multiple SetIntegerFlag annotations present
setAnnotatedFlag(SetIntegerFlags repeatedAnnotation)887     private void setAnnotatedFlag(SetIntegerFlags repeatedAnnotation) {
888         for (SetIntegerFlag annotation : repeatedAnnotation.value()) {
889             setAnnotatedFlag(annotation);
890         }
891     }
892 
893     // Single SetLongFlag annotations present
setAnnotatedFlag(SetLongFlag annotation)894     private void setAnnotatedFlag(SetLongFlag annotation) {
895         setFlag(annotation.name(), annotation.value());
896     }
897 
898     // Multiple SetLongFlag annotations present
setAnnotatedFlag(SetLongFlags repeatedAnnotation)899     private void setAnnotatedFlag(SetLongFlags repeatedAnnotation) {
900         for (SetLongFlag annotation : repeatedAnnotation.value()) {
901             setAnnotatedFlag(annotation);
902         }
903     }
904 
905     // Single SetLongFlag annotations present
setAnnotatedFlag(SetFloatFlag annotation)906     private void setAnnotatedFlag(SetFloatFlag annotation) {
907         setFlag(annotation.name(), annotation.value());
908     }
909 
910     // Multiple SetLongFlag annotations present
setAnnotatedFlag(SetFloatFlags repeatedAnnotation)911     private void setAnnotatedFlag(SetFloatFlags repeatedAnnotation) {
912         for (SetFloatFlag annotation : repeatedAnnotation.value()) {
913             setAnnotatedFlag(annotation);
914         }
915     }
916 
917     // Single SetDoubleFlag annotations present
setAnnotatedFlag(SetDoubleFlag annotation)918     private void setAnnotatedFlag(SetDoubleFlag annotation) {
919         setFlag(annotation.name(), annotation.value());
920     }
921 
922     // Multiple SetDoubleFlag annotations present
setAnnotatedFlag(SetDoubleFlags repeatedAnnotation)923     private void setAnnotatedFlag(SetDoubleFlags repeatedAnnotation) {
924         for (SetDoubleFlag annotation : repeatedAnnotation.value()) {
925             setAnnotatedFlag(annotation);
926         }
927     }
928 
929     // Single SetStringFlag annotations present
setAnnotatedFlag(SetStringFlag annotation)930     private void setAnnotatedFlag(SetStringFlag annotation) {
931         setFlag(annotation.name(), annotation.value());
932     }
933 
934     // Multiple SetStringFlag annotations present
setAnnotatedFlag(SetStringFlags repeatedAnnotation)935     private void setAnnotatedFlag(SetStringFlags repeatedAnnotation) {
936         for (SetStringFlag annotation : repeatedAnnotation.value()) {
937             setAnnotatedFlag(annotation);
938         }
939     }
940 
941     // Single SetStringArrayFlag annotations present
setAnnotatedFlag(SetStringArrayFlag annotation)942     private void setAnnotatedFlag(SetStringArrayFlag annotation) {
943         setArrayFlagWithExplicitSeparator(
944                 annotation.name(), annotation.separator(), annotation.value());
945     }
946 
947     // Multiple SetStringArrayFlag annotations present
setAnnotatedFlag(SetStringArrayFlags repeatedAnnotation)948     private void setAnnotatedFlag(SetStringArrayFlags repeatedAnnotation) {
949         for (SetStringArrayFlag annotation : repeatedAnnotation.value()) {
950             setAnnotatedFlag(annotation);
951         }
952     }
953 
954     // Single EnableDebugFlag annotations present
setAnnotatedFlag(EnableDebugFlag annotation)955     private void setAnnotatedFlag(EnableDebugFlag annotation) {
956         setDebugFlag(annotation.value(), true);
957     }
958 
959     // Multiple EnableDebugFlag annotations present
setAnnotatedFlag(EnableDebugFlags repeatedAnnotation)960     private void setAnnotatedFlag(EnableDebugFlags repeatedAnnotation) {
961         for (EnableDebugFlag annotation : repeatedAnnotation.value()) {
962             setAnnotatedFlag(annotation);
963         }
964     }
965 
966     // Single DisableDebugFlag annotations present
setAnnotatedFlag(DisableDebugFlag annotation)967     private void setAnnotatedFlag(DisableDebugFlag annotation) {
968         setDebugFlag(annotation.value(), false);
969     }
970 
971     // Multiple DisableDebugFlag annotations present
setAnnotatedFlag(DisableDebugFlags repeatedAnnotation)972     private void setAnnotatedFlag(DisableDebugFlags repeatedAnnotation) {
973         for (DisableDebugFlag annotation : repeatedAnnotation.value()) {
974             setAnnotatedFlag(annotation);
975         }
976     }
977 
978     // Single SetLongDebugFlag annotations present
setAnnotatedFlag(SetLongDebugFlag annotation)979     private void setAnnotatedFlag(SetLongDebugFlag annotation) {
980         setDebugFlag(annotation.name(), Long.toString(annotation.value()));
981     }
982 
983     // Multiple SetLongDebugFlag annotations present
setAnnotatedFlag(SetLongDebugFlags repeatedAnnotation)984     private void setAnnotatedFlag(SetLongDebugFlags repeatedAnnotation) {
985         for (SetLongDebugFlag annotation : repeatedAnnotation.value()) {
986             setAnnotatedFlag(annotation);
987         }
988     }
989 
990     // Single SetLogcatTag annotations present
setAnnotatedFlag(SetLogcatTag annotation)991     private void setAnnotatedFlag(SetLogcatTag annotation) {
992         setLogcatTag(annotation.tag(), annotation.level());
993     }
994 
995     // Multiple SetLogcatTag annotations present
setAnnotatedFlag(SetLogcatTags repeatedAnnotation)996     private void setAnnotatedFlag(SetLogcatTags repeatedAnnotation) {
997         for (SetLogcatTag annotation : repeatedAnnotation.value()) {
998             setAnnotatedFlag(annotation);
999         }
1000     }
1001 
1002     @SuppressWarnings("ClassCanBeStatic") // Subclasses reference enclosing class
1003     private abstract class Command {
1004         protected final String mDescription;
1005 
Command(String description)1006         Command(String description) {
1007             mDescription = description;
1008         }
1009 
execute()1010         abstract void execute();
1011 
1012         @Override
toString()1013         public final String toString() {
1014             return mDescription;
1015         }
1016     }
1017 
1018     private final class RunnableCommand extends Command {
1019         private final Runnable mRunnable;
1020 
RunnableCommand(String description, Runnable runnable)1021         RunnableCommand(String description, Runnable runnable) {
1022             super(description);
1023             mRunnable = runnable;
1024         }
1025 
1026         @Override
execute()1027         void execute() {
1028             runCommand(mDescription, mRunnable);
1029         }
1030     }
1031 
1032     private abstract class SetFlagOrSystemPropertyCommand extends Command {
1033         protected final NameValuePair mFlagOrSystemProperty;
1034 
SetFlagOrSystemPropertyCommand(String description, NameValuePair flagOrSystemProperty)1035         SetFlagOrSystemPropertyCommand(String description, NameValuePair flagOrSystemProperty) {
1036             super(description + "(" + flagOrSystemProperty + ")");
1037             mFlagOrSystemProperty = flagOrSystemProperty;
1038         }
1039     }
1040 
1041     private final class SetFlagCommand extends SetFlagOrSystemPropertyCommand {
SetFlagCommand(NameValuePair flag)1042         SetFlagCommand(NameValuePair flag) {
1043             super("SetFlag", flag);
1044         }
1045 
1046         @Override
execute()1047         void execute() {
1048             setFlag(mFlagOrSystemProperty);
1049         }
1050     }
1051 
1052     private final class SetSystemPropertyCommand extends SetFlagOrSystemPropertyCommand {
SetSystemPropertyCommand(NameValuePair flag)1053         SetSystemPropertyCommand(NameValuePair flag) {
1054             super("SetSystemProperty", flag);
1055         }
1056 
1057         @Override
execute()1058         void execute() {
1059             setSystemProperty(mFlagOrSystemProperty);
1060         }
1061     }
1062 }
1063