From b5f26d897d16c6533b956e320425139ccdfe19bb Mon Sep 17 00:00:00 2001
From: JW Wang <wangchun@google.com>
Date: Thu, 26 Mar 2020 17:14:26 +0800
Subject: [PATCH] Let #getUiAutomation return null if UiAutomation fails to
 connect (3/n)

Now #getUiAutomation returns null if the underlying UiAutomation fails
to connect for targetSdkVersion > R. This gives the caller a chance to
retry and decide when to give up.

Bug: 147785023
Test: presubmit
Change-Id: Ibf201fa042b80c067f2a0c3b0bde23adf9c60546
Merged-In: Ibf201fa042b80c067f2a0c3b0bde23adf9c60546
---
 core/java/android/app/Instrumentation.java | 29 ++++++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/core/java/android/app/Instrumentation.java b/core/java/android/app/Instrumentation.java
index 721525d9af9d7..2ef147b3e903e 100644
--- a/core/java/android/app/Instrumentation.java
+++ b/core/java/android/app/Instrumentation.java
@@ -62,6 +62,7 @@ import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.TimeoutException;
 
 /**
  * Base class for implementing application instrumentation code.  When running
@@ -90,6 +91,8 @@ public class Instrumentation {
 
     private static final String TAG = "Instrumentation";
 
+    private static final long CONNECT_TIMEOUT_MILLIS = 5000;
+
     /**
      * @hide
      */
@@ -2125,6 +2128,13 @@ public class Instrumentation {
      * Equivalent to {@code getUiAutomation(0)}. If a {@link UiAutomation} exists with different
      * flags, the flags on that instance will be changed, and then it will be returned.
      * </p>
+     * <p>
+     * Compatibility mode: This method is infallible for apps targeted for
+     * {@link Build.VERSION_CODES#R} and earlier versions; for apps targeted for later versions, it
+     * will return null if {@link UiAutomation} fails to connect. The caller can check the return
+     * value and retry on error.
+     * </p>
+     *
      * @return The UI automation instance.
      *
      * @see UiAutomation
@@ -2152,6 +2162,12 @@ public class Instrumentation {
      * If a {@link UiAutomation} exists with different flags, the flags on that instance will be
      * changed, and then it will be returned.
      * </p>
+     * <p>
+     * Compatibility mode: This method is infallible for apps targeted for
+     * {@link Build.VERSION_CODES#R} and earlier versions; for apps targeted for later versions, it
+     * will return null if {@link UiAutomation} fails to connect. The caller can check the return
+     * value and retry on error.
+     * </p>
      *
      * @param flags The flags to be passed to the UiAutomation, for example
      *        {@link UiAutomation#FLAG_DONT_SUPPRESS_ACCESSIBILITY_SERVICES}.
@@ -2173,8 +2189,17 @@ public class Instrumentation {
             } else {
                 mUiAutomation.disconnect();
             }
-            mUiAutomation.connect(flags);
-            return mUiAutomation;
+            if (getTargetContext().getApplicationInfo().targetSdkVersion <= Build.VERSION_CODES.R) {
+                mUiAutomation.connect(flags);
+                return mUiAutomation;
+            }
+            try {
+                mUiAutomation.connectWithTimeout(flags, CONNECT_TIMEOUT_MILLIS);
+                return mUiAutomation;
+            } catch (TimeoutException e) {
+                mUiAutomation.destroy();
+                mUiAutomation = null;
+            }
         }
         return null;
     }