From 8e4077f8cdeb721718269fca19eb5713d9af4261 Mon Sep 17 00:00:00 2001
From: Jagrut Desai <jagrutdesai@google.com>
Date: Tue, 10 Oct 2023 16:06:15 -0700
Subject: [PATCH] Adding test protocol support for getting num all apps columns

Test: Presubmit
Bug: 303257279
Flag: not needed
Change-Id: I9c6299fb3ff077bc42ca22609d2ff0232377a370
---
 .../android/launcher3/testing/TestInformationHandler.java   | 5 +++++
 .../com/android/launcher3/testing/shared/TestProtocol.java  | 1 +
 .../com/android/launcher3/tapl/LauncherInstrumentation.java | 6 ++++++
 3 files changed, 12 insertions(+)

diff --git a/src/com/android/launcher3/testing/TestInformationHandler.java b/src/com/android/launcher3/testing/TestInformationHandler.java
index e52e878d87..0438e575c6 100644
--- a/src/com/android/launcher3/testing/TestInformationHandler.java
+++ b/src/com/android/launcher3/testing/TestInformationHandler.java
@@ -165,6 +165,11 @@ public class TestInformationHandler implements ResourceBasedOverride {
                 response.putBoolean(TestProtocol.TEST_INFO_RESPONSE_FIELD, mDeviceProfile.isTablet);
                 return response;
 
+            case TestProtocol.REQUEST_NUM_ALL_APPS_COLUMNS:
+                response.putInt(TestProtocol.TEST_INFO_RESPONSE_FIELD,
+                        mDeviceProfile.numShownAllAppsColumns);
+                return response;
+
             case TestProtocol.REQUEST_IS_TWO_PANELS:
                 response.putBoolean(TestProtocol.TEST_INFO_RESPONSE_FIELD,
                         FOLDABLE_SINGLE_PAGE.get() ? false : mDeviceProfile.isTwoPanels);
diff --git a/tests/shared/com/android/launcher3/testing/shared/TestProtocol.java b/tests/shared/com/android/launcher3/testing/shared/TestProtocol.java
index 54a1c08e56..51f457d7f5 100644
--- a/tests/shared/com/android/launcher3/testing/shared/TestProtocol.java
+++ b/tests/shared/com/android/launcher3/testing/shared/TestProtocol.java
@@ -115,6 +115,7 @@ public final class TestProtocol {
     public static final String REQUEST_CLEAR_DATA = "clear-data";
     public static final String REQUEST_HOTSEAT_ICON_NAMES = "get-hotseat-icon-names";
     public static final String REQUEST_IS_TABLET = "is-tablet";
+    public static final String REQUEST_NUM_ALL_APPS_COLUMNS = "num-all-apps-columns";
     public static final String REQUEST_IS_TWO_PANELS = "is-two-panel";
     public static final String REQUEST_START_DRAG_THRESHOLD = "start-drag-threshold";
     public static final String REQUEST_SHELL_DRAG_READY = "shell-drag-ready";
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index d7f9c78ab6..307f192f64 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -27,6 +27,7 @@ import static android.view.MotionEvent.AXIS_GESTURE_SWIPE_FINGER_COUNT;
 import static com.android.launcher3.tapl.Folder.FOLDER_CONTENT_RES_ID;
 import static com.android.launcher3.tapl.TestHelpers.getOverviewPackageName;
 import static com.android.launcher3.testing.shared.TestProtocol.NORMAL_STATE_ORDINAL;
+import static com.android.launcher3.testing.shared.TestProtocol.REQUEST_NUM_ALL_APPS_COLUMNS;
 
 import android.app.ActivityManager;
 import android.app.Instrumentation;
@@ -353,6 +354,11 @@ public final class LauncherInstrumentation {
                 .getParcelable(TestProtocol.TEST_INFO_RESPONSE_FIELD);
     }
 
+    public int getNumAllAppsColumns() {
+        return getTestInfo(REQUEST_NUM_ALL_APPS_COLUMNS).getInt(
+                TestProtocol.TEST_INFO_RESPONSE_FIELD);
+    }
+
     public boolean isTablet() {
         return getTestInfo(TestProtocol.REQUEST_IS_TABLET)
                 .getBoolean(TestProtocol.TEST_INFO_RESPONSE_FIELD);