package com.squareup.instrumentation;

import android.app.Instrumentation;
import android.test.AndroidTestRunner;
import com.squareup.ShowTabletUi;
import com.squareup.instrumentation.constraints.PhoneOnly;
import com.squareup.instrumentation.constraints.TabletOnly;
import java.lang.reflect.Method;
import java.util.Enumeration;
import junit.framework.Test;
import junit.framework.TestCase;
import junit.framework.TestSuite;

import static com.squareup.instrumentation.TestSuiteUtils.maybeExplodeTest;

/**
 * Our own test runner which explodes tests annotated with @Variations.
 */
public class ExplodingTestRunner extends AndroidTestRunner {
  private final Instrumentation instrumentation;

  public ExplodingTestRunner(Instrumentation instrumentation) {
    this.instrumentation = instrumentation;
  }

  @Override
  public void setInstrumentation(Instrumentation instr) {
    super.setInstrumentation(instrumentation);
  }

  @Override
  public void setInstrumentaiton(Instrumentation instr) {
    super.setInstrumentation(instrumentation);
  }

  @Override public void setTest(Test test) {
    if (!(test instanceof TestSuite)) {
      throw new UnsupportedOperationException("Expected instance of TestSuite.");
    }

    TestSuite godTestSuite = new TestSuite();
    try {
      explodeSuite((TestSuite) test, godTestSuite);
    } catch (Exception e) {
      throw new RuntimeException(e);
    }
    super.setTest(godTestSuite);
  }

  private void explodeSuite(TestSuite testSuite, TestSuite result) throws Exception {
    ClassLoader classLoader = instrumentation.getTargetContext().getClassLoader();
    Class<?> testCaseClass = classLoader.loadClass(testSuite.getName());
    Enumeration<Test> testEnumerator = testSuite.tests();
    while (testEnumerator.hasMoreElements()) {
      Test test = testEnumerator.nextElement();
      if (test instanceof TestCase) {
        TestCase testCase = (TestCase) test;
        Method method = testCaseClass.getMethod(testCase.getName());
        if (!shouldSkipMethod(method)) {
          result.addTest(maybeExplodeTest(testCaseClass, method));
        }
      } else {
        explodeSuite((TestSuite) test, result); // Recursively explode this suite's tests.
      }
    }
  }

  private boolean shouldSkipMethod(Method method) {
    boolean tablet = instrumentation.getTargetContext().getResources().getBoolean(ShowTabletUi.ID);
    return method.isAnnotationPresent(PhoneOnly.class) && tablet
        || method.isAnnotationPresent(TabletOnly.class) && !tablet;
  }
}

@@ -2300,100 +2300,160 @@
 blet
- = instrumentation.getTargetContext().getResources().getBoolean(ShowTabletUi.ID);%0A    return
+Only = method.isAnnotationPresent(TabletOnly.class)%0A        %7C%7C method.getDeclaringClass().isAnnotationPresent(TabletOnly.class);%0A    boolean phoneOnly =
  met
@@ -2496,18 +2496,8 @@
 ass)
- && tablet
 %0A   
@@ -2503,32 +2503,52 @@
       %7C%7C method.
+getDeclaringClass().
 isAnnotationPres
@@ -2543,38 +2543,37 @@
 notationPresent(
-Tablet
+Phone
 Only.class) && !
@@ -2567,16 +2567,165 @@
 y.class)
+;%0A    boolean tablet = instrumentation.getTargetContext().getResources().getBoolean(ShowTabletUi.ID);%0A    return (phoneOnly && tablet) %7C%7C (tabletOnly
  && !tab
@@ -2727,16 +2727,17 @@
  !tablet
+)
 ;%0A  %7D%0A%7D%0A

Allow TabletOnly etc on whole test classes