1 files changed, 24 insertions, 2 deletions

diff --git a/tests/tapl/com/android/launcher3/tapl/Workspace.java b/tests/tapl/com/android/launcher3/tapl/Workspace.java
index 33754c125..b01b6f363 100644
--- a/tests/tapl/com/android/launcher3/tapl/Workspace.java
+++ b/tests/tapl/com/android/launcher3/tapl/Workspace.java
@@ -27,6 +27,7 @@ import android.view.MotionEvent;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
+import androidx.test.uiautomator.By;
 import androidx.test.uiautomator.Direction;
 import androidx.test.uiautomator.UiObject2;
 
@@ -142,11 +143,17 @@ public final class Workspace extends Home {
     }
 
     @NonNull
-    private AppIcon getHotseatAppIcon(String appName) {
+    public AppIcon getHotseatAppIcon(String appName) {
         return new AppIcon(mLauncher, mLauncher.getObjectInContainer(
                 mHotseat, AppIcon.getAppIconSelector(appName, mLauncher)));
     }
 
+    @NonNull
+    public Folder getHotseatFolder(String appName) {
+        return new Folder(mLauncher, mLauncher.getObjectInContainer(
+                mHotseat, Folder.getSelector(appName, mLauncher)));
+    }
+
     static void dragIconToWorkspace(
             LauncherInstrumentation launcher, Launchable launchable, Point dest,
             String longPressIndicator) {
@@ -213,6 +220,21 @@ public final class Workspace extends Home {
 
     @Override
     protected int getSwipeStartY() {
-        return mLauncher.waitForLauncherObject("hotseat").getVisibleBounds().top;
+        return mLauncher.getRealDisplaySize().y - 1;
+    }
+
+    @Nullable
+    public Widget tryGetWidget(String label, long timeout) {
+        final UiObject2 widget = mLauncher.tryWaitForLauncherObject(
+                By.clazz("com.android.launcher3.widget.LauncherAppWidgetHostView").desc(label),
+                timeout);
+        return widget != null ? new Widget(mLauncher, widget) : null;
+    }
+
+    @Nullable
+    public Widget tryGetPendingWidget(long timeout) {
+        final UiObject2 widget = mLauncher.tryWaitForLauncherObject(
+                By.clazz("com.android.launcher3.widget.PendingAppWidgetHostView"), timeout);
+        return widget != null ? new Widget(mLauncher, widget) : null;
     }
 }
\ No newline at end of file