Update instrumentation tests to use PreTestHooks for policies

components/policy/android/javatests/src/org/chromium/policy/test/annotations/Policies.java

Index: components/policy/android/javatests/src/org/chromium/policy/test/annotations/Policies.java
diff --git a/components/policy/android/javatests/src/org/chromium/policy/test/annotations/Policies.java b/components/policy/android/javatests/src/org/chromium/policy/test/annotations/Policies.java
new file mode 100644
index 0000000000000000000000000000000000000000..62a79fd0f99c3ef1092c2774ec344b4167a628ef
--- /dev/null
+++ b/components/policy/android/javatests/src/org/chromium/policy/test/annotations/Policies.java
@@ -0,0 +1,150 @@
+// Copyright 2015 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+package org.chromium.policy.test.annotations;
+
+import android.content.Context;
+import android.os.Bundle;
+
+import junit.framework.Assert;
+
+import org.chromium.base.test.BaseTestResult.PreTestHook;
+import org.chromium.policy.AbstractAppRestrictionsProvider;
+import org.chromium.policy.test.PolicyData;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Inherited;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+import java.lang.reflect.AnnotatedElement;
+import java.lang.reflect.Method;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+/**
+ * Annotations and utilities for testing code dependent on policies.
+ *
+ * Usage example:
+ * <pre>
+ * @Policies.Add({
+ *     @Policies.Item(key="Foo", type=Policies.Type.STRING, string="Bar"),
+ *     @Policies.Item(key="Baz", type=Policies.Type.STRING, string="Baz")
+ * })
+ * public class MyTestClass extends BaseActivityInstrumentationTestCase<ContentActivity> {
+ *
+ *     public void MyTest1() {
+ *         // Will run the Foo and Bar policies set
+ *     }
+ *
+ *     @Policies.Remove(@Policies.Item(key="Baz"))
+ *     public void MyTest2() {
+ *         // Will run with only the Foo policy set
+ *     }
+ * }
+ * </pre>
+ */
+public final class Policies {
+    public enum Type { UNDEFINED, STRING, STRING_ARRAY }
+
+    /** Items declared here will be added to the list of used policies. */
+    @Inherited
+    @Retention(RetentionPolicy.RUNTIME)
+    @Target({ElementType.METHOD, ElementType.TYPE})
+    public @interface Add {
+        Item[] value();
+    }
+
+    /** Items declared here will be removed from the list of used policies. */
+    @Inherited
+    @Retention(RetentionPolicy.RUNTIME)
+    @Target({ElementType.METHOD, ElementType.TYPE})
+    public @interface Remove {
+        Item[] value();
+    }
+
+    /**
+     * Individual policy item. Identified by a {@link #key}, and optional data {@link #type}s and
+     * values. When it is parsed, only the value associated to the declared type will be
+     * used.
+     */
+    @Inherited
+    @Retention(RetentionPolicy.RUNTIME)
+    @Target({ElementType.METHOD, ElementType.TYPE})
+    public @interface Item {
+        String key();
+
+        Type type() default Type.UNDEFINED;
+
+        String string() default "";
+
+        String[] stringArray() default {};
+    }
+
+    /** Parses the annotations to extract usable information as {@link PolicyData} objects. */
+    private static List<PolicyData> fromItems(Item[] items) {
+        List<PolicyData> result = new ArrayList<>();
+        for (Item item : items) {
+            PolicyData data;
+            switch (item.type()) {
+                case STRING:
+                    data = new PolicyData.Str(item.key(), item.string());
+                    break;
+                case STRING_ARRAY:
+                    data = new PolicyData.StrArray(item.key(), item.stringArray());
+                    break;
+                default:
+                    data = new PolicyData.Undefined(item.key());
+            }
+            result.add(data);
+        }
+        return result;
+    }
+
+    /** @see PreTestHook */
+    public static PreTestHook getRegistrationHook() {
+        return new RegistrationHook();
+    }
+
+    private static Set<PolicyData> getPolicies(AnnotatedElement element) {
+        AnnotatedElement parent = (element instanceof Method)
+                ? ((Method) element).getDeclaringClass()
+                : ((Class<?>) element).getSuperclass();
+        Set<PolicyData> flags = (parent == null) ? new HashSet<PolicyData>() : getPolicies(parent);
+
+        if (element.isAnnotationPresent(Policies.Add.class)) {
+            flags.addAll(fromItems(element.getAnnotation(Policies.Add.class).value()));
+        }
+
+        if (element.isAnnotationPresent(Policies.Remove.class)) {
+            flags.removeAll(fromItems(element.getAnnotation(Policies.Remove.class).value()));
+        }
+
+        return flags;
+    }
+
+    /**
+     * Registration hook for the {@link Policies} annotation family. Before a test, will parse
+     * the declared policies and use them as cached policies.
+     */
+    public static class RegistrationHook implements PreTestHook {
+        @Override
+        public void run(Context targetContext, Method testMethod) {
+
+            Assert.assertNotNull("Unable to get a non-null target context.", targetContext);
+            Context applicationContext = targetContext.getApplicationContext();
+
+            final Bundle policyBundle = PolicyData.asBundle(getPolicies(testMethod));
+
+            if (!policyBundle.isEmpty()) {
+                policyBundle.putBoolean(AbstractAppRestrictionsProvider.STICKY_CACHE_KEY, true);
+            }
+
+            AbstractAppRestrictionsProvider.setCachedPoliciesForTesting(
+                    applicationContext, policyBundle);
+        }
+    }
+}