1 files changed, 49 insertions, 16 deletions
diff --git a/test-runner/src/android/test/suitebuilder/TestPredicates.java b/test-runner/src/android/test/suitebuilder/TestPredicates.java
index 47aca55f9972..616d1a972612 100644
--- a/test-runner/src/android/test/suitebuilder/TestPredicates.java
+++ b/test-runner/src/android/test/suitebuilder/TestPredicates.java
@@ -17,30 +17,63 @@
 package android.test.suitebuilder;
 
 import android.test.InstrumentationTestCase;
-import android.test.suitebuilder.annotation.HasAnnotation;
-import android.test.suitebuilder.annotation.Suppress;
-import android.test.suitebuilder.annotation.LargeTest;
-import android.test.suitebuilder.annotation.MediumTest;
-import android.test.suitebuilder.annotation.SmallTest;
 import android.test.suitebuilder.annotation.Smoke;
+import android.test.suitebuilder.annotation.Suppress;
 import com.android.internal.util.Predicate;
-import com.android.internal.util.Predicates;
+import java.lang.annotation.Annotation;
 
 /**
  * {@hide} Not needed for 1.0 SDK.
  */
 public class TestPredicates {
 
-    public static final Predicate<TestMethod> SELECT_INSTRUMENTATION =
-            new AssignableFrom(InstrumentationTestCase.class);
-    public static final Predicate<TestMethod> REJECT_INSTRUMENTATION =
-            Predicates.not(SELECT_INSTRUMENTATION);
+    static final Predicate<TestMethod> REJECT_INSTRUMENTATION =
+            not(new AssignableFrom(InstrumentationTestCase.class));
+
+    static final Predicate<TestMethod> SELECT_SMOKE = hasAnnotation(Smoke.class);
+
+    static final Predicate<TestMethod> REJECT_SUPPRESSED = not(hasAnnotation(Suppress.class));
+
+    /**
+     * Return a predicate that checks to see if a {@link TestMethod} has an instance of the supplied
+     * annotation class, either on the method or on the containing class.
+     */
+    public static Predicate<TestMethod> hasAnnotation(Class<? extends Annotation> annotationClass) {
+        return new HasAnnotation(annotationClass);
+    }
+
+    private static class HasAnnotation implements Predicate<TestMethod> {
+
+        private final Class<? extends Annotation> annotationClass;
+
+        private HasAnnotation(Class<? extends Annotation> annotationClass) {
+            this.annotationClass = annotationClass;
+        }
+
+        @Override
+        public boolean apply(TestMethod testMethod) {
+            return testMethod.getAnnotation(annotationClass) != null ||
+                    testMethod.getEnclosingClass().getAnnotation(annotationClass) != null;
+        }
+    }
+
+    /**
+     * Returns a Predicate that evaluates to true iff the given Predicate
+     * evaluates to false.
+     */
+    public static <T> Predicate<T> not(Predicate<? super T> predicate) {
+        return new NotPredicate<T>(predicate);
+    }
+
+    private static class NotPredicate<T> implements Predicate<T> {
+        private final Predicate<? super T> predicate;
 
-    public static final Predicate<TestMethod> SELECT_SMOKE = new HasAnnotation(Smoke.class);
-    public static final Predicate<TestMethod> SELECT_SMALL = new HasAnnotation(SmallTest.class);
-    public static final Predicate<TestMethod> SELECT_MEDIUM = new HasAnnotation(MediumTest.class);
-    public static final Predicate<TestMethod> SELECT_LARGE = new HasAnnotation(LargeTest.class);
-    public static final Predicate<TestMethod> REJECT_SUPPRESSED =
-            Predicates.not(new HasAnnotation(Suppress.class));
+        private NotPredicate(Predicate<? super T> predicate) {
+            this.predicate = predicate;
+        }
 
+        public boolean apply(T t) {
+            return !predicate.apply(t);
+        }
+    }
 }