Add TAPL tests for PredictionRow.

This will help us catch and debug following issues
* Prediction row view is missing
* Prediction row view doesn't contain enough apps

Bug: 332171918
Test: Manual
Flag: NA

Change-Id: I11035997afb26799a2c3b69d5139b4f3c28fc405
diff --git a/quickstep/tests/src/com/android/quickstep/util/TaplTestsPredictionRow.java b/quickstep/tests/src/com/android/quickstep/util/TaplTestsPredictionRow.java
new file mode 100644
index 0000000..4030b01
--- /dev/null
+++ b/quickstep/tests/src/com/android/quickstep/util/TaplTestsPredictionRow.java
@@ -0,0 +1,35 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep.util;
+
+import com.android.quickstep.AbstractQuickStepTest;
+
+import org.junit.Test;
+
+public class TaplTestsPredictionRow extends AbstractQuickStepTest {
+
+    @Override
+    public void setUp() throws Exception {
+        super.setUp();
+        mLauncher.getWorkspace().switchToAllApps();
+    }
+
+    @Test
+    public void testPredictionRow() {
+        mLauncher.getAllApps().getPredictionRowView();
+    }
+}
diff --git a/tests/tapl/com/android/launcher3/tapl/AllApps.java b/tests/tapl/com/android/launcher3/tapl/AllApps.java
index 245ec09..b490124 100644
--- a/tests/tapl/com/android/launcher3/tapl/AllApps.java
+++ b/tests/tapl/com/android/launcher3/tapl/AllApps.java
@@ -291,6 +291,10 @@
         return mLauncher.waitForObjectInContainer(allAppsContainer, "apps_list_view");
     }
 
+    protected UiObject2 getAllAppsHeader(UiObject2 allAppsContainer) {
+        return mLauncher.waitForObjectInContainer(allAppsContainer, "all_apps_header");
+    }
+
     protected UiObject2 getSearchBox(UiObject2 allAppsContainer) {
         return mLauncher.waitForObjectInContainer(allAppsContainer, "search_container_all_apps");
     }
@@ -416,6 +420,14 @@
         }
     }
 
+    /** Returns PredictionRow if present in view. */
+    @NonNull
+    public PredictionRow getPredictionRowView() {
+        final UiObject2 allAppsContainer = verifyActiveContainer();
+        final UiObject2 allAppsHeader = getAllAppsHeader(allAppsContainer);
+        return new PredictionRow(mLauncher, allAppsHeader);
+    }
+
     /** Returns PrivateSpaceContainer if present in view. */
     @NonNull
     public PrivateSpaceContainer getPrivateSpaceUnlockedView() {
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index ba8121f..764a455 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -766,7 +766,7 @@
         fail(message + ". " + "Actual: " + actual);
     }
 
-    private void assertEquals(String message, int expected, int actual) {
+    void assertEquals(String message, int expected, int actual) {
         if (expected != actual) {
             fail(message + " expected: " + expected + " but was: " + actual);
         }
diff --git a/tests/tapl/com/android/launcher3/tapl/PredictionRow.java b/tests/tapl/com/android/launcher3/tapl/PredictionRow.java
new file mode 100644
index 0000000..11f71fd
--- /dev/null
+++ b/tests/tapl/com/android/launcher3/tapl/PredictionRow.java
@@ -0,0 +1,72 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.tapl;
+
+import androidx.annotation.NonNull;
+import androidx.test.uiautomator.UiObject2;
+
+/** View containing prediction app icons */
+public class PredictionRow {
+
+    private static final String PREDICTION_ROW_ID = "prediction_row";
+    private static final String PREDICTION_APP_ID = "icon";
+    private final LauncherInstrumentation mLauncher;
+    private final UiObject2 mAllAppsHeader;
+    private final UiObject2 mPredictionRow;
+
+    PredictionRow(LauncherInstrumentation launcherInstrumentation,
+            UiObject2 allAppsHeader) {
+        mLauncher = launcherInstrumentation;
+        mAllAppsHeader = allAppsHeader;
+        mPredictionRow = mLauncher.waitForObjectInContainer(mAllAppsHeader,
+                PREDICTION_ROW_ID);
+        verifyAppsPresentInsidePredictionRow();
+        verifyPredictionRowAppsCount();
+    }
+
+    /** Verify that one app is present in prediction row view. */
+    private void verifyAppsPresentInsidePredictionRow() {
+        mLauncher.waitForObjectInContainer(mPredictionRow,
+                PREDICTION_APP_ID);
+    }
+
+    /** Verify that prediction row apps count is same as launcher apps column count. */
+    private void verifyPredictionRowAppsCount() {
+        mLauncher.assertEquals("PredictionRow app count mismatch", mLauncher.getNumAllAppsColumns(),
+                getPredictionRowAppsCount());
+    }
+
+    /**
+     * Returns an app icon found in the prediction row. This fails if any icon is not
+     * found.
+     */
+    @NonNull
+    private HomeAppIcon getAnyAppIcon() {
+        return new AllAppsAppIcon(mLauncher,
+                mPredictionRow.findObject(AppIcon.getAnyAppIconSelector()));
+    }
+
+    /**
+     * Returns the size of prediction row apps count.
+     */
+    private int getPredictionRowAppsCount() {
+        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
+                "want to get all prediction row icons")) {
+            return mPredictionRow.findObjects(AppIcon.getAnyAppIconSelector()).size();
+        }
+    }
+}