add a time-limited version of runOnMainSync

Change-Id: I40b23d4f6cdbf05fb34d7252f1b98eef9d442a43

5 files changed, 72 insertions, 9 deletions

diff --git a/src/com/google/android/droiddriver/base/DroidDriverContext.java b/src/com/google/android/droiddriver/base/DroidDriverContext.java
index 17d681c..1c4b21f 100644
--- a/src/com/google/android/droiddriver/base/DroidDriverContext.java
+++ b/src/com/google/android/droiddriver/base/DroidDriverContext.java
@@ -22,6 +22,7 @@ import android.util.Log;
 
 import com.google.android.droiddriver.actions.InputInjector;
 import com.google.android.droiddriver.exceptions.DroidDriverException;
+import com.google.android.droiddriver.exceptions.TimeoutException;
 import com.google.android.droiddriver.util.Logs;
 
 import java.util.concurrent.ExecutionException;
@@ -76,6 +77,46 @@ public abstract class DroidDriverContext {
     return true;
   }
 
+  /**
+   * Tries to run {@code runnable} on the main thread on best-effort basis up to
+   * {@code timeoutMillis}. The {@code runnable} may never run, for example, in
+   * case that the main Looper has exited due to uncaught exception.
+   */
+  public boolean tryRunOnMainSync(Runnable runnable, long timeoutMillis) {
+    validateNotAppThread();
+    final FutureTask<?> futureTask = new FutureTask<Void>(runnable, null);
+    new Thread(new Runnable() {
+      @Override
+      public void run() {
+        instrumentation.runOnMainSync(futureTask);
+      }
+    }).start();
+
+    try {
+      futureTask.get(timeoutMillis, TimeUnit.MILLISECONDS);
+    } catch (InterruptedException e) {
+      throw new DroidDriverException(e);
+    } catch (ExecutionException e) {
+      throw new DroidDriverException(e);
+    } catch (java.util.concurrent.TimeoutException e) {
+      Logs.log(Log.WARN, getRunOnMainSyncTimeoutMessage(timeoutMillis));
+      return false;
+    }
+    return true;
+  }
+
+  public void runOnMainSync(Runnable runnable) {
+    long timeoutMillis = getDriver().getPoller().getTimeoutMillis();
+    if (!tryRunOnMainSync(runnable, timeoutMillis)) {
+      throw new TimeoutException(getRunOnMainSyncTimeoutMessage(timeoutMillis));
+    }
+  }
+
+  private String getRunOnMainSyncTimeoutMessage(long timeoutMillis) {
+    return String.format(
+        "Timed out after %d milliseconds waiting for Instrumentation.runOnMainSync", timeoutMillis);
+  }
+
   private void validateNotAppThread() {
     if (Looper.myLooper() == Looper.getMainLooper()) {
       throw new DroidDriverException(
diff --git a/src/com/google/android/droiddriver/helpers/BaseDroidDriverTest.java b/src/com/google/android/droiddriver/helpers/BaseDroidDriverTest.java
index 09b4a36..4f125c8 100644
--- a/src/com/google/android/droiddriver/helpers/BaseDroidDriverTest.java
+++ b/src/com/google/android/droiddriver/helpers/BaseDroidDriverTest.java
@@ -93,7 +93,6 @@ public abstract class BaseDroidDriverTest<T extends Activity> extends
     // Give uncaughtException (thrown by app instead of tests) high priority
     if (uncaughtException != null) {
       failure = uncaughtException;
-      uncaughtException = null;
       skipRemainingTests = true;
     }
 
@@ -103,7 +102,7 @@ public abstract class BaseDroidDriverTest<T extends Activity> extends
       }
       if (failure instanceof OutOfMemoryError) {
         dumpHprof();
-      } else {
+      } else if (uncaughtException == null) {
         String baseFileName = getBaseFileName();
         driver.dumpUiElementTree(baseFileName + ".xml");
         driver.getUiDevice().takeScreenshot(baseFileName + ".png");
@@ -112,7 +111,7 @@ public abstract class BaseDroidDriverTest<T extends Activity> extends
       // This method is for troubleshooting. Do not throw new error; we'll
       // throw the original failure.
       Logs.log(Log.WARN, e);
-      if (e instanceof OutOfMemoryError) {
+      if (e instanceof OutOfMemoryError && !(failure instanceof OutOfMemoryError)) {
         dumpHprof();
       }
     } finally {
@@ -145,6 +144,10 @@ public abstract class BaseDroidDriverTest<T extends Activity> extends
     if (skipRemainingTests) {
       return;
     }
+    if (uncaughtException != null) {
+      onFailure(uncaughtException);
+    }
+
     Throwable exception = null;
     try {
       setUp();
@@ -166,8 +169,5 @@ public abstract class BaseDroidDriverTest<T extends Activity> extends
     if (exception != null) {
       throw exception;
     }
-    if (uncaughtException != null) {
-      onFailure(uncaughtException);
-    }
   }
 }
diff --git a/src/com/google/android/droiddriver/instrumentation/InstrumentationDriver.java b/src/com/google/android/droiddriver/instrumentation/InstrumentationDriver.java
index 7684530..ebdc465 100644
--- a/src/com/google/android/droiddriver/instrumentation/InstrumentationDriver.java
+++ b/src/com/google/android/droiddriver/instrumentation/InstrumentationDriver.java
@@ -123,7 +123,7 @@ public class InstrumentationDriver extends BaseDroidDriver {
 
   Bitmap takeScreenshot() {
     ScreenshotRunnable screenshotRunnable = new ScreenshotRunnable(findRootView());
-    context.getInstrumentation().runOnMainSync(screenshotRunnable);
+    context.runOnMainSync(screenshotRunnable);
     return screenshotRunnable.screenshot;
   }
 
diff --git a/src/com/google/android/droiddriver/instrumentation/ViewElement.java b/src/com/google/android/droiddriver/instrumentation/ViewElement.java
index 55f9819..5c39dc4 100644
--- a/src/com/google/android/droiddriver/instrumentation/ViewElement.java
+++ b/src/com/google/android/droiddriver/instrumentation/ViewElement.java
@@ -216,7 +216,7 @@ public class ViewElement extends BaseUiElement {
     Preconditions.checkNotNull(view);
     this.parent = parent;
     SnapshotViewAttributesRunnable attributesSnapshot = new SnapshotViewAttributesRunnable(view);
-    context.getInstrumentation().runOnMainSync(attributesSnapshot);
+    context.runOnMainSync(attributesSnapshot);
     if (attributesSnapshot.exception != null) {
       throw new DroidDriverException(attributesSnapshot.exception);
     }
diff --git a/src/com/google/android/droiddriver/runner/TestRunner.java b/src/com/google/android/droiddriver/runner/TestRunner.java
index 9b9846e..7f6b0d1 100644
--- a/src/com/google/android/droiddriver/runner/TestRunner.java
+++ b/src/com/google/android/droiddriver/runner/TestRunner.java
@@ -40,6 +40,8 @@ import java.lang.annotation.Annotation;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.FutureTask;
+import java.util.concurrent.TimeUnit;
 
 /**
  * Adds activity watcher to InstrumentationTestRunner.
@@ -68,7 +70,7 @@ public class TestRunner extends InstrumentationTestRunner {
     getAndroidTestRunner().addTestListener(new TestListener() {
       @Override
       public void endTest(Test test) {
-        runOnMainSync(new Runnable() {
+        runOnMainSyncWithTimeLimit(new Runnable() {
           @Override
           public void run() {
             Iterator<Activity> iterator = activities.iterator();
@@ -167,4 +169,24 @@ public class TestRunner extends InstrumentationTestRunner {
       runningActivity = null;
     }
   }
+
+  private void runOnMainSyncWithTimeLimit(Runnable runnable) {
+    // Do we need it configurable? Now only used in endTest.
+    long timeoutMillis = 10000L;
+    final FutureTask<?> futureTask = new FutureTask<Void>(runnable, null);
+    new Thread(new Runnable() {
+      @Override
+      public void run() {
+        runOnMainSync(futureTask);
+      }
+    }).start();
+
+    try {
+      futureTask.get(timeoutMillis, TimeUnit.MILLISECONDS);
+    } catch (Throwable e) {
+      Logs.log(Log.WARN, e, String.format(
+          "Timed out after %d milliseconds waiting for Instrumentation.runOnMainSync",
+          timeoutMillis));
+    }
+  }
 }