1 /*
2  * Copyright (C) 2020 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.bedstead.harrier;
18 
19 import static android.Manifest.permission.INTERACT_ACROSS_USERS_FULL;
20 
21 import static com.android.bedstead.nene.permissions.Permissions.NOTIFY_PENDING_SYSTEM_UPDATE;
22 import static com.android.bedstead.nene.users.UserType.MANAGED_PROFILE_TYPE_NAME;
23 import static com.android.bedstead.nene.users.UserType.SECONDARY_USER_TYPE_NAME;
24 import static com.android.bedstead.nene.utils.Versions.meetsSdkVersionRequirements;
25 import static com.android.bedstead.remotedpc.Configuration.REMOTE_DPC_COMPONENT_NAME;
26 
27 import static com.google.common.truth.Truth.assertWithMessage;
28 
29 import static org.junit.Assert.assertFalse;
30 import static org.junit.Assume.assumeFalse;
31 import static org.junit.Assume.assumeTrue;
32 
33 import android.content.Context;
34 import android.content.Intent;
35 import android.os.Build;
36 import android.os.Bundle;
37 import android.util.Log;
38 
39 import androidx.test.core.app.ApplicationProvider;
40 import androidx.test.platform.app.InstrumentationRegistry;
41 
42 import com.android.bedstead.harrier.annotations.AfterClass;
43 import com.android.bedstead.harrier.annotations.BeforeClass;
44 import com.android.bedstead.harrier.annotations.EnsureDoesNotHavePermission;
45 import com.android.bedstead.harrier.annotations.EnsureHasPermission;
46 import com.android.bedstead.harrier.annotations.EnsurePackageNotInstalled;
47 import com.android.bedstead.harrier.annotations.FailureMode;
48 import com.android.bedstead.harrier.annotations.RequireDoesNotHaveFeature;
49 import com.android.bedstead.harrier.annotations.RequireFeature;
50 import com.android.bedstead.harrier.annotations.RequireGmsInstrumentation;
51 import com.android.bedstead.harrier.annotations.RequirePackageInstalled;
52 import com.android.bedstead.harrier.annotations.RequirePackageNotInstalled;
53 import com.android.bedstead.harrier.annotations.RequireSdkVersion;
54 import com.android.bedstead.harrier.annotations.RequireUserSupported;
55 import com.android.bedstead.harrier.annotations.enterprise.EnsureHasDeviceOwner;
56 import com.android.bedstead.harrier.annotations.enterprise.EnsureHasNoDeviceOwner;
57 import com.android.bedstead.harrier.annotations.enterprise.EnsureHasNoProfileOwner;
58 import com.android.bedstead.harrier.annotations.enterprise.EnsureHasProfileOwner;
59 import com.android.bedstead.harrier.annotations.meta.EnsureHasNoProfileAnnotation;
60 import com.android.bedstead.harrier.annotations.meta.EnsureHasNoUserAnnotation;
61 import com.android.bedstead.harrier.annotations.meta.EnsureHasProfileAnnotation;
62 import com.android.bedstead.harrier.annotations.meta.EnsureHasUserAnnotation;
63 import com.android.bedstead.harrier.annotations.meta.ParameterizedAnnotation;
64 import com.android.bedstead.harrier.annotations.meta.RequireRunOnProfileAnnotation;
65 import com.android.bedstead.harrier.annotations.meta.RequireRunOnUserAnnotation;
66 import com.android.bedstead.harrier.annotations.meta.RequiresBedsteadJUnit4;
67 import com.android.bedstead.nene.TestApis;
68 import com.android.bedstead.nene.devicepolicy.DeviceOwner;
69 import com.android.bedstead.nene.devicepolicy.DevicePolicyController;
70 import com.android.bedstead.nene.devicepolicy.ProfileOwner;
71 import com.android.bedstead.nene.exceptions.AdbException;
72 import com.android.bedstead.nene.exceptions.NeneException;
73 import com.android.bedstead.nene.packages.Package;
74 import com.android.bedstead.nene.permissions.PermissionContextImpl;
75 import com.android.bedstead.nene.users.User;
76 import com.android.bedstead.nene.users.UserBuilder;
77 import com.android.bedstead.nene.users.UserReference;
78 import com.android.bedstead.nene.utils.ShellCommand;
79 import com.android.bedstead.nene.utils.Versions;
80 import com.android.bedstead.remotedpc.RemoteDpc;
81 import com.android.compatibility.common.util.BlockingBroadcastReceiver;
82 
83 import com.google.common.base.Objects;
84 
85 import junit.framework.AssertionFailedError;
86 
87 import org.junit.AssumptionViolatedException;
88 import org.junit.rules.TestRule;
89 import org.junit.runner.Description;
90 import org.junit.runners.model.FrameworkMethod;
91 import org.junit.runners.model.Statement;
92 import org.junit.runners.model.TestClass;
93 
94 import java.lang.annotation.Annotation;
95 import java.lang.reflect.InvocationTargetException;
96 import java.lang.reflect.Method;
97 import java.lang.reflect.Modifier;
98 import java.util.ArrayList;
99 import java.util.Arrays;
100 import java.util.Collection;
101 import java.util.Collections;
102 import java.util.HashMap;
103 import java.util.HashSet;
104 import java.util.List;
105 import java.util.Map;
106 import java.util.Set;
107 import java.util.function.Function;
108 
109 
110 /**
111  * A Junit rule which exposes methods for efficiently changing and querying device state.
112  *
113  * <p>States set by the methods on this class will by default be cleaned up after the test.
114  *
115  *
116  * <p>Using this rule also enforces preconditions in annotations from the
117  * {@code com.android.comaptibility.common.util.enterprise.annotations} package.
118  *
119  * {@code assumeTrue} will be used, so tests which do not meet preconditions will be skipped.
120  */
121 public final class DeviceState implements TestRule {
122 
123     private static final String GMS_PKG = "com.google.android.gms";
124 
125     private final Context mContext = ApplicationProvider.getApplicationContext();
126     private static final TestApis sTestApis = new TestApis();
127     private static final String SKIP_TEST_TEARDOWN_KEY = "skip-test-teardown";
128     private static final String SKIP_CLASS_TEARDOWN_KEY = "skip-class-teardown";
129     private static final String SKIP_TESTS_REASON_KEY = "skip-tests-reason";
130     private boolean mSkipTestTeardown;
131     private boolean mSkipClassTeardown;
132     private boolean mSkipTests;
133     private boolean mFailTests;
134     private boolean mUsingBedsteadJUnit4 = false;
135     private String mSkipTestsReason;
136     private String mFailTestsReason;
137 
138     // Marks if the conditions for requiring running under GMS instrumentation have been set
139     // if not - we assume the test should never run under GMS instrumentation
140     private boolean mHasRequireGmsInstrumentation = false;
141 
142     private static final String TV_PROFILE_TYPE_NAME = "com.android.tv.profile";
143 
DeviceState()144     public DeviceState() {
145         Bundle arguments = InstrumentationRegistry.getArguments();
146         mSkipTestTeardown = Boolean.parseBoolean(
147                 arguments.getString(SKIP_TEST_TEARDOWN_KEY, "false"));
148         mSkipClassTeardown = Boolean.parseBoolean(
149                 arguments.getString(SKIP_CLASS_TEARDOWN_KEY, "false"));
150         mSkipTestsReason = arguments.getString(SKIP_TESTS_REASON_KEY, "");
151         mSkipTests = !mSkipTestsReason.isEmpty();
152     }
153 
setSkipTestTeardown(boolean skipTestTeardown)154     void setSkipTestTeardown(boolean skipTestTeardown) {
155         mSkipTestTeardown = skipTestTeardown;
156     }
157 
setUsingBedsteadJUnit4(boolean usingBedsteadJUnit4)158     void setUsingBedsteadJUnit4(boolean usingBedsteadJUnit4) {
159         mUsingBedsteadJUnit4 = usingBedsteadJUnit4;
160     }
161 
apply(final Statement base, final Description description)162     @Override public Statement apply(final Statement base,
163             final Description description) {
164 
165         if (description.isTest()) {
166             return applyTest(base, description);
167         } else if (description.isSuite()) {
168             return applySuite(base, description);
169         }
170         throw new IllegalStateException("Unknown description type: " + description);
171     }
172 
applyTest(final Statement base, final Description description)173     private Statement applyTest(final Statement base, final Description description) {
174         return new Statement() {
175             @Override public void evaluate() throws Throwable {
176                 Log.d(LOG_TAG, "Preparing state for test " + description.getMethodName());
177 
178                 assumeFalse(mSkipTestsReason, mSkipTests);
179                 assertFalse(mFailTestsReason, mFailTests);
180 
181                 List<Annotation> annotations = getAnnotations(description);
182                 PermissionContextImpl permissionContext = applyAnnotations(annotations);
183 
184                 Log.d(LOG_TAG,
185                         "Finished preparing state for test " + description.getMethodName());
186 
187                 try {
188                     base.evaluate();
189                 } finally {
190                     Log.d(LOG_TAG,
191                             "Tearing down state for test " + description.getMethodName());
192 
193                     if (permissionContext != null) {
194                         permissionContext.close();
195                     }
196 
197                     teardownNonShareableState();
198                     if (!mSkipTestTeardown) {
199                         teardownShareableState();
200                     }
201                     Log.d(LOG_TAG,
202                             "Finished tearing down state for test "
203                                     + description.getMethodName());
204                 }
205             }};
206     }
207 
208     private PermissionContextImpl applyAnnotations(List<Annotation> annotations)
209             throws Throwable {
210         PermissionContextImpl permissionContext = null;
211         for (Annotation annotation : annotations) {
212             Class<? extends Annotation> annotationType = annotation.annotationType();
213 
214             EnsureHasNoProfileAnnotation ensureHasNoProfileAnnotation =
215                     annotationType.getAnnotation(EnsureHasNoProfileAnnotation.class);
216             if (ensureHasNoProfileAnnotation != null) {
217                 UserType userType = (UserType) annotation.annotationType()
218                         .getMethod("forUser").invoke(annotation);
219                 ensureHasNoProfile(ensureHasNoProfileAnnotation.value(), userType);
220                 continue;
221             }
222 
223             EnsureHasProfileAnnotation ensureHasProfileAnnotation =
224                     annotationType.getAnnotation(EnsureHasProfileAnnotation.class);
225             if (ensureHasProfileAnnotation != null) {
226                 UserType forUser = (UserType) annotation.annotationType()
227                         .getMethod("forUser").invoke(annotation);
228                 OptionalBoolean installInstrumentedApp = (OptionalBoolean)
229                         annotation.annotationType()
230                                 .getMethod("installInstrumentedApp").invoke(annotation);
231 
232                 boolean dpcIsPrimary = false;
233                 if (ensureHasProfileAnnotation.hasProfileOwner()) {
234                     dpcIsPrimary = (boolean)
235                             annotation.annotationType()
236                                     .getMethod("dpcIsPrimary").invoke(annotation);
237                 }
238 
239                 ensureHasProfile(
240                         ensureHasProfileAnnotation.value(), installInstrumentedApp,
241                         forUser, ensureHasProfileAnnotation.hasProfileOwner(),
242                         dpcIsPrimary);
243                 continue;
244             }
245 
246             EnsureHasNoUserAnnotation ensureHasNoUserAnnotation =
247                     annotationType.getAnnotation(EnsureHasNoUserAnnotation.class);
248             if (ensureHasNoUserAnnotation != null) {
249                 ensureHasNoUser(ensureHasNoUserAnnotation.value());
250                 continue;
251             }
252 
253             EnsureHasUserAnnotation ensureHasUserAnnotation =
254                     annotationType.getAnnotation(EnsureHasUserAnnotation.class);
255             if (ensureHasUserAnnotation != null) {
256                 OptionalBoolean installInstrumentedApp = (OptionalBoolean)
257                         annotation.getClass()
258                                 .getMethod("installInstrumentedApp").invoke(annotation);
259                 ensureHasUser(ensureHasUserAnnotation.value(), installInstrumentedApp);
260                 continue;
261             }
262 
263             RequireRunOnUserAnnotation requireRunOnUserAnnotation =
264                     annotationType.getAnnotation(RequireRunOnUserAnnotation.class);
265             if (requireRunOnUserAnnotation != null) {
266                 requireRunOnUser(requireRunOnUserAnnotation.value());
267                 continue;
268             }
269 
270             RequireRunOnProfileAnnotation requireRunOnProfileAnnotation =
271                     annotationType.getAnnotation(RequireRunOnProfileAnnotation.class);
272             if (requireRunOnProfileAnnotation != null) {
273                 OptionalBoolean installInstrumentedAppInParent = (OptionalBoolean)
274                         annotation.getClass()
275                                 .getMethod("installInstrumentedAppInParent")
276                                 .invoke(annotation);
277 
278 
279                 boolean dpcIsPrimary = false;
280                 Set<String> affiliationIds = null;
281                 if (requireRunOnProfileAnnotation.hasProfileOwner()) {
282                     dpcIsPrimary = (boolean)
283                             annotation.annotationType()
284                                     .getMethod("dpcIsPrimary").invoke(annotation);
285                     affiliationIds = new HashSet<>(Arrays.asList((String[])
286                             annotation.annotationType()
287                                     .getMethod("affiliationIds").invoke(annotation)));
288                 }
289 
290                 requireRunOnProfile(requireRunOnProfileAnnotation.value(),
291                         installInstrumentedAppInParent,
292                         requireRunOnProfileAnnotation.hasProfileOwner(),
293                         dpcIsPrimary, affiliationIds);
294                 continue;
295             }
296 
297             if (annotation instanceof EnsureHasDeviceOwner) {
298                 EnsureHasDeviceOwner ensureHasDeviceOwnerAnnotation =
299                         (EnsureHasDeviceOwner) annotation;
300                 ensureHasDeviceOwner(ensureHasDeviceOwnerAnnotation.onUser(),
301                         ensureHasDeviceOwnerAnnotation.failureMode(),
302                         ensureHasDeviceOwnerAnnotation.isPrimary(),
303                         new HashSet<>(Arrays.asList(ensureHasDeviceOwnerAnnotation.affiliationIds())));
304                 continue;
305             }
306 
307             if (annotation instanceof EnsureHasNoDeviceOwner) {
308                 ensureHasNoDeviceOwner();
309                 continue;
310             }
311 
312             if (annotation instanceof RequireFeature) {
313                 RequireFeature requireFeatureAnnotation = (RequireFeature) annotation;
314                 requireFeature(
315                         requireFeatureAnnotation.value(),
316                         requireFeatureAnnotation.failureMode());
317                 continue;
318             }
319 
320             if (annotation instanceof RequireDoesNotHaveFeature) {
321                 RequireDoesNotHaveFeature requireDoesNotHaveFeatureAnnotation =
322                         (RequireDoesNotHaveFeature) annotation;
323                 requireDoesNotHaveFeature(
324                         requireDoesNotHaveFeatureAnnotation.value(),
325                         requireDoesNotHaveFeatureAnnotation.failureMode());
326                 continue;
327             }
328 
329             if (annotation instanceof EnsureHasProfileOwner) {
330                 EnsureHasProfileOwner ensureHasProfileOwnerAnnotation =
331                         (EnsureHasProfileOwner) annotation;
332                 ensureHasProfileOwner(ensureHasProfileOwnerAnnotation.onUser(),
333                         ensureHasProfileOwnerAnnotation.isPrimary(),
334                         new HashSet<>(Arrays.asList(ensureHasProfileOwnerAnnotation.affiliationIds())));
335                 continue;
336             }
337 
338             if (annotationType.equals(EnsureHasNoProfileOwner.class)) {
339                 EnsureHasNoProfileOwner ensureHasNoProfileOwnerAnnotation =
340                         (EnsureHasNoProfileOwner) annotation;
341                 ensureHasNoProfileOwner(ensureHasNoProfileOwnerAnnotation.onUser());
342                 continue;
343             }
344 
345             if (annotation instanceof RequireUserSupported) {
346                 RequireUserSupported requireUserSupportedAnnotation =
347                         (RequireUserSupported) annotation;
348                 requireUserSupported(
349                         requireUserSupportedAnnotation.value(),
350                         requireUserSupportedAnnotation.failureMode());
351                 continue;
352             }
353 
354             if (annotation instanceof RequireGmsInstrumentation) {
355                 RequireGmsInstrumentation requireGmsInstrumentationAnnotation =
356                         (RequireGmsInstrumentation) annotation;
357                 requireGmsInstrumentation(requireGmsInstrumentationAnnotation.min(),
358                         requireGmsInstrumentationAnnotation.max());
359                 continue;
360             }
361 
362             if (annotation instanceof RequireSdkVersion) {
363                 RequireSdkVersion requireSdkVersionAnnotation =
364                         (RequireSdkVersion) annotation;
365 
366                 requireSdkVersion(
367                         requireSdkVersionAnnotation.min(),
368                         requireSdkVersionAnnotation.max(),
369                         requireSdkVersionAnnotation.failureMode());
370                 continue;
371             }
372 
373             if (annotation instanceof RequirePackageInstalled) {
374                 RequirePackageInstalled requirePackageInstalledAnnotation =
375                         (RequirePackageInstalled) annotation;
376                 requirePackageInstalled(
377                         requirePackageInstalledAnnotation.value(),
378                         requirePackageInstalledAnnotation.onUser(),
379                         requirePackageInstalledAnnotation.failureMode());
380                 continue;
381             }
382 
383             if (annotation instanceof RequirePackageNotInstalled) {
384                 RequirePackageNotInstalled requirePackageNotInstalledAnnotation =
385                         (RequirePackageNotInstalled) annotation;
386                 requirePackageNotInstalled(
387                         requirePackageNotInstalledAnnotation.value(),
388                         requirePackageNotInstalledAnnotation.onUser(),
389                         requirePackageNotInstalledAnnotation.failureMode()
390                 );
391                 continue;
392             }
393 
394             if (annotation instanceof EnsurePackageNotInstalled) {
395                 EnsurePackageNotInstalled ensurePackageNotInstalledAnnotation =
396                         (EnsurePackageNotInstalled) annotation;
397                 ensurePackageNotInstalled(
398                         ensurePackageNotInstalledAnnotation.value(),
399                         ensurePackageNotInstalledAnnotation.onUser()
400                 );
401                 continue;
402             }
403 
404             if (annotation instanceof EnsureHasPermission) {
405                 EnsureHasPermission ensureHasPermissionAnnotation =
406                         (EnsureHasPermission) annotation;
407 
408                 for (String permission : ensureHasPermissionAnnotation.value()) {
409                     ensureCanGetPermission(permission);
410                 }
411 
412 
413                 try {
414                     if (permissionContext == null) {
415                         permissionContext = sTestApis.permissions().withPermission(
416                                 ensureHasPermissionAnnotation.value());
417                     } else {
418                         permissionContext = permissionContext.withPermission(
419                                 ensureHasPermissionAnnotation.value());
420                     }
421                 } catch (NeneException e) {
422                     failOrSkip("Error getting permission: " + e,
423                             ensureHasPermissionAnnotation.failureMode());
424                 }
425                 continue;
426             }
427 
428             if (annotation instanceof EnsureDoesNotHavePermission) {
429                 EnsureDoesNotHavePermission ensureDoesNotHavePermission =
430                         (EnsureDoesNotHavePermission) annotation;
431 
432                 try {
433                     if (permissionContext == null) {
434                         permissionContext = sTestApis.permissions().withoutPermission(
435                                 ensureDoesNotHavePermission.value());
436                     } else {
437                         permissionContext = permissionContext.withoutPermission(
438                                 ensureDoesNotHavePermission.value());
439                     }
440                 } catch (NeneException e) {
441                     failOrSkip("Error denying permission: " + e,
442                             ensureDoesNotHavePermission.failureMode());
443                 }
444                 continue;
445             }
446         }
447 
448         if (!mHasRequireGmsInstrumentation) {
449             // TODO(scottjonathan): Only enforce if we've configured GMS Instrumentation
450             requireNoGmsInstrumentation();
451         }
452 
453         return permissionContext;
454     }
455 
456     private List<Annotation> getAnnotations(Description description) {
457         if (mUsingBedsteadJUnit4 && description.isTest()) {
458             // The annotations are already exploded for tests
459             return new ArrayList<>(description.getAnnotations());
460         }
461 
462         // Otherwise we should build a new collection by recursively gathering annotations
463         // if we find any which don't work without the runner we should error and fail the test
464         List<Annotation> annotations = new ArrayList<>();
465 
466         if (description.isTest()) {
467             annotations =
468                     new ArrayList<>(Arrays.asList(description.getTestClass().getAnnotations()));
469         }
470 
471         annotations.addAll(description.getAnnotations());
472 
473         checkAnnotations(annotations);
474 
475         BedsteadJUnit4.resolveRecursiveAnnotations(annotations,
476                 /* parameterizedAnnotation= */ null);
477 
478         checkAnnotations(annotations);
479 
480         return annotations;
481     }
482 
483     private void checkAnnotations(Collection<Annotation> annotations) {
484         for (Annotation annotation : annotations) {
485             if (annotation.annotationType().getAnnotation(RequiresBedsteadJUnit4.class) != null
486                     || annotation.annotationType().getAnnotation(
487                             ParameterizedAnnotation.class) != null) {
488                 throw new AssertionFailedError("Test is annotated "
489                         + annotation.annotationType().getSimpleName()
490                         + " which requires using the BedsteadJUnit4 test runner");
491             }
492         }
493     }
494 
495     private Statement applySuite(final Statement base, final Description description) {
496         return new Statement() {
497             @Override
498             public void evaluate() throws Throwable {
499                 checkValidAnnotations(description);
500 
501                 TestClass testClass = new TestClass(description.getTestClass());
502 
503                 boolean skipAfterClass = true;
504                 PermissionContextImpl permissionContext = null;
505 
506                 if (!mSkipTests && !mFailTests) {
507                     skipAfterClass = false;
508                     Log.d(LOG_TAG, "Preparing state for suite " + description.getClassName());
509 
510                     try {
511                         List<Annotation> annotations =
512                                 new ArrayList<>(getAnnotations(description));
513                         permissionContext = applyAnnotations(annotations);
514                     } catch (AssumptionViolatedException e) {
515                         mSkipTests = true;
516                         mSkipTestsReason = e.getMessage();
517                     } catch (AssertionError e) {
518                         mFailTests = true;
519                         mFailTestsReason = e.getMessage();
520                     }
521 
522                     Log.d(LOG_TAG,
523                             "Finished preparing state for suite "
524                                     + description.getClassName());
525                 }
526 
527                 if (!mSkipTests && !mFailTests) {
528                     runAnnotatedMethods(testClass, BeforeClass.class);
529                 }
530 
531                 base.evaluate();
532 
533                 if (!skipAfterClass) {
534                     runAnnotatedMethods(testClass, AfterClass.class);
535                 }
536 
537                 if (permissionContext != null) {
538                     permissionContext.close();
539                 }
540 
541                 if (!mSkipClassTeardown) {
542                     teardownShareableState();
543                 }
544             }
545         };
546     }
547 
548     private static final Map<Class<? extends Annotation>, Class<? extends Annotation>>
549             BANNED_ANNOTATIONS_TO_REPLACEMENTS = getBannedAnnotationsToReplacements();
550     private static Map<
551             Class<? extends Annotation>,
552             Class<? extends Annotation>> getBannedAnnotationsToReplacements() {
553         Map<
554                 Class<? extends Annotation>,
555                 Class<? extends Annotation>> bannedAnnotationsToReplacements = new HashMap<>();
556         bannedAnnotationsToReplacements.put(org.junit.BeforeClass.class, BeforeClass.class);
557         bannedAnnotationsToReplacements.put(org.junit.AfterClass.class, AfterClass.class);
558         return bannedAnnotationsToReplacements;
559     }
560 
561     private void checkValidAnnotations(Description classDescription) {
562         for (Method method : classDescription.getTestClass().getMethods()) {
563             for (Map.Entry<
564                     Class<? extends Annotation>,
565                     Class<? extends Annotation>> bannedAnnotation
566                     : BANNED_ANNOTATIONS_TO_REPLACEMENTS.entrySet()) {
567                 if (method.isAnnotationPresent(bannedAnnotation.getKey())) {
568                     throw new IllegalStateException("Do not use "
569                             + bannedAnnotation.getKey().getCanonicalName()
570                             + " when using DeviceState, replace with "
571                             + bannedAnnotation.getValue().getCanonicalName());
572                 }
573             }
574 
575             if (method.getAnnotation(BeforeClass.class) != null
576                     || method.getAnnotation(AfterClass.class) != null) {
577                 checkPublicStaticVoidNoArgs(method);
578             }
579         }
580     }
581 
582     private void checkPublicStaticVoidNoArgs(Method method) {
583         if (method.getParameterTypes().length > 0) {
584             throw new IllegalStateException(
585                     "Method " + method.getName() + " should have no parameters");
586         }
587         if (method.getReturnType() != Void.TYPE) {
588             throw new IllegalStateException("Method " + method.getName() + "() should be void");
589         }
590         if (!Modifier.isStatic(method.getModifiers())) {
591             throw new IllegalStateException("Method " + method.getName() + "() should be static");
592         }
593         if (!Modifier.isPublic(method.getModifiers())) {
594             throw new IllegalStateException("Method " + method.getName() + "() should be public");
595         }
596     }
597 
598     private void runAnnotatedMethods(
599             TestClass testClass, Class<? extends Annotation> annotation) throws Throwable {
600 
601         List<FrameworkMethod> methods = new ArrayList<>(testClass.getAnnotatedMethods(annotation));
602         Collections.reverse(methods);
603         for (FrameworkMethod method : methods) {
604             try {
605                 method.invokeExplosively(testClass.getJavaClass());
606             } catch (InvocationTargetException e) {
607                 throw e.getCause();
608             }
609         }
610     }
611 
612     private void requireRunOnUser(String userType) {
613         User instrumentedUser = sTestApis.users().instrumented().resolve();
614 
615         assumeTrue("This test only runs on users of type " + userType,
616                 instrumentedUser.type().name().equals(userType));
617 
618         mUsers.put(instrumentedUser.type(), instrumentedUser);
619     }
620 
621     private void requireRunOnProfile(String userType,
622             OptionalBoolean installInstrumentedAppInParent,
623             boolean hasProfileOwner, boolean dpcIsPrimary, Set<String> affiliationIds) {
624         User instrumentedUser = sTestApis.users().instrumented().resolve();
625 
626         assumeTrue("This test only runs on users of type " + userType,
627                 instrumentedUser.type().name().equals(userType));
628 
629         if (!mProfiles.containsKey(instrumentedUser.type())) {
630             mProfiles.put(instrumentedUser.type(), new HashMap<>());
631         }
632 
633         mProfiles.get(instrumentedUser.type()).put(instrumentedUser.parent(), instrumentedUser);
634 
635         if (installInstrumentedAppInParent.equals(OptionalBoolean.TRUE)) {
636             sTestApis.packages().find(sContext.getPackageName()).install(
637                     instrumentedUser.parent());
638         } else if (installInstrumentedAppInParent.equals(OptionalBoolean.FALSE)) {
639             sTestApis.packages().find(sContext.getPackageName()).uninstall(
640                     instrumentedUser.parent());
641         }
642 
643         if (hasProfileOwner) {
644             ensureHasProfileOwner(instrumentedUser, dpcIsPrimary, affiliationIds);
645         } else {
646             ensureHasNoProfileOwner(instrumentedUser);
647         }
648     }
649 
650     private void requireFeature(String feature, FailureMode failureMode) {
651         checkFailOrSkip("Device must have feature " + feature,
652                 sTestApis.packages().features().contains(feature), failureMode);
653     }
654 
655     private void requireDoesNotHaveFeature(String feature, FailureMode failureMode) {
656         checkFailOrSkip("Device must not have feature " + feature,
657                 !sTestApis.packages().features().contains(feature), failureMode);
658     }
659 
660     private void requireNoGmsInstrumentation() {
661         boolean instrumentingGms =
662                 sTestApis.context().instrumentedContext().getPackageName().equals(GMS_PKG);
663 
664         checkFailOrSkip(
665                 "This test never runs using gms instrumentation",
666                 !instrumentingGms,
667                 FailureMode.SKIP
668         );
669     }
670 
671     private void requireGmsInstrumentation(int min, int max) {
672         mHasRequireGmsInstrumentation = true;
673         boolean instrumentingGms =
674                 sTestApis.context().instrumentedContext().getPackageName().equals(GMS_PKG);
675 
676         if (meetsSdkVersionRequirements(min, max)) {
677             checkFailOrSkip(
678                     "For SDK versions between " + min +  " and " + max
679                             + " (inclusive), this test only runs when using gms instrumentation",
680                     instrumentingGms,
681                     FailureMode.SKIP
682             );
683         } else {
684             checkFailOrSkip(
685                     "For SDK versions between " + min +  " and " + max
686                             + " (inclusive), this test only runs when not using gms "
687                             + "instrumentation",
688                     !instrumentingGms,
689                     FailureMode.SKIP
690             );
691         }
692     }
693 
694     private void requireSdkVersion(int min, int max, FailureMode failureMode) {
695         requireSdkVersion(min, max, failureMode,
696                 "Sdk version must be between " + min +  " and " + max + " (inclusive)");
697     }
698 
699     private void requireSdkVersion(
700             int min, int max, FailureMode failureMode, String failureMessage) {
701         checkFailOrSkip(
702                 failureMessage,
703                 meetsSdkVersionRequirements(min, max),
704                 failureMode
705         );
706     }
707 
708     private com.android.bedstead.nene.users.UserType requireUserSupported(
709             String userType, FailureMode failureMode) {
710         com.android.bedstead.nene.users.UserType resolvedUserType =
711                 sTestApis.users().supportedType(userType);
712 
713         checkFailOrSkip(
714                 "Device must support user type " + userType
715                 + " only supports: " + sTestApis.users().supportedTypes(),
716                 resolvedUserType != null, failureMode);
717 
718         return resolvedUserType;
719     }
720 
721     private void checkFailOrSkip(String message, boolean value, FailureMode failureMode) {
722         if (failureMode.equals(FailureMode.FAIL)) {
723             assertWithMessage(message).that(value).isTrue();
724         } else if (failureMode.equals(FailureMode.SKIP)) {
725             assumeTrue(message, value);
726         } else {
727             throw new IllegalStateException("Unknown failure mode: " + failureMode);
728         }
729     }
730 
731     private void failOrSkip(String message, FailureMode failureMode) {
732         if (failureMode.equals(FailureMode.FAIL)) {
733             throw new AssertionError(message);
734         } else if (failureMode.equals(FailureMode.SKIP)) {
735             throw new AssumptionViolatedException(message);
736         } else {
737             throw new IllegalStateException("Unknown failure mode: " + failureMode);
738         }
739     }
740 
741     public enum UserType {
742         /** Only to be used with annotations. */
743         ANY,
744         SYSTEM_USER,
745         CURRENT_USER,
746         PRIMARY_USER,
747         SECONDARY_USER,
748         WORK_PROFILE,
749         TV_PROFILE,
750     }
751 
752     private static final String LOG_TAG = "DeviceState";
753 
754     private static final Context sContext = sTestApis.context().instrumentedContext();
755 
756     private final Map<com.android.bedstead.nene.users.UserType, UserReference> mUsers =
757             new HashMap<>();
758     private final Map<com.android.bedstead.nene.users.UserType, Map<UserReference, UserReference>>
759             mProfiles = new HashMap<>();
760     private DevicePolicyController mDeviceOwner;
761     private Map<UserReference, DevicePolicyController> mProfileOwners = new HashMap<>();
762     private DevicePolicyController mPrimaryDpc;
763 
764     private final List<UserReference> mCreatedUsers = new ArrayList<>();
765     private final List<UserBuilder> mRemovedUsers = new ArrayList<>();
766     private final List<BlockingBroadcastReceiver> mRegisteredBroadcastReceivers = new ArrayList<>();
767     private boolean mHasChangedDeviceOwner = false;
768     private DevicePolicyController mOriginalDeviceOwner = null;
769     private Map<UserReference, DevicePolicyController> mChangedProfileOwners = new HashMap<>();
770 
771     /**
772      * Get the {@link UserReference} of the work profile for the current user.
773      *
774      * <p>If the current user is a work profile, then the current user will be returned.
775      *
776      * <p>This should only be used to get work profiles managed by Harrier (using either the
777      * annotations or calls to the {@link DeviceState} class.
778      *
779      * @throws IllegalStateException if there is no harrier-managed work profile
780      */
781     public UserReference workProfile() {
782         return workProfile(/* forUser= */ UserType.CURRENT_USER);
783     }
784 
785     /**
786      * Get the {@link UserReference} of the work profile.
787      *
788      * <p>This should only be used to get work profiles managed by Harrier (using either the
789      * annotations or calls to the {@link DeviceState} class.
790      *
791      * @throws IllegalStateException if there is no harrier-managed work profile for the given user
792      */
793     public UserReference workProfile(UserType forUser) {
794         return workProfile(resolveUserTypeToUser(forUser));
795     }
796 
797     /**
798      * Get the {@link UserReference} of the work profile.
799      *
800      * <p>This should only be used to get work profiles managed by Harrier (using either the
801      * annotations or calls to the {@link DeviceState} class.
802      *
803      * @throws IllegalStateException if there is no harrier-managed work profile for the given user
804      */
805     public UserReference workProfile(UserReference forUser) {
806         return profile(MANAGED_PROFILE_TYPE_NAME, forUser);
807     }
808 
809     /**
810      * Get the {@link UserReference} of the profile of the given type for the given user.
811      *
812      * <p>This should only be used to get profiles managed by Harrier (using either the
813      * annotations or calls to the {@link DeviceState} class.
814      *
815      * @throws IllegalStateException if there is no harrier-managed profile for the given user
816      */
817     public UserReference profile(String profileType, UserType forUser) {
818         return profile(profileType, resolveUserTypeToUser(forUser));
819     }
820 
821     /**
822      * Get the {@link UserReference} of the profile for the current user.
823      *
824      * <p>If the current user is a profile of the correct type, then the current user will be
825      * returned.
826      *
827      * <p>This should only be used to get profiles managed by Harrier (using either the
828      * annotations or calls to the {@link DeviceState} class.
829      *
830      * @throws IllegalStateException if there is no harrier-managed profile
831      */
832     public UserReference profile(String profileType) {
833         return profile(profileType, /* forUser= */ UserType.CURRENT_USER);
834     }
835 
836     /**
837      * Get the {@link UserReference} of the profile of the given type for the given user.
838      *
839      * <p>This should only be used to get profiles managed by Harrier (using either the
840      * annotations or calls to the {@link DeviceState} class.
841      *
842      * @throws IllegalStateException if there is no harrier-managed profile for the given user
843      */
844     public UserReference profile(String profileType, UserReference forUser) {
845         com.android.bedstead.nene.users.UserType resolvedUserType =
846                 sTestApis.users().supportedType(profileType);
847 
848         if (resolvedUserType == null) {
849             throw new IllegalStateException("Can not have a profile of type " + profileType
850                     + " as they are not supported on this device");
851         }
852 
853         return profile(resolvedUserType, forUser);
854     }
855 
856     /**
857      * Get the {@link UserReference} of the profile of the given type for the given user.
858      *
859      * <p>This should only be used to get profiles managed by Harrier (using either the
860      * annotations or calls to the {@link DeviceState} class.
861      *
862      * @throws IllegalStateException if there is no harrier-managed profile for the given user
863      */
864     public UserReference profile(
865             com.android.bedstead.nene.users.UserType userType, UserReference forUser) {
866         if (userType == null || forUser == null) {
867             throw new NullPointerException();
868         }
869 
870         if (!mProfiles.containsKey(userType) || !mProfiles.get(userType).containsKey(forUser)) {
871             UserReference parentUser = sTestApis.users().instrumented().resolve().parent();
872 
873             if (parentUser != null) {
874                 if (mProfiles.containsKey(userType)
875                         && mProfiles.get(userType).containsKey(parentUser)) {
876                     return mProfiles.get(userType).get(parentUser);
877                 }
878             }
879 
880             throw new IllegalStateException(
881                     "No harrier-managed profile of type " + userType + ". This method should only"
882                             + " be used when Harrier has been used to create the profile.");
883         }
884 
885         return mProfiles.get(userType).get(forUser);
886     }
887 
888     /**
889      * Get the {@link UserReference} of the tv profile for the current user
890      *
891      * <p>This should only be used to get tv profiles managed by Harrier (using either the
892      * annotations or calls to the {@link DeviceState} class.
893      *
894      * @throws IllegalStateException if there is no harrier-managed tv profile
895      */
896     public UserReference tvProfile() {
897         return tvProfile(/* forUser= */ UserType.CURRENT_USER);
898     }
899 
900     /**
901      * Get the {@link UserReference} of the tv profile.
902      *
903      * <p>This should only be used to get tv profiles managed by Harrier (using either the
904      * annotations or calls to the {@link DeviceState} class.
905      *
906      * @throws IllegalStateException if there is no harrier-managed tv profile
907      */
908     public UserReference tvProfile(UserType forUser) {
909         return tvProfile(resolveUserTypeToUser(forUser));
910     }
911 
912     /**
913      * Get the {@link UserReference} of the tv profile.
914      *
915      * <p>This should only be used to get tv profiles managed by Harrier (using either the
916      * annotations or calls to the {@link DeviceState} class.
917      *
918      * @throws IllegalStateException if there is no harrier-managed tv profile
919      */
920     public UserReference tvProfile(UserReference forUser) {
921         return profile(TV_PROFILE_TYPE_NAME, forUser);
922     }
923 
924     /**
925      * Get the user ID of the first human user on the device.
926      */
927     public UserReference primaryUser() {
928         return sTestApis.users().all()
929                 .stream().filter(User::isPrimary).findFirst()
930                 .orElseThrow(IllegalStateException::new);
931     }
932 
933     /**
934      * Get a secondary user.
935      *
936      * <p>This should only be used to get secondary users managed by Harrier (using either the
937      * annotations or calls to the {@link DeviceState} class.
938      *
939      * @throws IllegalStateException if there is no harrier-managed secondary user
940      */
941     public UserReference secondaryUser() {
942         return user(SECONDARY_USER_TYPE_NAME);
943     }
944 
945     /**
946      * Get a user of the given type.
947      *
948      * <p>This should only be used to get users managed by Harrier (using either the
949      * annotations or calls to the {@link DeviceState} class.
950      *
951      * @throws IllegalStateException if there is no harrier-managed user of the correct type
952      */
953     public UserReference user(String userType) {
954         com.android.bedstead.nene.users.UserType resolvedUserType =
955                 sTestApis.users().supportedType(userType);
956 
957         if (resolvedUserType == null) {
958             throw new IllegalStateException("Can not have a user of type " + userType
959                     + " as they are not supported on this device");
960         }
961 
962         return user(resolvedUserType);
963     }
964 
965     /**
966      * Get a user of the given type.
967      *
968      * <p>This should only be used to get users managed by Harrier (using either the
969      * annotations or calls to the {@link DeviceState} class.
970      *
971      * @throws IllegalStateException if there is no harrier-managed user of the correct type
972      */
973     public UserReference user(com.android.bedstead.nene.users.UserType userType) {
974         if (userType == null) {
975             throw new NullPointerException();
976         }
977 
978         if (!mUsers.containsKey(userType)) {
979             throw new IllegalStateException(
980                     "No harrier-managed user of type " + userType + ". This method should only be"
981                             + "used when Harrier has been used to create the user.");
982         }
983 
984         return mUsers.get(userType);
985     }
986 
987     private UserReference ensureHasProfile(
988             String profileType,
989             OptionalBoolean installInstrumentedApp,
990             UserType forUser,
991             boolean hasProfileOwner,
992             boolean profileOwnerIsPrimary) {
993         requireFeature("android.software.managed_users", FailureMode.SKIP);
994         com.android.bedstead.nene.users.UserType resolvedUserType =
995                 requireUserSupported(profileType, FailureMode.SKIP);
996 
997         UserReference forUserReference = resolveUserTypeToUser(forUser);
998 
999         UserReference profile =
1000                 sTestApis.users().findProfileOfType(resolvedUserType, forUserReference);
1001         if (profile == null) {
1002             profile = createProfile(resolvedUserType, forUserReference);
1003         }
1004 
1005         profile.start();
1006 
1007         if (installInstrumentedApp.equals(OptionalBoolean.TRUE)) {
1008             sTestApis.packages().find(sContext.getPackageName()).install(profile);
1009         } else if (installInstrumentedApp.equals(OptionalBoolean.FALSE)) {
1010             sTestApis.packages().find(sContext.getPackageName()).uninstall(profile);
1011         }
1012 
1013         if (!mProfiles.containsKey(resolvedUserType)) {
1014             mProfiles.put(resolvedUserType, new HashMap<>());
1015         }
1016 
1017         mProfiles.get(resolvedUserType).put(forUserReference, profile);
1018 
1019         if (hasProfileOwner) {
1020             ensureHasProfileOwner(profile, profileOwnerIsPrimary, /* affiliationIds= */ null);
1021         }
1022         return profile;
1023     }
1024 
1025     private void ensureHasNoProfile(String profileType, UserType forUser) {
1026         requireFeature("android.software.managed_users", FailureMode.SKIP);
1027 
1028         UserReference forUserReference = resolveUserTypeToUser(forUser);
1029         com.android.bedstead.nene.users.UserType resolvedProfileType =
1030                 sTestApis.users().supportedType(profileType);
1031 
1032         if (resolvedProfileType == null) {
1033             // These profile types don't exist so there can't be any
1034             return;
1035         }
1036 
1037         UserReference profile =
1038                 sTestApis.users().findProfileOfType(
1039                         resolvedProfileType,
1040                         forUserReference);
1041         if (profile != null) {
1042             removeAndRecordUser(profile);
1043         }
1044     }
1045 
1046     private void ensureHasUser(String userType, OptionalBoolean installInstrumentedApp) {
1047         com.android.bedstead.nene.users.UserType resolvedUserType =
1048                 requireUserSupported(userType, FailureMode.SKIP);
1049 
1050         Collection<UserReference> users = sTestApis.users().findUsersOfType(resolvedUserType);
1051 
1052         UserReference user = users.isEmpty() ? createUser(resolvedUserType)
1053                 : users.iterator().next();
1054 
1055         user.start();
1056 
1057         if (installInstrumentedApp.equals(OptionalBoolean.TRUE)) {
1058             sTestApis.packages().find(sContext.getPackageName()).install(user);
1059         } else if (installInstrumentedApp.equals(OptionalBoolean.FALSE)) {
1060             sTestApis.packages().find(sContext.getPackageName()).uninstall(user);
1061         }
1062 
1063         mUsers.put(resolvedUserType, user);
1064     }
1065 
1066     /**
1067      * Ensure that there is no user of the given type.
1068      */
1069     private void ensureHasNoUser(String userType) {
1070         com.android.bedstead.nene.users.UserType resolvedUserType =
1071                 sTestApis.users().supportedType(userType);
1072 
1073         if (resolvedUserType == null) {
1074             // These user types don't exist so there can't be any
1075             return;
1076         }
1077 
1078         for (UserReference secondaryUser : sTestApis.users().findUsersOfType(resolvedUserType)) {
1079             removeAndRecordUser(secondaryUser);
1080         }
1081     }
1082 
1083     private void removeAndRecordUser(UserReference userReference) {
1084         if (userReference == null) {
1085             return; // Nothing to remove
1086         }
1087 
1088         User user = userReference.resolve();
1089 
1090         if (!mCreatedUsers.remove(user)) {
1091             mRemovedUsers.add(sTestApis.users().createUser()
1092                     .name(user.name())
1093                     .type(user.type())
1094                     .parent(user.parent()));
1095         }
1096 
1097         user.remove();
1098     }
1099 
1100     public void requireCanSupportAdditionalUser() {
1101         int maxUsers = getMaxNumberOfUsersSupported();
1102         int currentUsers = sTestApis.users().all().size();
1103 
1104         assumeTrue("The device does not have space for an additional user (" + currentUsers +
1105                 " current users, " + maxUsers + " max users)", currentUsers + 1 <= maxUsers);
1106     }
1107 
1108     /**
1109      * Create and register a {@link BlockingBroadcastReceiver} which will be unregistered after the
1110      * test has run.
1111      */
1112     public BlockingBroadcastReceiver registerBroadcastReceiver(String action) {
1113         return registerBroadcastReceiver(action, /* checker= */ null);
1114     }
1115 
1116     /**
1117      * Create and register a {@link BlockingBroadcastReceiver} which will be unregistered after the
1118      * test has run.
1119      */
1120     public BlockingBroadcastReceiver registerBroadcastReceiver(
1121             String action, Function<Intent, Boolean> checker) {
1122         BlockingBroadcastReceiver broadcastReceiver =
1123                 new BlockingBroadcastReceiver(mContext, action, checker);
1124         broadcastReceiver.register();
1125         mRegisteredBroadcastReceivers.add(broadcastReceiver);
1126 
1127         return broadcastReceiver;
1128     }
1129 
1130     private UserReference resolveUserTypeToUser(UserType userType) {
1131         switch (userType) {
1132             case SYSTEM_USER:
1133                 return sTestApis.users().system();
1134             case CURRENT_USER:
1135                 return sTestApis.users().instrumented();
1136             case PRIMARY_USER:
1137                 return primaryUser();
1138             case SECONDARY_USER:
1139                 return secondaryUser();
1140             case WORK_PROFILE:
1141                 return workProfile();
1142             case TV_PROFILE:
1143                 return tvProfile();
1144             case ANY:
1145                 throw new IllegalStateException("ANY UserType can not be used here");
1146             default:
1147                 throw new IllegalArgumentException("Unknown user type " + userType);
1148         }
1149     }
1150 
1151     private void teardownNonShareableState() {
1152         mProfiles.clear();
1153         mUsers.clear();
1154 
1155         for (BlockingBroadcastReceiver broadcastReceiver : mRegisteredBroadcastReceivers) {
1156             broadcastReceiver.unregisterQuietly();
1157         }
1158         mRegisteredBroadcastReceivers.clear();
1159         mPrimaryDpc = null;
1160     }
1161 
1162     private void teardownShareableState() {
1163         if (mHasChangedDeviceOwner) {
1164             if (mOriginalDeviceOwner == null) {
1165                 if (mDeviceOwner != null) {
1166                     mDeviceOwner.remove();
1167                 }
1168             } else if (!mOriginalDeviceOwner.equals(mDeviceOwner)) {
1169                 mDeviceOwner.remove();
1170                 sTestApis.devicePolicy().setDeviceOwner(
1171                         mOriginalDeviceOwner.user(), mOriginalDeviceOwner.componentName());
1172             }
1173             mHasChangedDeviceOwner = false;
1174             mOriginalDeviceOwner = null;
1175         }
1176 
1177         for (Map.Entry<UserReference, DevicePolicyController> originalProfileOwner :
1178                 mChangedProfileOwners.entrySet()) {
1179 
1180             ProfileOwner currentProfileOwner =
1181                     sTestApis.devicePolicy().getProfileOwner(originalProfileOwner.getKey());
1182 
1183             if (Objects.equal(currentProfileOwner, originalProfileOwner.getValue())) {
1184                 continue; // No need to restore
1185             }
1186 
1187             if (currentProfileOwner != null) {
1188                 currentProfileOwner.remove();
1189             }
1190 
1191             if (originalProfileOwner.getValue() != null) {
1192                 sTestApis.devicePolicy().setProfileOwner(originalProfileOwner.getKey(),
1193                         originalProfileOwner.getValue().componentName());
1194             }
1195         }
1196         mChangedProfileOwners.clear();
1197 
1198         for (UserReference user : mCreatedUsers) {
1199             user.remove();
1200         }
1201 
1202         mCreatedUsers.clear();
1203 
1204         for (UserBuilder userBuilder : mRemovedUsers) {
1205             userBuilder.create();
1206         }
1207 
1208         mRemovedUsers.clear();
1209     }
1210 
1211     private UserReference createProfile(
1212             com.android.bedstead.nene.users.UserType profileType, UserReference parent) {
1213         requireCanSupportAdditionalUser();
1214         try {
1215             UserReference user = sTestApis.users().createUser()
1216                     .parent(parent)
1217                     .type(profileType)
1218                     .createAndStart();
1219             mCreatedUsers.add(user);
1220             return user;
1221         } catch (NeneException e) {
1222             throw new IllegalStateException("Error creating profile of type " + profileType, e);
1223         }
1224     }
1225 
1226     private UserReference createUser(com.android.bedstead.nene.users.UserType userType) {
1227         requireCanSupportAdditionalUser();
1228         try {
1229             UserReference user = sTestApis.users().createUser()
1230                     .type(userType)
1231                     .createAndStart();
1232             mCreatedUsers.add(user);
1233             return user;
1234         } catch (NeneException e) {
1235             throw new IllegalStateException("Error creating user of type " + userType, e);
1236         }
1237     }
1238 
1239     private int getMaxNumberOfUsersSupported() {
1240         try {
1241             return ShellCommand.builder("pm get-max-users")
1242                     .validate((output) -> output.startsWith("Maximum supported users:"))
1243                     .executeAndParseOutput(
1244                             (output) -> Integer.parseInt(output.split(": ", 2)[1].trim()));
1245         } catch (AdbException e) {
1246             throw new IllegalStateException("Invalid command output", e);
1247         }
1248     }
1249 
1250     private void ensureHasDeviceOwner(UserType onUser, FailureMode failureMode, boolean isPrimary, Set<String> affiliationIds) {
1251         // TODO(scottjonathan): Should support non-remotedpc device owner (default to remotedpc)
1252         // TODO(scottjonathan): Should allow setting the device owner on a different user
1253         UserReference userReference = resolveUserTypeToUser(onUser);
1254         if (isPrimary && mPrimaryDpc != null && !userReference.equals(mPrimaryDpc.user())) {
1255             throw new IllegalStateException("Only one DPC can be marked as primary per test (current primary is " + mPrimaryDpc + ")");
1256         }
1257         if (!userReference.equals(sTestApis.users().instrumented())) {
1258             // INTERACT_ACROSS_USERS_FULL is required for RemoteDPC
1259             ensureCanGetPermission(INTERACT_ACROSS_USERS_FULL);
1260         }
1261 
1262         DeviceOwner currentDeviceOwner = sTestApis.devicePolicy().getDeviceOwner();
1263 
1264         if (currentDeviceOwner != null
1265                 && currentDeviceOwner.componentName().equals(RemoteDpc.DPC_COMPONENT_NAME)) {
1266             mDeviceOwner = currentDeviceOwner;
1267         } else {
1268             UserReference instrumentedUser = sTestApis.users().instrumented();
1269 
1270             if (!Versions.meetsMinimumSdkVersionRequirement(Build.VERSION_CODES.S)) {
1271                 // Prior to S we can't set device owner if there are other users on the device
1272                 for (UserReference u : sTestApis.users().all()) {
1273                     if (u.equals(instrumentedUser)) {
1274                         continue;
1275                     }
1276                     try {
1277                         removeAndRecordUser(u);
1278                     } catch (NeneException e) {
1279                         failOrSkip(
1280                                 "Error removing user to prepare for DeviceOwner: " + e.toString(),
1281                                 failureMode);
1282                     }
1283                 }
1284             }
1285 
1286             // TODO(scottjonathan): Remove accounts
1287             ensureHasNoProfileOwner(userReference);
1288 
1289             if (!mHasChangedDeviceOwner) {
1290                 mOriginalDeviceOwner = currentDeviceOwner;
1291                 mHasChangedDeviceOwner = true;
1292             }
1293 
1294             mDeviceOwner = RemoteDpc.setAsDeviceOwner(userReference)
1295                     .devicePolicyController();
1296         }
1297 
1298         if (isPrimary) {
1299             mPrimaryDpc = mDeviceOwner;
1300         }
1301 
1302         RemoteDpc.forDevicePolicyController(mDeviceOwner).devicePolicyManager()
1303                 .setAffiliationIds(affiliationIds);
1304     }
1305 
1306     private void ensureHasProfileOwner(UserType onUser, boolean isPrimary, Set<String> affiliationIds) {
1307         // TODO(scottjonathan): Should support non-remotedpc profile owner (default to remotedpc)
1308         UserReference user = resolveUserTypeToUser(onUser);
1309         ensureHasProfileOwner(user, isPrimary, affiliationIds);
1310     }
1311 
1312     private void ensureHasProfileOwner(
1313             UserReference user, boolean isPrimary, Set<String> affiliationIds) {
1314         if (isPrimary && mPrimaryDpc != null && !user.equals(mPrimaryDpc.user())) {
1315             throw new IllegalStateException("Only one DPC can be marked as primary per test");
1316         }
1317 
1318         if (!user.equals(sTestApis.users().instrumented())) {
1319             // INTERACT_ACROSS_USERS_FULL is required for RemoteDPC
1320             ensureCanGetPermission(INTERACT_ACROSS_USERS_FULL);
1321         }
1322 
1323         ProfileOwner currentProfileOwner = sTestApis.devicePolicy().getProfileOwner(user);
1324         DeviceOwner currentDeviceOwner = sTestApis.devicePolicy().getDeviceOwner();
1325 
1326         if (currentDeviceOwner != null && currentDeviceOwner.user().equals(user)) {
1327             // Can't have DO and PO on the same user
1328             ensureHasNoDeviceOwner();
1329         }
1330 
1331         if (currentProfileOwner != null
1332                 && currentProfileOwner.componentName().equals(
1333                 RemoteDpc.DPC_COMPONENT_NAME)) {
1334             mProfileOwners.put(user, currentProfileOwner);
1335         } else {
1336             if (!mChangedProfileOwners.containsKey(user)) {
1337                 mChangedProfileOwners.put(user, currentProfileOwner);
1338             }
1339 
1340             mProfileOwners.put(user, RemoteDpc.setAsProfileOwner(user).devicePolicyController());
1341         }
1342 
1343         if (isPrimary) {
1344             mPrimaryDpc = mProfileOwners.get(user);
1345         }
1346 
1347         if (affiliationIds != null) {
1348             RemoteDpc profileOwner = profileOwner(user);
1349             profileOwner.devicePolicyManager()
1350                     .setAffiliationIds(affiliationIds);
1351         }
1352     }
1353 
1354     private void ensureHasNoDeviceOwner() {
1355         DeviceOwner deviceOwner = sTestApis.devicePolicy().getDeviceOwner();
1356 
1357         if (deviceOwner == null) {
1358             return;
1359         }
1360 
1361         if (!mHasChangedDeviceOwner) {
1362             mOriginalDeviceOwner = deviceOwner;
1363             mHasChangedDeviceOwner = true;
1364         }
1365 
1366         mDeviceOwner = null;
1367         deviceOwner.remove();
1368     }
1369 
1370     private void ensureHasNoProfileOwner(UserType onUser) {
1371         UserReference user = resolveUserTypeToUser(onUser);
1372 
1373         ensureHasNoProfileOwner(user);
1374     }
1375 
1376     private void ensureHasNoProfileOwner(UserReference user) {
1377 
1378         ProfileOwner currentProfileOwner = sTestApis.devicePolicy().getProfileOwner(user);
1379 
1380         if (currentProfileOwner == null) {
1381             return;
1382         }
1383 
1384         if (!mChangedProfileOwners.containsKey(user)) {
1385             mChangedProfileOwners.put(user, currentProfileOwner);
1386         }
1387 
1388         sTestApis.devicePolicy().getProfileOwner(user).remove();
1389         mProfileOwners.remove(user);
1390     }
1391 
1392     /**
1393      * Get the {@link RemoteDpc} for the device owner controlled by Harrier.
1394      *
1395      * <p>If no Harrier-managed device owner exists, an exception will be thrown.
1396      *
1397      * <p>If the device owner is not a RemoteDPC then an exception will be thrown
1398      */
1399     public RemoteDpc deviceOwner() {
1400         if (mDeviceOwner == null) {
1401             throw new IllegalStateException("No Harrier-managed device owner. This method should "
1402                     + "only be used when Harrier was used to set the Device Owner.");
1403         }
1404         if (!mDeviceOwner.componentName().equals(REMOTE_DPC_COMPONENT_NAME)) {
1405             throw new IllegalStateException("The device owner is not a RemoteDPC."
1406                     + " You must use Nene to query for this device owner.");
1407         }
1408 
1409         return RemoteDpc.forDevicePolicyController(mDeviceOwner);
1410     }
1411 
1412     /**
1413      * Get the {@link RemoteDpc} for the profile owner on the current user controlled by Harrier.
1414      *
1415      * <p>If no Harrier-managed profile owner exists, an exception will be thrown.
1416      *
1417      * <p>If the profile owner is not a RemoteDPC then an exception will be thrown.
1418      */
1419     public RemoteDpc profileOwner() {
1420         return profileOwner(UserType.CURRENT_USER);
1421     }
1422 
1423     /**
1424      * Get the {@link RemoteDpc} for the profile owner on the given user controlled by Harrier.
1425      *
1426      * <p>If no Harrier-managed profile owner exists, an exception will be thrown.
1427      *
1428      * <p>If the profile owner is not a RemoteDPC then an exception will be thrown.
1429      */
1430     public RemoteDpc profileOwner(UserType onUser) {
1431         if (onUser == null) {
1432             throw new NullPointerException();
1433         }
1434 
1435         return profileOwner(resolveUserTypeToUser(onUser));
1436     }
1437 
1438     /**
1439      * Get the {@link RemoteDpc} for the profile owner on the given user controlled by Harrier.
1440      *
1441      * <p>If no Harrier-managed profile owner exists, an exception will be thrown.
1442      *
1443      * <p>If the profile owner is not a RemoteDPC then an exception will be thrown.
1444      */
1445     public RemoteDpc profileOwner(UserReference onUser) {
1446         if (onUser == null) {
1447             throw new NullPointerException();
1448         }
1449 
1450         if (!mProfileOwners.containsKey(onUser)) {
1451             throw new IllegalStateException("No Harrier-managed profile owner. This method should "
1452                     + "only be used when Harrier was used to set the Profile Owner.");
1453         }
1454 
1455         DevicePolicyController profileOwner = mProfileOwners.get(onUser);
1456 
1457         if (!profileOwner.componentName().equals(REMOTE_DPC_COMPONENT_NAME)) {
1458             throw new IllegalStateException("The profile owner is not a RemoteDPC."
1459                     + " You must use Nene to query for this profile owner.");
1460         }
1461 
1462         return RemoteDpc.forDevicePolicyController(profileOwner);
1463     }
1464 
1465     private void requirePackageInstalled(
1466             String packageName, UserType forUser, FailureMode failureMode) {
1467 
1468         Package pkg = sTestApis.packages().find(packageName).resolve();
1469         checkFailOrSkip(
1470                 packageName + " is required to be installed for " + forUser,
1471                 pkg != null,
1472                 failureMode);
1473 
1474         if (forUser.equals(UserType.ANY)) {
1475             checkFailOrSkip(
1476                     packageName + " is required to be installed",
1477                     !pkg.installedOnUsers().isEmpty(),
1478                     failureMode);
1479         } else {
1480             checkFailOrSkip(
1481                     packageName + " is required to be installed for " + forUser,
1482                     pkg.installedOnUsers().contains(resolveUserTypeToUser(forUser)),
1483                     failureMode);
1484         }
1485     }
1486 
1487     private void requirePackageNotInstalled(
1488             String packageName, UserType forUser, FailureMode failureMode) {
1489         Package pkg = sTestApis.packages().find(packageName).resolve();
1490         if (pkg == null) {
1491             // Definitely not installed
1492             return;
1493         }
1494 
1495         if (forUser.equals(UserType.ANY)) {
1496             checkFailOrSkip(
1497                     packageName + " is required to be not installed",
1498                     pkg.installedOnUsers().isEmpty(),
1499                     failureMode);
1500         } else {
1501             checkFailOrSkip(
1502                     packageName + " is required to be not installed for " + forUser,
1503                     !pkg.installedOnUsers().contains(resolveUserTypeToUser(forUser)),
1504                     failureMode);
1505         }
1506     }
1507 
1508     private void ensurePackageNotInstalled(
1509             String packageName, UserType forUser) {
1510 
1511         Package pkg = sTestApis.packages().find(packageName).resolve();
1512         if (pkg == null) {
1513             // Definitely not installed
1514             return;
1515         }
1516 
1517         if (forUser.equals(UserType.ANY)) {
1518             if (!pkg.installedOnUsers().isEmpty()) {
1519                 pkg.uninstallFromAllUsers();
1520             }
1521         } else {
1522             UserReference user = resolveUserTypeToUser(forUser);
1523             if (pkg.installedOnUsers().contains(user)) {
1524                 pkg.uninstall(user);
1525             }
1526         }
1527     }
1528 
1529     /**
1530      * Get the most appropriate {@link RemoteDpc} instance for the device state.
1531      *
1532      * <p>This method should only be used by tests which are annotated with {@link PolicyTest}.
1533      *
1534      * <p>If no DPC is set as the "primary" DPC for the device state, then this method will first
1535      * check for a profile owner in the current user, or else check for a device owner.
1536      *
1537      * <p>If no Harrier-managed profile owner or device owner exists, an exception will be thrown.
1538      *
1539      * <p>If the profile owner or device owner is not a RemoteDPC then an exception will be thrown.
1540      */
1541     public RemoteDpc dpc() {
1542         if (mPrimaryDpc != null) {
1543             return RemoteDpc.forDevicePolicyController(mPrimaryDpc);
1544         }
1545 
1546         if (mProfileOwners.containsKey(sTestApis.users().instrumented())) {
1547             DevicePolicyController profileOwner =
1548                     mProfileOwners.get(sTestApis.users().instrumented());
1549 
1550 
1551             if (profileOwner.componentName().equals(REMOTE_DPC_COMPONENT_NAME)) {
1552                 return RemoteDpc.forDevicePolicyController(profileOwner);
1553             }
1554         }
1555 
1556         if (mDeviceOwner != null) {
1557             if (mDeviceOwner.componentName().equals(REMOTE_DPC_COMPONENT_NAME)) {
1558                 return RemoteDpc.forDevicePolicyController(mDeviceOwner);
1559             }
1560 
1561         }
1562 
1563         throw new IllegalStateException("No Harrier-managed profile owner or device owner.");
1564     }
1565 
1566     private void ensureCanGetPermission(String permission) {
1567         // TODO(scottjonathan): Apply gms permission switches automatically rather than hard-coding
1568         // TODO(scottjonathan): Add a config to only enforce gms permission when needed
1569         if (permission.equals(NOTIFY_PENDING_SYSTEM_UPDATE)) {
1570             requireGmsInstrumentation(1, Build.VERSION_CODES.R);
1571         }
1572         // TODO(scottjonathan): Apply version-specific constraints automatically
1573         if (permission.equals(INTERACT_ACROSS_USERS_FULL)) {
1574             requireSdkVersion(
1575                     Build.VERSION_CODES.Q, Integer.MAX_VALUE, FailureMode.SKIP,
1576                     "This test requires INTERACT_ACROSS_USERS_FULL which can only be used on Q+");
1577         }
1578     }
1579 }
1580