1 files changed, 14 insertions, 19 deletions

diff --git a/tests/src/com/android/launcher3/ui/TestViewHelpers.java b/tests/src/com/android/launcher3/ui/TestViewHelpers.java
index 083f58015c..4b2bade452 100644
--- a/tests/src/com/android/launcher3/ui/TestViewHelpers.java
+++ b/tests/src/com/android/launcher3/ui/TestViewHelpers.java
@@ -15,11 +15,14 @@
  */
 package com.android.launcher3.ui;
 
-import static androidx.test.InstrumentationRegistry.getInstrumentation;
-import static androidx.test.InstrumentationRegistry.getTargetContext;
+import static android.os.Process.myUserHandle;
 
+import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
+
+import static com.android.launcher3.util.TestUtil.getOnUiThread;
+
+import android.app.Instrumentation;
 import android.content.ComponentName;
-import android.os.Process;
 import android.util.Log;
 import android.view.View;
 import android.view.ViewGroup;
@@ -29,7 +32,6 @@ import com.android.launcher3.testcomponent.AppWidgetWithConfig;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.WidgetManagerHelper;
 
-import java.util.concurrent.Callable;
 import java.util.function.Function;
 
 public class TestViewHelpers {
@@ -38,23 +40,16 @@ public class TestViewHelpers {
     /**
      * Finds a widget provider which can fit on the home screen.
      *
-     * @param test               test suite.
      * @param hasConfigureScreen if true, a provider with a config screen is returned.
      */
-    public static LauncherAppWidgetProviderInfo findWidgetProvider(AbstractLauncherUiTest test,
-            final boolean hasConfigureScreen) {
-        LauncherAppWidgetProviderInfo info =
-                test.getOnUiThread(new Callable<LauncherAppWidgetProviderInfo>() {
-                    @Override
-                    public LauncherAppWidgetProviderInfo call() throws Exception {
-                        ComponentName cn = new ComponentName(getInstrumentation().getContext(),
-                                hasConfigureScreen ? AppWidgetWithConfig.class
-                                        : AppWidgetNoConfig.class);
-                        Log.d(TAG, "findWidgetProvider componentName=" + cn.flattenToString());
-                        return new WidgetManagerHelper(getTargetContext())
-                                .findProvider(cn, Process.myUserHandle());
-                    }
-                });
+    public static LauncherAppWidgetProviderInfo findWidgetProvider(boolean hasConfigureScreen) {
+        LauncherAppWidgetProviderInfo info = getOnUiThread(() -> {
+            Instrumentation i = getInstrumentation();
+            ComponentName cn = new ComponentName(i.getContext(),
+                    hasConfigureScreen ? AppWidgetWithConfig.class : AppWidgetNoConfig.class);
+            Log.d(TAG, "findWidgetProvider componentName=" + cn.flattenToString());
+            return new WidgetManagerHelper(i.getTargetContext()).findProvider(cn, myUserHandle());
+        });
         if (info == null) {
             throw new IllegalArgumentException("No valid widget provider");
         }