1 files changed, 16 insertions, 10 deletions
diff --git a/tests/tapl/com/android/launcher3/tapl/SearchResultFromQsb.java b/tests/tapl/com/android/launcher3/tapl/SearchResultFromQsb.java
index 80176e993f..d02e747900 100644
--- a/tests/tapl/com/android/launcher3/tapl/SearchResultFromQsb.java
+++ b/tests/tapl/com/android/launcher3/tapl/SearchResultFromQsb.java
@@ -32,7 +32,7 @@ public class SearchResultFromQsb {
 
     // This particular ID change should happen with caution
     private static final String SEARCH_CONTAINER_RES_ID = "search_results_list_view";
-    private final LauncherInstrumentation mLauncher;
+    protected final LauncherInstrumentation mLauncher;
 
     SearchResultFromQsb(LauncherInstrumentation launcher) {
         mLauncher = launcher;
@@ -49,27 +49,33 @@ public class SearchResultFromQsb {
     }
 
     /** Find the app from search results with app name. */
-    public Launchable findAppIcon(String appName) {
+    public AppIcon findAppIcon(String appName) {
         UiObject2 icon = mLauncher.waitForLauncherObject(By.clazz(TextView.class).text(appName));
+        return createAppIcon(icon);
+    }
+
+    protected AppIcon createAppIcon(UiObject2 icon) {
         return new AllAppsAppIcon(mLauncher, icon);
     }
 
     /** Find the web suggestion from search suggestion's title text */
-    public void verifyWebSuggestIsPresent(String text) {
-        ArrayList<UiObject2> goldenGateResults =
+    public SearchWebSuggestion findWebSuggestion(String text) {
+        ArrayList<UiObject2> webSuggestions =
                 new ArrayList<>(mLauncher.waitForObjectsInContainer(
                         mLauncher.waitForSystemLauncherObject(SEARCH_CONTAINER_RES_ID),
                         By.clazz(TextView.class)));
-        boolean found = false;
-        for(UiObject2 uiObject: goldenGateResults) {
+        for (UiObject2 uiObject: webSuggestions) {
             String currentString = uiObject.getText();
             if (currentString.equals(text)) {
-                found = true;
+                return createWebSuggestion(uiObject);
             }
         }
-        if (!found) {
-            throw new IllegalStateException("Web suggestion title: " + text + " not found");
-        }
+        mLauncher.fail("Web suggestion title: " + text + " not found");
+        return null;
+    }
+
+    protected SearchWebSuggestion createWebSuggestion(UiObject2 webSuggestion) {
+        return new SearchWebSuggestion(mLauncher, webSuggestion);
     }
 
     /** Find the total amount of views being displayed and return the size */