Using a runtime generated layout for tests instead of defining xml

This allows support for easily setting up default layouts

Bug: 277345535
Test: Presubmit
Flag: N/A
Change-Id: I1c089d60ac3f8add8d7e1060d343e04d30afe094
diff --git a/ext_tests/src/com/android/launcher3/testing/DebugTestInformationHandler.java b/ext_tests/src/com/android/launcher3/testing/DebugTestInformationHandler.java
index 4f9c32a..3fddd9d 100644
--- a/ext_tests/src/com/android/launcher3/testing/DebugTestInformationHandler.java
+++ b/ext_tests/src/com/android/launcher3/testing/DebugTestInformationHandler.java
@@ -218,29 +218,6 @@
                 }
             }
 
-            case TestProtocol.REQUEST_USE_TEST_WORKSPACE_LAYOUT: {
-                useTestWorkspaceLayout(
-                        LauncherSettings.Settings.ARG_DEFAULT_WORKSPACE_LAYOUT_TEST);
-                return response;
-            }
-
-            case TestProtocol.REQUEST_USE_TEST2_WORKSPACE_LAYOUT: {
-                useTestWorkspaceLayout(
-                        LauncherSettings.Settings.ARG_DEFAULT_WORKSPACE_LAYOUT_TEST2);
-                return response;
-            }
-
-            case TestProtocol.REQUEST_USE_TAPL_WORKSPACE_LAYOUT: {
-                useTestWorkspaceLayout(
-                        LauncherSettings.Settings.ARG_DEFAULT_WORKSPACE_LAYOUT_TAPL);
-                return response;
-            }
-
-            case TestProtocol.REQUEST_USE_DEFAULT_WORKSPACE_LAYOUT: {
-                useTestWorkspaceLayout(null);
-                return response;
-            }
-
             case TestProtocol.REQUEST_HOTSEAT_ICON_NAMES: {
                 return getLauncherUIProperty(Bundle::putStringArrayList, l -> {
                     ShortcutAndWidgetContainer hotseatIconsContainer =
@@ -278,20 +255,4 @@
                 return super.call(method, arg, extras);
         }
     }
-
-    private void useTestWorkspaceLayout(String layout) {
-        final long identity = Binder.clearCallingIdentity();
-        try {
-            if (layout != null) {
-                LauncherSettings.Settings.call(mContext.getContentResolver(),
-                        LauncherSettings.Settings.METHOD_SET_USE_TEST_WORKSPACE_LAYOUT_FLAG,
-                        layout);
-            } else {
-                LauncherSettings.Settings.call(mContext.getContentResolver(),
-                        LauncherSettings.Settings.METHOD_CLEAR_USE_TEST_WORKSPACE_LAYOUT_FLAG);
-            }
-        } finally {
-            Binder.restoreCallingIdentity(identity);
-        }
-    }
 }
diff --git a/quickstep/tests/src/com/android/quickstep/TaplTestsTaskbar.java b/quickstep/tests/src/com/android/quickstep/TaplTestsTaskbar.java
index 735c5e6..6243471 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplTestsTaskbar.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplTestsTaskbar.java
@@ -29,6 +29,8 @@
 
 import com.android.launcher3.tapl.Taskbar;
 import com.android.launcher3.ui.TaplTestsLauncher3;
+import com.android.launcher3.util.LauncherLayoutBuilder;
+import com.android.launcher3.util.TestUtil;
 import com.android.launcher3.util.rule.ScreenRecordRule.ScreenRecord;
 import com.android.quickstep.TaskbarModeSwitchRule.TaskbarModeSwitch;
 
@@ -49,11 +51,17 @@
     private static final String CALCULATOR_APP_PACKAGE =
             resolveSystemApp(Intent.CATEGORY_APP_CALCULATOR);
 
+    private AutoCloseable mLauncherLayout;
+
     @Override
     public void setUp() throws Exception {
         Assume.assumeTrue(mLauncher.isTablet());
         super.setUp();
-        mLauncher.useTestWorkspaceLayoutOnReload();
+
+        LauncherLayoutBuilder layoutBuilder = new LauncherLayoutBuilder().atHotseat(0).putApp(
+                "com.google.android.apps.nexuslauncher.tests",
+                "com.android.launcher3.testcomponent.BaseTestingActivity");
+        mLauncherLayout = TestUtil.setLauncherDefaultLayout(mTargetContext, layoutBuilder);
         TaplTestsLauncher3.initialize(this);
 
         startAppFast(CALCULATOR_APP_PACKAGE);
@@ -62,9 +70,11 @@
     }
 
     @After
-    public void tearDown() {
-        mLauncher.useDefaultWorkspaceLayoutOnReload();
+    public void tearDown() throws Exception {
         mLauncher.enableBlockTimeout(false);
+        if (mLauncherLayout != null) {
+            mLauncherLayout.close();
+        }
     }
 
     @Test
diff --git a/res/xml/default_tapl_test_workspace.xml b/res/xml/default_tapl_test_workspace.xml
deleted file mode 100644
index 24d76f3..0000000
--- a/res/xml/default_tapl_test_workspace.xml
+++ /dev/null
@@ -1,28 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-<!-- Copyright (C) 2023 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.
--->
-<!-- Split display specific version of Launcher3/res/xml/default_workspace_4x4.xml -->
-<favorites xmlns:launcher="http://schemas.android.com/apk/res-auto/com.android.launcher3" >
-
-    <!-- Hotseat (We use the screen as the position of the item in the hotseat) -->
-    <favorite
-        launcher:container="-101"
-        launcher:screen="0"
-        launcher:x="0"
-        launcher:y="0"
-        launcher:className="com.google.android.apps.chrome.Main"
-        launcher:packageName="com.android.chrome" />
-
-</favorites>
diff --git a/res/xml/default_test2_workspace.xml b/res/xml/default_test2_workspace.xml
deleted file mode 100644
index c560104..0000000
--- a/res/xml/default_test2_workspace.xml
+++ /dev/null
@@ -1,57 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-<!-- Split display specific version of Launcher3/res/xml/default_workspace_4x4.xml -->
-<favorites xmlns:launcher="http://schemas.android.com/apk/res-auto/com.android.launcher3" >
-
-    <!-- Hotseat (We use the screen as the position of the item in the hotseat) -->
-    <!-- Dialer Messaging Chrome Camera -->
-    <favorite
-        launcher:container="-101"
-        launcher:screen="0"
-        launcher:x="0"
-        launcher:y="0"
-        launcher:className="com.google.android.dialer.extensions.GoogleDialtactsActivity"
-        launcher:packageName="com.google.android.dialer" />
-
-    <favorite
-        launcher:container="-101"
-        launcher:screen="1"
-        launcher:x="1"
-        launcher:y="0"
-        launcher:className="com.google.android.apps.messaging.ui.ConversationListActivity"
-        launcher:packageName="com.google.android.apps.messaging" />
-
-    <favorite
-        launcher:container="-101"
-        launcher:screen="2"
-        launcher:x="2"
-        launcher:y="0"
-        launcher:className="com.google.android.apps.chrome.Main"
-        launcher:packageName="com.android.chrome" />
-
-    <favorite
-        launcher:container="-101"
-        launcher:screen="3"
-        launcher:x="3"
-        launcher:y="0"
-        launcher:className="com.android.camera.CameraLauncher"
-        launcher:packageName="com.google.android.GoogleCamera" />
-
-    <!-- Bottom row -->
-    <!-- Maps [space] [space] Play -->
-    <favorite
-        launcher:className="com.google.android.maps.MapsActivity"
-        launcher:packageName="com.google.android.apps.maps"
-        launcher:screen="0"
-        launcher:x="0"
-        launcher:y="-1" />
-
-    <favorite
-        launcher:className="com.android.vending.AssetBrowserActivity"
-        launcher:packageName="com.android.vending"
-        launcher:screen="0"
-        launcher:x="3"
-        launcher:y="-1" />
-
-    <!-- TODO: Place weather widget when it's available -->
-
-</favorites>
diff --git a/res/xml/default_test_workspace.xml b/res/xml/default_test_workspace.xml
deleted file mode 100644
index bd718b3..0000000
--- a/res/xml/default_test_workspace.xml
+++ /dev/null
@@ -1,29 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-<!-- Copyright (C) 2022 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.
--->
-
-<!-- Split display specific version of Launcher3/res/xml/default_workspace_4x4.xml -->
-<favorites xmlns:launcher="http://schemas.android.com/apk/res-auto/com.android.launcher3" >
-
-    <!-- Launcher Test Activity -->
-    <resolve
-        launcher:container="-101"
-        launcher:screen="0"
-        launcher:x="0"
-        launcher:y="0" >
-        <favorite launcher:uri="#Intent;action=android.intent.action.MAIN;category=android.intent.category.CATEGORY_TEST;component=com.google.android.apps.nexuslauncher.tests/com.android.launcher3.testcomponent.BaseTestingActivity;end" />
-    </resolve>
-
-</favorites>
diff --git a/src/com/android/launcher3/AutoInstallsLayout.java b/src/com/android/launcher3/AutoInstallsLayout.java
index 197aa5a..ede7e2f 100644
--- a/src/com/android/launcher3/AutoInstallsLayout.java
+++ b/src/com/android/launcher3/AutoInstallsLayout.java
@@ -27,6 +27,8 @@
 import android.content.pm.LauncherApps;
 import android.content.pm.PackageManager;
 import android.content.res.Resources;
+import android.content.res.Resources.NotFoundException;
+import android.content.res.XmlResourceParser;
 import android.database.sqlite.SQLiteDatabase;
 import android.net.Uri;
 import android.os.Bundle;
@@ -38,7 +40,9 @@
 import android.util.Xml;
 
 import androidx.annotation.Nullable;
+import androidx.annotation.StringRes;
 import androidx.annotation.WorkerThread;
+import androidx.annotation.XmlRes;
 
 import com.android.launcher3.LauncherProvider.SqlArguments;
 import com.android.launcher3.LauncherSettings.Favorites;
@@ -161,7 +165,7 @@
     protected final LayoutParserCallback mCallback;
 
     protected final PackageManager mPackageManager;
-    protected final Resources mSourceRes;
+    protected final SourceResources mSourceRes;
     protected final Supplier<XmlPullParser> mInitialLayoutSupplier;
 
     private final InvariantDeviceProfile mIdp;
@@ -178,11 +182,12 @@
     public AutoInstallsLayout(Context context, LauncherWidgetHolder appWidgetHolder,
             LayoutParserCallback callback, Resources res,
             int layoutId, String rootTag) {
-        this(context, appWidgetHolder, callback, res, () -> res.getXml(layoutId), rootTag);
+        this(context, appWidgetHolder, callback, SourceResources.wrap(res),
+                () -> res.getXml(layoutId), rootTag);
     }
 
     public AutoInstallsLayout(Context context, LauncherWidgetHolder appWidgetHolder,
-            LayoutParserCallback callback, Resources res,
+            LayoutParserCallback callback, SourceResources res,
             Supplier<XmlPullParser> initialLayoutSupplier, String rootTag) {
         mContext = context;
         mAppWidgetHolder = appWidgetHolder;
@@ -703,4 +708,42 @@
         to.put(key, from.getAsInteger(key));
     }
 
+    /**
+     * Wrapper over resources for easier abstraction
+     */
+    public interface SourceResources {
+
+        /**
+         * Refer {@link Resources#getXml(int)}
+         */
+        default XmlResourceParser getXml(@XmlRes int id) throws NotFoundException {
+            throw new NotFoundException();
+        }
+
+        /**
+         * Refer {@link Resources#getString(int)}
+         */
+        default String getString(@StringRes int id) throws NotFoundException {
+            throw new NotFoundException();
+        }
+
+        /**
+         * Returns a {@link SourceResources} corresponding to the provided resources
+         */
+        static SourceResources wrap(Resources res) {
+            return new SourceResources() {
+                @Override
+                public XmlResourceParser getXml(int id) {
+                    return res.getXml(id);
+                }
+
+                @Override
+                public String getString(int id) {
+                    return res.getString(id);
+                }
+            };
+        }
+    }
+
+
 }
diff --git a/src/com/android/launcher3/LauncherProvider.java b/src/com/android/launcher3/LauncherProvider.java
index dee3205..d30d23c 100644
--- a/src/com/android/launcher3/LauncherProvider.java
+++ b/src/com/android/launcher3/LauncherProvider.java
@@ -244,14 +244,6 @@
                 mModelDbController.createEmptyDB();
                 return null;
             }
-            case LauncherSettings.Settings.METHOD_SET_USE_TEST_WORKSPACE_LAYOUT_FLAG: {
-                mModelDbController.setUseTestWorkspaceLayout(arg);
-                return null;
-            }
-            case LauncherSettings.Settings.METHOD_CLEAR_USE_TEST_WORKSPACE_LAYOUT_FLAG: {
-                mModelDbController.setUseTestWorkspaceLayout(null);
-                return null;
-            }
             case LauncherSettings.Settings.METHOD_LOAD_DEFAULT_FAVORITES: {
                 mModelDbController.loadDefaultFavoritesIfNecessary();
                 return null;
diff --git a/src/com/android/launcher3/LauncherSettings.java b/src/com/android/launcher3/LauncherSettings.java
index 1bbb09a..1ca3747 100644
--- a/src/com/android/launcher3/LauncherSettings.java
+++ b/src/com/android/launcher3/LauncherSettings.java
@@ -374,15 +374,6 @@
 
         public static final String METHOD_CREATE_EMPTY_DB = "create_empty_db";
 
-        public static final String METHOD_SET_USE_TEST_WORKSPACE_LAYOUT_FLAG =
-                "set_use_test_workspace_layout_flag";
-        public static final String ARG_DEFAULT_WORKSPACE_LAYOUT_TEST = "default_test_workspace";
-        public static final String ARG_DEFAULT_WORKSPACE_LAYOUT_TEST2 = "default_test2_workspace";
-        public static final String ARG_DEFAULT_WORKSPACE_LAYOUT_TAPL = "default_tapl_workspace";
-
-        public static final String METHOD_CLEAR_USE_TEST_WORKSPACE_LAYOUT_FLAG =
-                "clear_use_test_workspace_layout_flag";
-
         public static final String METHOD_LOAD_DEFAULT_FAVORITES = "load_default_favorites";
 
         public static final String METHOD_REMOVE_GHOST_WIDGETS = "remove_ghost_widgets";
@@ -399,6 +390,10 @@
 
         public static final String EXTRA_DB_NAME = "db_name";
 
+        public static final String LAYOUT_DIGEST_KEY = "launcher3.layout.provider.blob";
+        public static final String LAYOUT_DIGEST_LABEL = "launcher-layout";
+        public static final String LAYOUT_DIGEST_TAG = "ignore";
+
         public static Bundle call(ContentResolver cr, String method) {
             return call(cr, method, null /* arg */);
         }
diff --git a/src/com/android/launcher3/model/ModelDbController.java b/src/com/android/launcher3/model/ModelDbController.java
index 7452bcd..9b54ce1 100644
--- a/src/com/android/launcher3/model/ModelDbController.java
+++ b/src/com/android/launcher3/model/ModelDbController.java
@@ -15,38 +15,49 @@
  */
 package com.android.launcher3.model;
 
+import static android.util.Base64.NO_PADDING;
+import static android.util.Base64.NO_WRAP;
+
 import static com.android.launcher3.DefaultLayoutParser.RES_PARTNER_DEFAULT_LAYOUT;
+import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_DIGEST_KEY;
+import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_DIGEST_LABEL;
+import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_DIGEST_TAG;
 import static com.android.launcher3.model.DatabaseHelper.EMPTY_DATABASE_CREATED;
 import static com.android.launcher3.provider.LauncherDbUtils.copyTable;
 import static com.android.launcher3.provider.LauncherDbUtils.tableExists;
 
+import android.app.blob.BlobHandle;
+import android.app.blob.BlobStoreManager;
+import android.content.ContentResolver;
 import android.content.ContentValues;
 import android.content.Context;
 import android.content.SharedPreferences;
+import android.content.pm.PackageManager;
 import android.content.pm.ProviderInfo;
+import android.content.res.Resources;
 import android.database.Cursor;
 import android.database.SQLException;
 import android.database.sqlite.SQLiteDatabase;
 import android.net.Uri;
 import android.os.Binder;
 import android.os.Bundle;
+import android.os.ParcelFileDescriptor;
 import android.os.Process;
 import android.os.UserManager;
 import android.provider.Settings;
 import android.text.TextUtils;
+import android.util.Base64;
 import android.util.Log;
 import android.util.Xml;
 
-import androidx.annotation.Nullable;
-
 import com.android.launcher3.AutoInstallsLayout;
+import com.android.launcher3.AutoInstallsLayout.SourceResources;
 import com.android.launcher3.DefaultLayoutParser;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.LauncherSettings.Favorites;
-import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.provider.LauncherDbUtils;
 import com.android.launcher3.provider.LauncherDbUtils.SQLiteTransaction;
@@ -69,14 +80,7 @@
 public class ModelDbController {
     private static final String TAG = "LauncherProvider";
 
-    private static final int TEST_WORKSPACE_LAYOUT_RES_XML = R.xml.default_test_workspace;
-    private static final int TEST2_WORKSPACE_LAYOUT_RES_XML = R.xml.default_test2_workspace;
-    private static final int TAPL_WORKSPACE_LAYOUT_RES_XML = R.xml.default_tapl_test_workspace;
-
     protected DatabaseHelper mOpenHelper;
-    protected String mProviderAuthority;
-
-    private int mDefaultWorkspaceLayoutOverride = 0;
 
     private final Context mContext;
 
@@ -224,21 +228,6 @@
     }
 
     /**
-     * Overrides the default xml to be used for setting up workspace
-     */
-    public void setUseTestWorkspaceLayout(@Nullable String layout) {
-        if (LauncherSettings.Settings.ARG_DEFAULT_WORKSPACE_LAYOUT_TEST.equals(layout)) {
-            mDefaultWorkspaceLayoutOverride = TEST_WORKSPACE_LAYOUT_RES_XML;
-        } else if (LauncherSettings.Settings.ARG_DEFAULT_WORKSPACE_LAYOUT_TEST2.equals(layout)) {
-            mDefaultWorkspaceLayoutOverride = TEST2_WORKSPACE_LAYOUT_RES_XML;
-        } else if (LauncherSettings.Settings.ARG_DEFAULT_WORKSPACE_LAYOUT_TAPL.equals(layout)) {
-            mDefaultWorkspaceLayoutOverride = TAPL_WORKSPACE_LAYOUT_RES_XML;
-        } else {
-            mDefaultWorkspaceLayoutOverride = 0;
-        }
-    }
-
-    /**
      * Removes any widget which are present in the framework, but not in out internal DB
      */
     public void removeGhostWidgets() {
@@ -403,39 +392,55 @@
      */
     private AutoInstallsLayout createWorkspaceLoaderFromAppRestriction(
             LauncherWidgetHolder widgetHolder) {
-        final String authority;
-        if (!TextUtils.isEmpty(mProviderAuthority)) {
-            authority = mProviderAuthority;
-        } else {
-            authority = Settings.Secure.getString(mContext.getContentResolver(),
-                    "launcher3.layout.provider");
+        ContentResolver cr = mContext.getContentResolver();
+        String blobHandlerDigest = Settings.Secure.getString(cr, LAYOUT_DIGEST_KEY);
+        if (Utilities.ATLEAST_R && !TextUtils.isEmpty(blobHandlerDigest)) {
+            BlobStoreManager blobManager = mContext.getSystemService(BlobStoreManager.class);
+            try (InputStream in = new ParcelFileDescriptor.AutoCloseInputStream(
+                    blobManager.openBlob(BlobHandle.createWithSha256(
+                            Base64.decode(blobHandlerDigest, NO_WRAP | NO_PADDING),
+                            LAYOUT_DIGEST_LABEL, 0, LAYOUT_DIGEST_TAG)))) {
+                return getAutoInstallsLayoutFromIS(in, widgetHolder, new SourceResources() { });
+            } catch (Exception e) {
+                Log.e(TAG, "Error getting layout from blob handle" , e);
+                return null;
+            }
         }
+
+        String authority = Settings.Secure.getString(cr, "launcher3.layout.provider");
         if (TextUtils.isEmpty(authority)) {
             return null;
         }
 
-        ProviderInfo pi = mContext.getPackageManager().resolveContentProvider(authority, 0);
+        PackageManager pm = mContext.getPackageManager();
+        ProviderInfo pi = pm.resolveContentProvider(authority, 0);
         if (pi == null) {
             Log.e(TAG, "No provider found for authority " + authority);
             return null;
         }
         Uri uri = getLayoutUri(authority, mContext);
-        try (InputStream in = mContext.getContentResolver().openInputStream(uri)) {
-            // Read the full xml so that we fail early in case of any IO error.
-            String layout = new String(IOUtils.toByteArray(in));
-            XmlPullParser parser = Xml.newPullParser();
-            parser.setInput(new StringReader(layout));
-
+        try (InputStream in = cr.openInputStream(uri)) {
             Log.d(TAG, "Loading layout from " + authority);
-            return new AutoInstallsLayout(mContext, widgetHolder, mOpenHelper,
-                    mContext.getPackageManager().getResourcesForApplication(pi.applicationInfo),
-                    () -> parser, AutoInstallsLayout.TAG_WORKSPACE);
+
+            Resources res = pm.getResourcesForApplication(pi.applicationInfo);
+            return getAutoInstallsLayoutFromIS(in, widgetHolder, SourceResources.wrap(res));
         } catch (Exception e) {
             Log.e(TAG, "Error getting layout stream from: " + authority , e);
             return null;
         }
     }
 
+    private AutoInstallsLayout getAutoInstallsLayoutFromIS(InputStream in,
+            LauncherWidgetHolder widgetHolder, SourceResources res) throws Exception {
+        // Read the full xml so that we fail early in case of any IO error.
+        String layout = new String(IOUtils.toByteArray(in));
+        XmlPullParser parser = Xml.newPullParser();
+        parser.setInput(new StringReader(layout));
+
+        return new AutoInstallsLayout(mContext, widgetHolder, mOpenHelper, res,
+                () -> parser, AutoInstallsLayout.TAG_WORKSPACE);
+    }
+
     private static Uri getLayoutUri(String authority, Context ctx) {
         InvariantDeviceProfile grid = LauncherAppState.getIDP(ctx);
         return new Uri.Builder().scheme("content").authority(authority).path("launcher_layout")
@@ -448,13 +453,9 @@
 
     private DefaultLayoutParser getDefaultLayoutParser(LauncherWidgetHolder widgetHolder) {
         InvariantDeviceProfile idp = LauncherAppState.getIDP(mContext);
-        int defaultLayout = mDefaultWorkspaceLayoutOverride > 0
-                ? mDefaultWorkspaceLayoutOverride : idp.defaultLayoutId;
-
-        if (mContext.getSystemService(UserManager.class).isDemoUser()
-                && idp.demoModeLayoutId != 0) {
-            defaultLayout = idp.demoModeLayoutId;
-        }
+        int defaultLayout = idp.demoModeLayoutId != 0
+                && mContext.getSystemService(UserManager.class).isDemoUser()
+                ? idp.demoModeLayoutId : idp.defaultLayoutId;
 
         return new DefaultLayoutParser(mContext, widgetHolder,
                 mOpenHelper, mContext.getResources(), defaultLayout);
diff --git a/tests/Android.bp b/tests/Android.bp
index 81853d1..da8d844 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -45,6 +45,7 @@
       "src/com/android/launcher3/ui/AbstractLauncherUiTest.java",
       "src/com/android/launcher3/ui/PortraitLandscapeRunner.java",
       "src/com/android/launcher3/ui/TaplTestsLauncher3.java",
+      "src/com/android/launcher3/util/LauncherLayoutBuilder.java",
       "src/com/android/launcher3/util/TestUtil.java",
       "src/com/android/launcher3/util/Wait.java",
       "src/com/android/launcher3/util/WidgetUtils.java",
diff --git a/tests/shared/com/android/launcher3/testing/shared/TestProtocol.java b/tests/shared/com/android/launcher3/testing/shared/TestProtocol.java
index a37c3cd..dc835e2 100644
--- a/tests/shared/com/android/launcher3/testing/shared/TestProtocol.java
+++ b/tests/shared/com/android/launcher3/testing/shared/TestProtocol.java
@@ -106,11 +106,6 @@
     public static final String REQUEST_GET_HAD_NONTEST_EVENTS = "get-had-nontest-events";
     public static final String REQUEST_STOP_EVENT_LOGGING = "stop-event-logging";
     public static final String REQUEST_CLEAR_DATA = "clear-data";
-    public static final String REQUEST_USE_TEST_WORKSPACE_LAYOUT = "use-test-workspace-layout";
-    public static final String REQUEST_USE_TEST2_WORKSPACE_LAYOUT = "use-test2-workspace-layout";
-    public static final String REQUEST_USE_TAPL_WORKSPACE_LAYOUT = "use-tapl-workspace-layout";
-    public static final String REQUEST_USE_DEFAULT_WORKSPACE_LAYOUT =
-            "use-default-workspace-layout";
     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_IS_TWO_PANELS = "is-two-panel";
diff --git a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
index 217bec3..81f1525 100644
--- a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
+++ b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
@@ -54,12 +54,14 @@
 import com.android.launcher3.tapl.HomeAppIconMenuItem;
 import com.android.launcher3.tapl.Widgets;
 import com.android.launcher3.tapl.Workspace;
+import com.android.launcher3.util.LauncherLayoutBuilder;
 import com.android.launcher3.util.TestUtil;
 import com.android.launcher3.util.rule.ScreenRecordRule.ScreenRecord;
 import com.android.launcher3.util.rule.TISBindRule;
 import com.android.launcher3.widget.picker.WidgetsFullSheet;
 import com.android.launcher3.widget.picker.WidgetsRecyclerView;
 
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Ignore;
 import org.junit.Rule;
@@ -83,6 +85,8 @@
     @Rule
     public TISBindRule mTISBindRule = new TISBindRule();
 
+    private AutoCloseable mLauncherLayout;
+
     @Before
     public void setUp() throws Exception {
         super.setUp();
@@ -101,6 +105,13 @@
         AbstractLauncherUiTest.checkDetectedLeaks(test.mLauncher);
     }
 
+    @After
+    public void tearDown() throws Exception {
+        if (mLauncherLayout != null) {
+            mLauncherLayout.close();
+        }
+    }
+
     // Please don't add negative test cases for methods that fail only after a long wait.
     public static void expectFail(String message, Runnable action) {
         boolean failed = false;
@@ -230,8 +241,10 @@
     @Test
     @ScreenRecord // b/202433017
     public void testWorkspace() throws Exception {
-        // Make sure there is an instance of chrome on the hotseat
-        mLauncher.useTaplWorkspaceLayoutOnReload();
+        // Set workspace  that includes the chrome Activity app icon on the hotseat.
+        LauncherLayoutBuilder builder = new LauncherLayoutBuilder()
+                .atHotseat(0).putApp("com.android.chrome", "com.google.android.apps.chrome.Main");
+        mLauncherLayout = TestUtil.setLauncherDefaultLayout(mTargetContext, builder);
         clearLauncherData();
 
         final Workspace workspace = mLauncher.getWorkspace();
diff --git a/tests/src/com/android/launcher3/ui/workspace/TwoPanelWorkspaceTest.java b/tests/src/com/android/launcher3/ui/workspace/TwoPanelWorkspaceTest.java
index 3abafdf..c4b6d43 100644
--- a/tests/src/com/android/launcher3/ui/workspace/TwoPanelWorkspaceTest.java
+++ b/tests/src/com/android/launcher3/ui/workspace/TwoPanelWorkspaceTest.java
@@ -30,6 +30,8 @@
 import com.android.launcher3.tapl.Workspace;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.ui.TaplTestsLauncher3;
+import com.android.launcher3.util.LauncherLayoutBuilder;
+import com.android.launcher3.util.TestUtil;
 
 import org.junit.After;
 import org.junit.Before;
@@ -49,12 +51,24 @@
 @RunWith(AndroidJUnit4.class)
 public class TwoPanelWorkspaceTest extends AbstractLauncherUiTest {
 
+    private AutoCloseable mLauncherLayout;
+
     @Before
     public void setUp() throws Exception {
         super.setUp();
-        mLauncher.useTest2WorkspaceLayoutOnReload();
-        TaplTestsLauncher3.initialize(this);
 
+        // Set layout that includes Maps/Play on workspace, and Messaging/Chrome on hotseat.
+        LauncherLayoutBuilder builder = new LauncherLayoutBuilder()
+                .atHotseat(0).putApp(
+                        "com.google.android.apps.messaging",
+                        "com.google.android.apps.messaging.ui.ConversationListActivity")
+                .atHotseat(1).putApp("com.android.chrome", "com.google.android.apps.chrome.Main")
+                .atWorkspace(0, -1, 0).putApp(
+                        "com.google.android.apps.maps", "com.google.android.maps.MapsActivity")
+                .atWorkspace(3, -1, 0).putApp(
+                        "com.android.vending", "com.android.vending.AssetBrowserActivity");
+        mLauncherLayout = TestUtil.setLauncherDefaultLayout(mTargetContext, builder);
+        TaplTestsLauncher3.initialize(this);
         assumeTrue(mLauncher.isTwoPanels());
 
         // Pre verifying the screens
@@ -67,9 +81,11 @@
     }
 
     @After
-    public void tearDown() {
+    public void tearDown() throws Exception {
         executeOnLauncher(launcher -> launcher.enableHotseatEdu(true));
-        mLauncher.useDefaultWorkspaceLayoutOnReload();
+        if (mLauncherLayout != null) {
+            mLauncherLayout.close();
+        }
     }
 
     @Test
diff --git a/tests/src/com/android/launcher3/util/TestUtil.java b/tests/src/com/android/launcher3/util/TestUtil.java
index 433fd31..4981795 100644
--- a/tests/src/com/android/launcher3/util/TestUtil.java
+++ b/tests/src/com/android/launcher3/util/TestUtil.java
@@ -15,15 +15,29 @@
  */
 package com.android.launcher3.util;
 
+import static android.util.Base64.NO_PADDING;
+import static android.util.Base64.NO_WRAP;
+
 import static androidx.test.InstrumentationRegistry.getContext;
 import static androidx.test.InstrumentationRegistry.getInstrumentation;
 import static androidx.test.InstrumentationRegistry.getTargetContext;
 
+import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_DIGEST_KEY;
+import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_DIGEST_LABEL;
+import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_DIGEST_TAG;
+
+import android.app.blob.BlobHandle;
+import android.app.blob.BlobStoreManager;
+import android.content.Context;
 import android.content.pm.LauncherApps;
 import android.content.res.Resources;
+import android.os.AsyncTask;
 import android.os.Handler;
 import android.os.Looper;
+import android.os.ParcelFileDescriptor.AutoCloseOutputStream;
 import android.os.UserHandle;
+import android.provider.Settings;
+import android.util.Base64;
 
 import androidx.test.uiautomator.UiDevice;
 
@@ -36,6 +50,8 @@
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.OutputStream;
+import java.security.MessageDigest;
 import java.util.concurrent.CountDownLatch;
 import java.util.function.Predicate;
 import java.util.function.ToIntFunction;
@@ -109,6 +125,35 @@
                 "pm uninstall " + DUMMY_PACKAGE);
     }
 
+    /**
+     * Sets the default layout for Launcher and returns an object which can be used to clear
+     * the data
+     */
+    public static AutoCloseable setLauncherDefaultLayout(
+            Context context, LauncherLayoutBuilder layoutBuilder) throws Exception {
+        byte[] data = layoutBuilder.build().getBytes();
+        byte[] digest = MessageDigest.getInstance("SHA-256").digest(data);
+
+        BlobHandle handle = BlobHandle.createWithSha256(
+                digest, LAYOUT_DIGEST_LABEL, 0, LAYOUT_DIGEST_TAG);
+        BlobStoreManager blobManager = context.getSystemService(BlobStoreManager.class);
+        final long sessionId = blobManager.createSession(handle);
+        CountDownLatch wait = new CountDownLatch(1);
+        try (BlobStoreManager.Session session = blobManager.openSession(sessionId)) {
+            try (OutputStream out = new AutoCloseOutputStream(session.openWrite(0, -1))) {
+                out.write(data);
+            }
+            session.allowPublicAccess();
+            session.commit(AsyncTask.THREAD_POOL_EXECUTOR, i -> wait.countDown());
+        }
+
+        String key = Base64.encodeToString(digest, NO_WRAP | NO_PADDING);
+        Settings.Secure.putString(context.getContentResolver(), LAYOUT_DIGEST_KEY, key);
+        wait.await();
+        return () ->
+            Settings.Secure.putString(context.getContentResolver(), LAYOUT_DIGEST_KEY, null);
+    }
+
     private static class PackageInstallCheck extends LauncherApps.Callback
             implements AutoCloseable {
 
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index 9905603..ba8f070 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -1851,36 +1851,6 @@
         getTestInfo(TestProtocol.REQUEST_CLEAR_DATA);
     }
 
-    /**
-     * Reloads the workspace with a test layout that includes the Test Activity app icon on the
-     * hotseat.
-     */
-    public void useTestWorkspaceLayoutOnReload() {
-        getTestInfo(TestProtocol.REQUEST_USE_TEST_WORKSPACE_LAYOUT);
-    }
-
-    /**
-     * Reloads the workspace with a test layout that includes Maps/Play on workspace, and
-     * Dialer/Messaging/Chrome/Camera on hotseat.
-     */
-    public void useTest2WorkspaceLayoutOnReload() {
-        getTestInfo(TestProtocol.REQUEST_USE_TEST2_WORKSPACE_LAYOUT);
-    }
-
-
-    /**
-     * Reloads the workspace with a test layout that includes the chrome Activity app icon on the
-     * hotseat.
-     */
-    public void useTaplWorkspaceLayoutOnReload() {
-        getTestInfo(TestProtocol.REQUEST_USE_TAPL_WORKSPACE_LAYOUT);
-    }
-
-    /** Reloads the workspace with the default layout defined by the user's grid size selection. */
-    public void useDefaultWorkspaceLayoutOnReload() {
-        getTestInfo(TestProtocol.REQUEST_USE_DEFAULT_WORKSPACE_LAYOUT);
-    }
-
     /** Shows the taskbar if it is hidden, otherwise does nothing. */
     public void showTaskbarIfHidden() {
         getTestInfo(TestProtocol.REQUEST_UNSTASH_TASKBAR_IF_STASHED);