Fix a bug where image viewer custom tab has a wrong return stack

chrome/android/javatests/src/org/chromium/chrome/browser/customtabs/CustomTabActivityTestBase.java

Index: chrome/android/javatests/src/org/chromium/chrome/browser/customtabs/CustomTabActivityTestBase.java
diff --git a/chrome/android/javatests/src/org/chromium/chrome/browser/customtabs/CustomTabActivityTestBase.java b/chrome/android/javatests/src/org/chromium/chrome/browser/customtabs/CustomTabActivityTestBase.java
index aea13ca1ef181728514f03bf35d6c398dabbbf9c..210d053bd0776eb3a97d498f6a9ab4ad6fa6fcbd 100644
--- a/chrome/android/javatests/src/org/chromium/chrome/browser/customtabs/CustomTabActivityTestBase.java
+++ b/chrome/android/javatests/src/org/chromium/chrome/browser/customtabs/CustomTabActivityTestBase.java
@@ -7,9 +7,9 @@ package org.chromium.chrome.browser.customtabs;
 import static org.chromium.base.test.util.ScalableTimeout.scaleTimeout;
 
 import android.app.Activity;
-import android.app.Instrumentation;
 import android.content.Intent;
 
+import org.chromium.base.ApplicationStatus;
 import org.chromium.base.test.util.CallbackHelper;
 import org.chromium.chrome.browser.DeferredStartupHandler;
 import org.chromium.chrome.browser.tab.EmptyTabObserver;
@@ -18,6 +18,8 @@ import org.chromium.chrome.test.ChromeActivityTestCaseBase;
 import org.chromium.content.browser.test.util.Criteria;
 import org.chromium.content.browser.test.util.CriteriaHelper;
 
+import java.lang.ref.WeakReference;
+import java.util.List;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 
@@ -40,15 +42,23 @@ public abstract class CustomTabActivityTestBase extends
 
     @Override
     protected void startActivityCompletely(Intent intent) {
-        Instrumentation.ActivityMonitor monitor = getInstrumentation().addMonitor(
-                CustomTabActivity.class.getName(), null, false);
         Activity activity = getInstrumentation().startActivitySync(intent);
         assertNotNull("Main activity did not start", activity);
-        CustomTabActivity customTabActivity =
-                (CustomTabActivity) monitor.waitForActivityWithTimeout(
-                ACTIVITY_START_TIMEOUT_MS);
-        assertNotNull("CustomTabActivity did not start", customTabActivity);
-        setActivity(customTabActivity);
+        CriteriaHelper.pollUiThread(new Criteria() {
+            @Override
+            public boolean isSatisfied() {
+                List<WeakReference<Activity>> list = ApplicationStatus.getRunningActivities();
+                for (WeakReference<Activity> ref : list) {
+                    Activity activity = ref.get();
+                    if (activity == null) continue;
+                    if (activity instanceof CustomTabActivity) {
+                        setActivity(activity);
+                        return true;
+                    }
+                }
+                return false;
+            }
+        });
     }
 
     /**