2 files changed, 60 insertions, 43 deletions

diff --git a/src/io/appium/droiddriver/util/ActivityUtils.java b/src/io/appium/droiddriver/util/ActivityUtils.java
index 248f6b3..ff06ab5 100644
--- a/src/io/appium/droiddriver/util/ActivityUtils.java
+++ b/src/io/appium/droiddriver/util/ActivityUtils.java
@@ -17,37 +17,54 @@
 package io.appium.droiddriver.util;
 
 import android.app.Activity;
-import io.appium.droiddriver.exceptions.UnrecoverableException;
+import android.os.Looper;
+import android.support.test.runner.lifecycle.ActivityLifecycleMonitorRegistry;
+import android.support.test.runner.lifecycle.Stage;
+import android.util.Log;
+import java.util.Iterator;
+import java.util.concurrent.Callable;
 
-/**
- * Static helper methods for retrieving activities.
- */
+/** Static helper methods for retrieving activities. */
 public class ActivityUtils {
-  public interface Supplier<T> {
-    /**
-     * Retrieves an instance of the appropriate type. The returned object may or
-     * may not be a new instance, depending on the implementation.
-     *
-     * @return an instance of the appropriate type
-     */
-    T get();
-  }
+  private static final Callable<Activity> GET_RUNNING_ACTIVITY =
+      new Callable<Activity>() {
+        @Override
+        public Activity call() {
+          Iterator<Activity> activityIterator =
+              ActivityLifecycleMonitorRegistry.getInstance()
+                  .getActivitiesInStage(Stage.RESUMED)
+                  .iterator();
+          return activityIterator.hasNext() ? activityIterator.next() : null;
+        }
+      };
+  private static Supplier<Activity> runningActivitySupplier =
+      new Supplier<Activity>() {
+        @Override
+        public Activity get() {
+          try {
+            // If this is called on main (UI) thread, don't call runOnMainSync
+            if (Looper.myLooper() == Looper.getMainLooper()) {
+              return GET_RUNNING_ACTIVITY.call();
+            }
 
-  private static Supplier<Activity> runningActivitySupplier;
+            return InstrumentationUtils.runOnMainSyncWithTimeout(GET_RUNNING_ACTIVITY);
+          } catch (Exception e) {
+            Logs.log(Log.WARN, e);
+            return null;
+          }
+        }
+      };
 
   /**
-   * Sets the Supplier for the running (a.k.a. resumed or foreground) activity.
-   * Called from {@link io.appium.droiddriver.runner.TestRunner}. If a
-   * custom runner is used, this method must be called appropriately, otherwise
-   * {@link #getRunningActivity} won't work.
+   * Sets the Supplier for the running (a.k.a. resumed or foreground) activity. If a custom runner
+   * is used, this method must be called appropriately, otherwise {@link #getRunningActivity} won't
+   * work.
    */
   public static synchronized void setRunningActivitySupplier(Supplier<Activity> activitySupplier) {
-    runningActivitySupplier = activitySupplier;
+    runningActivitySupplier = Preconditions.checkNotNull(activitySupplier);
   }
 
-  /**
-   * Shorthand to {@link #getRunningActivity(long)} with {@code timeoutMillis=30_000}.
-   */
+  /** Shorthand to {@link #getRunningActivity(long)} with {@code timeoutMillis=30_000}. */
   public static Activity getRunningActivity() {
     return getRunningActivity(30_000L);
   }
@@ -74,11 +91,16 @@ public class ActivityUtils {
    * @return the currently running activity, or null if no activity has focus.
    */
   public static synchronized Activity getRunningActivityNoWait() {
-    if (runningActivitySupplier == null) {
-      throw new UnrecoverableException(
-          "If you don't use DroidDriver TestRunner, you need to call"
-              + " ActivityUtils.setRunningActivitySupplier appropriately");
-    }
     return runningActivitySupplier.get();
   }
+
+  public interface Supplier<T> {
+    /**
+     * Retrieves an instance of the appropriate type. The returned object may or may not be a new
+     * instance, depending on the implementation.
+     *
+     * @return an instance of the appropriate type
+     */
+    T get();
+  }
 }
diff --git a/src/io/appium/droiddriver/util/InstrumentationUtils.java b/src/io/appium/droiddriver/util/InstrumentationUtils.java
index 0ca087a..06ac5ab 100644
--- a/src/io/appium/droiddriver/util/InstrumentationUtils.java
+++ b/src/io/appium/droiddriver/util/InstrumentationUtils.java
@@ -20,10 +20,10 @@ import android.app.Instrumentation;
 import android.content.Context;
 import android.os.Bundle;
 import android.os.Looper;
+import android.support.test.InstrumentationRegistry;
 import android.util.Log;
 import io.appium.droiddriver.exceptions.DroidDriverException;
 import io.appium.droiddriver.exceptions.TimeoutException;
-import io.appium.droiddriver.exceptions.UnrecoverableException;
 import java.util.concurrent.Callable;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
@@ -43,11 +43,11 @@ public class InstrumentationUtils {
   private static long runOnMainSyncTimeoutMillis;
 
   /**
-   * Initializes this class. If you use a runner that is not DroidDriver-aware, you need to call
-   * this method appropriately. See {@link io.appium.droiddriver.runner.TestRunner#onCreate} for
-   * example.
+   * Initializes this class. If you don't use android.support.test.runner.AndroidJUnitRunner or a
+   * runner that supports {link InstrumentationRegistry}, you need to call this method
+   * appropriately.
    */
-  public static void init(Instrumentation instrumentation, Bundle arguments) {
+  public static synchronized void init(Instrumentation instrumentation, Bundle arguments) {
     if (InstrumentationUtils.instrumentation != null) {
       throw new DroidDriverException("init() can only be called once");
     }
@@ -55,14 +55,13 @@ public class InstrumentationUtils {
     options = arguments;
 
     String timeoutString = getD2Option("runOnMainSyncTimeout");
-    runOnMainSyncTimeoutMillis = timeoutString == null ? 10000L : Long.parseLong(timeoutString);
+    runOnMainSyncTimeoutMillis = timeoutString == null ? 10_000L : Long.parseLong(timeoutString);
   }
 
-  private static void checkInitialized() {
+  private static synchronized void checkInitialized() {
     if (instrumentation == null) {
-      throw new UnrecoverableException(
-          "If you use a runner that is not DroidDriver-aware, you"
-              + " need to call InstrumentationUtils.init appropriately");
+      // Assume android.support.test.runner.InstrumentationRegistry is valid
+      init(InstrumentationRegistry.getInstrumentation(), InstrumentationRegistry.getArguments());
     }
   }
 
@@ -85,11 +84,7 @@ public class InstrumentationUtils {
     return options;
   }
 
-  /**
-   * Gets the string value associated with the given key. This is preferred over using {@link
-   * #getOptions} because the returned {@link Bundle} contains only string values - am instrument
-   * options do not support value types other than string.
-   */
+  /** Gets the string value associated with the given key. */
   public static String getOption(String key) {
     return getOptions().getString(key);
   }
@@ -110,7 +105,7 @@ public class InstrumentationUtils {
   public static boolean tryWaitForIdleSync(long timeoutMillis) {
     checkNotMainThread();
     FutureTask<Void> emptyTask = new FutureTask<Void>(EMPTY_RUNNABLE, null);
-    instrumentation.waitForIdle(emptyTask);
+    getInstrumentation().waitForIdle(emptyTask);
 
     try {
       emptyTask.get(timeoutMillis, TimeUnit.MILLISECONDS);