Check running activity again on main thread during findRootView

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=136169242

3 files changed, 124 insertions, 89 deletions

diff --git a/src/io/appium/droiddriver/instrumentation/InstrumentationDriver.java b/src/io/appium/droiddriver/instrumentation/InstrumentationDriver.java
index d3e5dd2..03b2123 100644
--- a/src/io/appium/droiddriver/instrumentation/InstrumentationDriver.java
+++ b/src/io/appium/droiddriver/instrumentation/InstrumentationDriver.java
@@ -16,14 +16,11 @@
 
 package io.appium.droiddriver.instrumentation;
 
+import android.app.Activity;
 import android.app.Instrumentation;
 import android.os.SystemClock;
 import android.util.Log;
 import android.view.View;
-
-import java.util.List;
-import java.util.concurrent.Callable;
-
 import io.appium.droiddriver.actions.InputInjector;
 import io.appium.droiddriver.base.BaseDroidDriver;
 import io.appium.droiddriver.base.DroidDriverContext;
@@ -31,11 +28,36 @@ import io.appium.droiddriver.exceptions.NoRunningActivityException;
 import io.appium.droiddriver.util.ActivityUtils;
 import io.appium.droiddriver.util.InstrumentationUtils;
 import io.appium.droiddriver.util.Logs;
+import java.util.List;
+import java.util.concurrent.Callable;
 
-/**
- * Implementation of DroidDriver that is driven via instrumentation.
- */
+/** Implementation of DroidDriver that is driven via instrumentation. */
 public class InstrumentationDriver extends BaseDroidDriver<View, ViewElement> {
+  private static final Callable<View> FIND_ROOT_VIEW =
+      new Callable<View>() {
+        @Override
+        public View call() {
+          InstrumentationUtils.checkMainThread();
+          Activity runningActivity = ActivityUtils.getRunningActivityNoWait();
+          if (runningActivity == null) {
+            // runningActivity changed since last call!
+            return null;
+          }
+
+          List<View> views = RootFinder.getRootViews();
+          if (views.size() > 1) {
+            Logs.log(Log.VERBOSE, "views.size()=" + views.size());
+            for (View view : views) {
+              if (view.hasWindowFocus()) {
+                return view;
+              }
+            }
+          }
+          // Fall back to DecorView.
+          // TODO(kjin): Should wait until a view hasWindowFocus?
+          return runningActivity.getWindow().getDecorView();
+        }
+      };
   private final DroidDriverContext<View, ViewElement> context;
   private final InputInjector injector;
   private final InstrumentationUiDevice uiDevice;
@@ -61,41 +83,22 @@ public class InstrumentationDriver extends BaseDroidDriver<View, ViewElement> {
     return new ViewElement(context, rawElement, parent);
   }
 
-  private static final Callable<View> FIND_ROOT_VIEW = new Callable<View>() {
-    @Override
-    public View call() {
-      List<View> views = RootFinder.getRootViews();
-      if (views.size() > 1) {
-        Logs.log(Log.VERBOSE, "views.size()=" + views.size());
-        for (View view : views) {
-          if (view.hasWindowFocus()) {
-            return view;
-          }
-        }
-      }
-      // Fall back to DecorView.
-      return ActivityUtils.getRunningActivity().getWindow().getDecorView();
-    }
-  };
-
   private View findRootView() {
-    waitForRunningActivity();
-    return InstrumentationUtils.runOnMainSyncWithTimeout(FIND_ROOT_VIEW);
-  }
-
-  private void waitForRunningActivity() {
     long timeoutMillis = getPoller().getTimeoutMillis();
     long end = SystemClock.uptimeMillis() + timeoutMillis;
     while (true) {
-      if (ActivityUtils.getRunningActivity() != null) {
-        return;
-      }
       long remainingMillis = end - SystemClock.uptimeMillis();
       if (remainingMillis < 0) {
-        throw new NoRunningActivityException(String.format(
-            "Cannot find the running activity after %d milliseconds", timeoutMillis));
+        throw new NoRunningActivityException(
+            String.format("Cannot find the running activity after %d milliseconds", timeoutMillis));
+      }
+
+      if (ActivityUtils.getRunningActivity(remainingMillis) != null) {
+        View view = InstrumentationUtils.runOnMainSyncWithTimeout(FIND_ROOT_VIEW);
+        if (view != null) {
+          return view;
+        }
       }
-      SystemClock.sleep(Math.min(250, remainingMillis));
     }
   }
 
diff --git a/src/io/appium/droiddriver/util/ActivityUtils.java b/src/io/appium/droiddriver/util/ActivityUtils.java
index 1e35de8..248f6b3 100644
--- a/src/io/appium/droiddriver/util/ActivityUtils.java
+++ b/src/io/appium/droiddriver/util/ActivityUtils.java
@@ -17,9 +17,7 @@
 package io.appium.droiddriver.util;
 
 import android.app.Activity;
-
 import io.appium.droiddriver.exceptions.UnrecoverableException;
-import io.appium.droiddriver.instrumentation.InstrumentationDriver;
 
 /**
  * Static helper methods for retrieving activities.
@@ -48,15 +46,38 @@ public class ActivityUtils {
   }
 
   /**
-   * Gets the running (a.k.a. resumed or foreground) activity.
-   * {@link InstrumentationDriver} depends on this.
+   * Shorthand to {@link #getRunningActivity(long)} with {@code timeoutMillis=30_000}.
+   */
+  public static Activity getRunningActivity() {
+    return getRunningActivity(30_000L);
+  }
+
+  /**
+   * Waits for idle on main looper, then gets the running (a.k.a. resumed or foreground) activity.
+   *
+   * @return the currently running activity, or null if no activity has focus.
+   */
+  public static Activity getRunningActivity(long timeoutMillis) {
+    // It's safe to check running activity only when the main looper is idle.
+    // If the AUT is in background, its main looper should be idle already.
+    // If the AUT is in foreground, its main looper should be idle eventually.
+    if (InstrumentationUtils.tryWaitForIdleSync(timeoutMillis)) {
+      return getRunningActivityNoWait();
+    }
+    return null;
+  }
+
+  /**
+   * Gets the running (a.k.a. resumed or foreground) activity without waiting for idle on main
+   * looper.
    *
    * @return the currently running activity, or null if no activity has focus.
    */
-  public static synchronized Activity getRunningActivity() {
+  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");
+      throw new UnrecoverableException(
+          "If you don't use DroidDriver TestRunner, you need to call"
+              + " ActivityUtils.setRunningActivitySupplier appropriately");
     }
     return runningActivitySupplier.get();
   }
diff --git a/src/io/appium/droiddriver/util/InstrumentationUtils.java b/src/io/appium/droiddriver/util/InstrumentationUtils.java
index c4f280d..0ca087a 100644
--- a/src/io/appium/droiddriver/util/InstrumentationUtils.java
+++ b/src/io/appium/droiddriver/util/InstrumentationUtils.java
@@ -21,30 +21,26 @@ import android.content.Context;
 import android.os.Bundle;
 import android.os.Looper;
 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;
 import java.util.concurrent.FutureTask;
 import java.util.concurrent.TimeUnit;
 
-import io.appium.droiddriver.exceptions.DroidDriverException;
-import io.appium.droiddriver.exceptions.TimeoutException;
-import io.appium.droiddriver.exceptions.UnrecoverableException;
-
-/**
- * Static utility methods pertaining to {@link Instrumentation}.
- */
+/** Static utility methods pertaining to {@link Instrumentation}. */
 public class InstrumentationUtils {
+  private static final Runnable EMPTY_RUNNABLE =
+      new Runnable() {
+        @Override
+        public void run() {}
+      };
+  private static final Executor RUN_ON_MAIN_SYNC_EXECUTOR = Executors.newSingleThreadExecutor();
   private static Instrumentation instrumentation;
   private static Bundle options;
   private static long runOnMainSyncTimeoutMillis;
-  private static final Runnable EMPTY_RUNNABLE = new Runnable() {
-    @Override
-    public void run() {
-    }
-  };
-  private static final Executor RUN_ON_MAIN_SYNC_EXECUTOR = Executors.newSingleThreadExecutor();
 
   /**
    * Initializes this class. If you use a runner that is not DroidDriver-aware, you need to call
@@ -64,8 +60,9 @@ public class InstrumentationUtils {
 
   private static 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");
+      throw new UnrecoverableException(
+          "If you use a runner that is not DroidDriver-aware, you"
+              + " need to call InstrumentationUtils.init appropriately");
     }
   }
 
@@ -79,8 +76,9 @@ public class InstrumentationUtils {
   }
 
   /**
-   * Gets the <a href= "http://developer.android.com/tools/testing/testing_otheride.html#AMOptionsSyntax"
-   * >am instrument options</a>.
+   * Gets the <a href=
+   * "http://developer.android.com/tools/testing/testing_otheride.html#AMOptionsSyntax" >am
+   * instrument options</a>.
    */
   public static Bundle getOptions() {
     checkInitialized();
@@ -110,14 +108,15 @@ public class InstrumentationUtils {
    * example, the ProgressBar.
    */
   public static boolean tryWaitForIdleSync(long timeoutMillis) {
-    validateNotAppThread();
+    checkNotMainThread();
     FutureTask<Void> emptyTask = new FutureTask<Void>(EMPTY_RUNNABLE, null);
     instrumentation.waitForIdle(emptyTask);
 
     try {
       emptyTask.get(timeoutMillis, TimeUnit.MILLISECONDS);
     } catch (java.util.concurrent.TimeoutException e) {
-      Logs.log(Log.INFO,
+      Logs.log(
+          Log.INFO,
           "Timed out after " + timeoutMillis + " milliseconds waiting for idle on main looper");
       return false;
     } catch (Throwable t) {
@@ -127,44 +126,51 @@ public class InstrumentationUtils {
   }
 
   public static void runOnMainSyncWithTimeout(final Runnable runnable) {
-    runOnMainSyncWithTimeout(new Callable<Void>() {
-      @Override
-      public Void call() throws Exception {
-        runnable.run();
-        return null;
-      }
-    });
+    runOnMainSyncWithTimeout(
+        new Callable<Void>() {
+          @Override
+          public Void call() throws Exception {
+            runnable.run();
+            return null;
+          }
+        });
   }
 
   /**
    * Runs {@code callable} on the main thread on best-effort basis up to a time limit, which
    * defaults to {@code 10000L} and can be set as an am instrument option under the key {@code
-   * dd.runOnMainSyncTimeout}. <p>This is a safer variation of {@link Instrumentation#runOnMainSync}
-   * because the latter may hang. You may turn off this behavior by setting {@code "-e
-   * dd.runOnMainSyncTimeout 0"} on the am command line.</p>The {@code callable} may never run, for
-   * example, if the main Looper has exited due to uncaught exception.
+   * dd.runOnMainSyncTimeout}.
+   *
+   * <p>This is a safer variation of {@link Instrumentation#runOnMainSync} because the latter may
+   * hang. You may turn off this behavior by setting {@code "-e dd.runOnMainSyncTimeout 0"} on the
+   * am command line.The {@code callable} may never run, for example, if the main Looper has exited
+   * due to uncaught exception.
    */
   public static <V> V runOnMainSyncWithTimeout(Callable<V> callable) {
-    validateNotAppThread();
+    checkNotMainThread();
     final RunOnMainSyncFutureTask<V> futureTask = new RunOnMainSyncFutureTask<>(callable);
 
     if (runOnMainSyncTimeoutMillis <= 0L) {
       // Call runOnMainSync on current thread without time limit.
       futureTask.runOnMainSyncNoThrow();
     } else {
-      RUN_ON_MAIN_SYNC_EXECUTOR.execute(new Runnable() {
-        @Override
-        public void run() {
-          futureTask.runOnMainSyncNoThrow();
-        }
-      });
+      RUN_ON_MAIN_SYNC_EXECUTOR.execute(
+          new Runnable() {
+            @Override
+            public void run() {
+              futureTask.runOnMainSyncNoThrow();
+            }
+          });
     }
 
     try {
       return futureTask.get(runOnMainSyncTimeoutMillis, TimeUnit.MILLISECONDS);
     } catch (java.util.concurrent.TimeoutException e) {
-      throw new TimeoutException("Timed out after " + runOnMainSyncTimeoutMillis
-          + " milliseconds waiting for Instrumentation.runOnMainSync", e);
+      throw new TimeoutException(
+          "Timed out after "
+              + runOnMainSyncTimeoutMillis
+              + " milliseconds waiting for Instrumentation.runOnMainSync",
+          e);
     } catch (Throwable t) {
       throw DroidDriverException.propagate(t);
     } finally {
@@ -172,6 +178,18 @@ public class InstrumentationUtils {
     }
   }
 
+  public static void checkMainThread() {
+    if (Looper.myLooper() != Looper.getMainLooper()) {
+      throw new DroidDriverException("This method must be called on the main thread");
+    }
+  }
+
+  public static void checkNotMainThread() {
+    if (Looper.myLooper() == Looper.getMainLooper()) {
+      throw new DroidDriverException("This method cannot be called on the main thread");
+    }
+  }
+
   private static class RunOnMainSyncFutureTask<V> extends FutureTask<V> {
     public RunOnMainSyncFutureTask(Callable<V> callable) {
       super(callable);
@@ -185,11 +203,4 @@ public class InstrumentationUtils {
       }
     }
   }
-
-  private static void validateNotAppThread() {
-    if (Looper.myLooper() == Looper.getMainLooper()) {
-      throw new DroidDriverException(
-          "This method can not be called from the main application thread");
-    }
-  }
 }