Create a separate class for widget related methods from popup provider

Child cl removes those methods from popup provider and updates its
references

Bug: 353347512
Flag: EXEMPT BUGFIX
Test: Unit test
Change-Id: Id559e8c8d32a40adab961e7c710ec835d5a92184
diff --git a/src/com/android/launcher3/widget/picker/WidgetRecommendationCategory.java b/src/com/android/launcher3/widget/picker/WidgetRecommendationCategory.java
index 072d1d5..a68effd 100644
--- a/src/com/android/launcher3/widget/picker/WidgetRecommendationCategory.java
+++ b/src/com/android/launcher3/widget/picker/WidgetRecommendationCategory.java
@@ -19,6 +19,8 @@
 import androidx.annotation.Nullable;
 import androidx.annotation.StringRes;
 
+import com.android.launcher3.R;
+
 import java.util.Objects;
 
 /**
@@ -26,6 +28,10 @@
  * option in the pop-up opened on long press of launcher workspace).
  */
 public class WidgetRecommendationCategory implements Comparable<WidgetRecommendationCategory> {
+    public static WidgetRecommendationCategory DEFAULT_WIDGET_RECOMMENDATION_CATEGORY =
+            new WidgetRecommendationCategory(
+                    R.string.others_widget_recommendation_category_label, /*order=*/0);
+
     /** Resource id that holds the user friendly label for the category. */
     @StringRes
     public final int categoryTitleRes;
diff --git a/src/com/android/launcher3/widget/picker/model/WidgetPickerDataProvider.kt b/src/com/android/launcher3/widget/picker/model/WidgetPickerDataProvider.kt
new file mode 100644
index 0000000..46d3e7a
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/model/WidgetPickerDataProvider.kt
@@ -0,0 +1,84 @@
+/*
+ * 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.widget.picker.model
+
+import com.android.launcher3.model.data.ItemInfo
+import com.android.launcher3.widget.model.WidgetsListBaseEntry
+import com.android.launcher3.widget.picker.model.data.WidgetPickerData
+import com.android.launcher3.widget.picker.model.data.WidgetPickerDataUtils.withRecommendedWidgets
+import com.android.launcher3.widget.picker.model.data.WidgetPickerDataUtils.withWidgets
+import java.io.PrintWriter
+
+/**
+ * Provides [WidgetPickerData] to various views such as widget picker, app-specific widget picker,
+ * widgets shortcut.
+ */
+class WidgetPickerDataProvider {
+    /** All the widgets data provided for the views */
+    private var mWidgetPickerData: WidgetPickerData = WidgetPickerData()
+
+    private var changeListener: WidgetPickerDataChangeListener? = null
+
+    /** Sets a listener to be called back when widget data is updated. */
+    fun setChangeListener(changeListener: WidgetPickerDataChangeListener?) {
+        this.changeListener = changeListener
+    }
+
+    /** Returns the current snapshot of [WidgetPickerData]. */
+    fun get(): WidgetPickerData {
+        return mWidgetPickerData
+    }
+
+    /**
+     * Updates the widgets available to the widget picker.
+     *
+     * Generally called when the widgets model has new data.
+     */
+    @JvmOverloads
+    fun setWidgets(
+        allWidgets: List<WidgetsListBaseEntry>,
+        defaultWidgets: List<WidgetsListBaseEntry> = listOf()
+    ) {
+        mWidgetPickerData =
+            mWidgetPickerData.withWidgets(allWidgets = allWidgets, defaultWidgets = defaultWidgets)
+        changeListener?.onWidgetsBound()
+    }
+
+    /**
+     * Makes the widget recommendations available to the widget picker
+     *
+     * Generally called when new widget predictions are available.
+     */
+    fun setWidgetRecommendations(recommendations: List<ItemInfo>) {
+        mWidgetPickerData = mWidgetPickerData.withRecommendedWidgets(recommendations)
+        changeListener?.onRecommendedWidgetsBound()
+    }
+
+    /** Writes the current state to the provided writer. */
+    fun dump(prefix: String, writer: PrintWriter) {
+        writer.println(prefix + "WidgetPickerDataProvider:")
+        writer.println("$prefix\twidgetPickerData:$mWidgetPickerData")
+    }
+
+    interface WidgetPickerDataChangeListener {
+        /** A callback to get notified when widgets are bound. */
+        fun onWidgetsBound()
+
+        /** A callback to get notified when recommended widgets are bound. */
+        fun onRecommendedWidgetsBound()
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/model/data/WidgetPickerData.kt b/src/com/android/launcher3/widget/picker/model/data/WidgetPickerData.kt
new file mode 100644
index 0000000..3332ef0
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/model/data/WidgetPickerData.kt
@@ -0,0 +1,110 @@
+/*
+ * 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.widget.picker.model.data
+
+import com.android.launcher3.model.WidgetItem
+import com.android.launcher3.model.data.ItemInfo
+import com.android.launcher3.util.ComponentKey
+import com.android.launcher3.util.PackageUserKey
+import com.android.launcher3.widget.PendingAddWidgetInfo
+import com.android.launcher3.widget.model.WidgetsListBaseEntry
+import com.android.launcher3.widget.model.WidgetsListContentEntry
+import com.android.launcher3.widget.picker.WidgetRecommendationCategory
+import com.android.launcher3.widget.picker.WidgetRecommendationCategory.DEFAULT_WIDGET_RECOMMENDATION_CATEGORY
+
+// This file contains WidgetPickerData and utility functions to operate on it.
+
+/** Widget data for display in the widget picker. */
+data class WidgetPickerData(
+    val allWidgets: List<WidgetsListBaseEntry> = listOf(),
+    val defaultWidgets: List<WidgetsListBaseEntry> = listOf(),
+    val recommendations: Map<WidgetRecommendationCategory, List<WidgetItem>> = mapOf(),
+)
+
+/** Provides utility methods to work with a [WidgetPickerData] object. */
+object WidgetPickerDataUtils {
+    /**
+     * Returns a [WidgetPickerData] with the provided widgets.
+     *
+     * When [defaultWidgets] is not passed, defaults from previous object are not copied over.
+     * Defaults (if supported) should be updated when all widgets are updated.
+     */
+    fun WidgetPickerData.withWidgets(
+        allWidgets: List<WidgetsListBaseEntry>,
+        defaultWidgets: List<WidgetsListBaseEntry> = listOf()
+    ): WidgetPickerData {
+        return copy(allWidgets = allWidgets, defaultWidgets = defaultWidgets)
+    }
+
+    /** Returns a [WidgetPickerData] with the given recommendations set. */
+    fun WidgetPickerData.withRecommendedWidgets(recommendations: List<ItemInfo>): WidgetPickerData {
+        val allWidgetsMap: Map<ComponentKey, WidgetItem> =
+            allWidgets
+                .filterIsInstance<WidgetsListContentEntry>()
+                .flatMap { it.mWidgets }
+                .filterNotNull()
+                .distinct()
+                .associateBy { it } // as ComponentKey
+
+        val categoriesMap =
+            recommendations
+                .filterIsInstance<PendingAddWidgetInfo>()
+                .filter { allWidgetsMap.containsKey(ComponentKey(it.targetComponent, it.user)) }
+                .groupBy { it.recommendationCategory ?: DEFAULT_WIDGET_RECOMMENDATION_CATEGORY }
+                .mapValues { (_, pendingAddWidgetInfos) ->
+                    pendingAddWidgetInfos.map {
+                        allWidgetsMap[ComponentKey(it.targetComponent, it.user)] as WidgetItem
+                    }
+                }
+
+        return copy(recommendations = categoriesMap)
+    }
+
+    /** Finds all [WidgetItem]s available for the provided package user. */
+    @JvmStatic
+    fun findAllWidgetsForPackageUser(
+        widgetPickerData: WidgetPickerData,
+        packageUserKey: PackageUserKey
+    ): List<WidgetItem> {
+        return findContentEntryForPackageUser(widgetPickerData, packageUserKey)?.mWidgets
+            ?: emptyList()
+    }
+
+    /**
+     * Finds and returns the [WidgetsListContentEntry] for the given package user.
+     *
+     * Set [fromDefaultWidgets] to true to limit the content entry to default widgets.
+     */
+    @JvmOverloads
+    @JvmStatic
+    fun findContentEntryForPackageUser(
+        widgetPickerData: WidgetPickerData,
+        packageUserKey: PackageUserKey,
+        fromDefaultWidgets: Boolean = false
+    ): WidgetsListContentEntry? {
+        val widgetsListBaseEntries =
+            if (fromDefaultWidgets) {
+                widgetPickerData.defaultWidgets
+            } else {
+                widgetPickerData.allWidgets
+            }
+
+        return widgetsListBaseEntries.filterIsInstance<WidgetsListContentEntry>().firstOrNull {
+            PackageUserKey.fromPackageItemInfo(it.mPkgItem) == packageUserKey
+        }
+    }
+}
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/WidgetPickerDataProviderTest.kt b/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/WidgetPickerDataProviderTest.kt
new file mode 100644
index 0000000..1822639
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/WidgetPickerDataProviderTest.kt
@@ -0,0 +1,170 @@
+/*
+ * 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.widget.picker.model
+
+import android.content.ComponentName
+import android.content.Context
+import android.os.UserHandle
+import android.platform.test.rule.AllowedDevices
+import android.platform.test.rule.DeviceProduct
+import android.platform.test.rule.LimitDevicesRule
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.launcher3.InvariantDeviceProfile
+import com.android.launcher3.LauncherAppState
+import com.android.launcher3.LauncherSettings
+import com.android.launcher3.icons.ComponentWithLabel
+import com.android.launcher3.icons.IconCache
+import com.android.launcher3.model.WidgetItem
+import com.android.launcher3.model.data.PackageItemInfo
+import com.android.launcher3.util.ActivityContextWrapper
+import com.android.launcher3.util.WidgetUtils
+import com.android.launcher3.widget.LauncherAppWidgetProviderInfo
+import com.android.launcher3.widget.PendingAddWidgetInfo
+import com.android.launcher3.widget.model.WidgetsListBaseEntry
+import com.android.launcher3.widget.model.WidgetsListContentEntry
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry
+import com.android.launcher3.widget.picker.model.WidgetPickerDataProvider.WidgetPickerDataChangeListener
+import com.google.common.truth.Truth.assertThat
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.invocation.InvocationOnMock
+import org.mockito.junit.MockitoJUnit
+import org.mockito.junit.MockitoRule
+import org.mockito.kotlin.any
+import org.mockito.kotlin.doAnswer
+import org.mockito.kotlin.times
+import org.mockito.kotlin.verify
+import org.mockito.kotlin.verifyNoMoreInteractions
+
+// Tests for the WidgetPickerDataProvider class
+
+@RunWith(AndroidJUnit4::class)
+@AllowedDevices(allowed = [DeviceProduct.ROBOLECTRIC])
+class WidgetPickerDataProviderTest {
+    @Rule @JvmField val limitDevicesRule = LimitDevicesRule()
+    @Rule @JvmField val mockitoRule: MockitoRule = MockitoJUnit.rule()
+
+    @Mock private lateinit var changeListener: WidgetPickerDataChangeListener
+
+    @Mock private lateinit var iconCache: IconCache
+
+    private lateinit var userHandle: UserHandle
+    private lateinit var context: Context
+    private lateinit var testInvariantProfile: InvariantDeviceProfile
+
+    private lateinit var appWidgetItem: WidgetItem
+
+    private var underTest = WidgetPickerDataProvider()
+
+    @Before
+    fun setUp() {
+        userHandle = UserHandle.CURRENT
+        context = ActivityContextWrapper(ApplicationProvider.getApplicationContext())
+        testInvariantProfile = LauncherAppState.getIDP(context)
+
+        doAnswer { invocation: InvocationOnMock ->
+                val componentWithLabel = invocation.getArgument<Any>(0) as ComponentWithLabel
+                componentWithLabel.getComponent().shortClassName
+            }
+            .`when`(iconCache)
+            .getTitleNoCache(any<ComponentWithLabel>())
+
+        appWidgetItem = createWidgetItem()
+    }
+
+    @Test
+    fun setWidgets_invokesTheListener_andUpdatedWidgetsAvailable() {
+        assertThat(underTest.get().allWidgets).isEmpty()
+
+        underTest.setChangeListener(changeListener)
+        val allWidgets = appWidgetListBaseEntries()
+        underTest.setWidgets(allWidgets = allWidgets)
+
+        assertThat(underTest.get().allWidgets).containsExactlyElementsIn(allWidgets)
+        verify(changeListener, times(1)).onWidgetsBound()
+        verifyNoMoreInteractions(changeListener)
+    }
+
+    @Test
+    fun setWidgetRecommendations_callsBackTheListener_andUpdatedRecommendationsAvailable() {
+        underTest.setWidgets(allWidgets = appWidgetListBaseEntries())
+        assertThat(underTest.get().recommendations).isEmpty()
+
+        underTest.setChangeListener(changeListener)
+        val recommendations =
+            listOf(
+                PendingAddWidgetInfo(
+                    appWidgetItem.widgetInfo,
+                    LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION
+                ),
+            )
+        underTest.setWidgetRecommendations(recommendations)
+
+        assertThat(underTest.get().recommendations).hasSize(1)
+        verify(changeListener, times(1)).onRecommendedWidgetsBound()
+        verifyNoMoreInteractions(changeListener)
+    }
+
+    @Test
+    fun setChangeListener_null_noCallback() {
+        underTest.setChangeListener(changeListener)
+        underTest.setChangeListener(null) // reset
+
+        underTest.setWidgets(allWidgets = appWidgetListBaseEntries())
+        val recommendations =
+            listOf(
+                PendingAddWidgetInfo(
+                    appWidgetItem.widgetInfo,
+                    LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION
+                ),
+            )
+        underTest.setWidgetRecommendations(recommendations)
+
+        verifyNoMoreInteractions(changeListener)
+    }
+
+    private fun createWidgetItem(): WidgetItem {
+        val providerInfo =
+            WidgetUtils.createAppWidgetProviderInfo(
+                ComponentName.createRelative(APP_PACKAGE_NAME, APP_PROVIDER_1_CLASS_NAME)
+            )
+        val widgetInfo = LauncherAppWidgetProviderInfo.fromProviderInfo(context, providerInfo)
+        return WidgetItem(widgetInfo, testInvariantProfile, iconCache, context)
+    }
+
+    private fun appWidgetListBaseEntries(): List<WidgetsListBaseEntry> {
+        val packageItemInfo = PackageItemInfo(APP_PACKAGE_NAME, userHandle)
+        packageItemInfo.title = APP_PACKAGE_TITLE
+        val widgets = listOf(appWidgetItem)
+
+        return buildList {
+            add(WidgetsListHeaderEntry.create(packageItemInfo, APP_SECTION_NAME, widgets))
+            add(WidgetsListContentEntry(packageItemInfo, APP_SECTION_NAME, widgets))
+        }
+    }
+
+    companion object {
+        const val APP_PACKAGE_NAME = "com.example.app"
+        const val APP_PACKAGE_TITLE = "SomeApp"
+        const val APP_SECTION_NAME = "S" // for fast popup
+        const val APP_PROVIDER_1_CLASS_NAME = "appProvider1"
+    }
+}
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/data/WidgetPickerDataTest.kt b/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/data/WidgetPickerDataTest.kt
new file mode 100644
index 0000000..e59e211
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/data/WidgetPickerDataTest.kt
@@ -0,0 +1,379 @@
+/*
+ * 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.widget.picker.model.data
+
+import android.content.ComponentName
+import android.content.Context
+import android.os.UserHandle
+import android.platform.test.rule.AllowedDevices
+import android.platform.test.rule.DeviceProduct
+import android.platform.test.rule.LimitDevicesRule
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.launcher3.InvariantDeviceProfile
+import com.android.launcher3.LauncherAppState
+import com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION
+import com.android.launcher3.icons.ComponentWithLabel
+import com.android.launcher3.icons.IconCache
+import com.android.launcher3.model.WidgetItem
+import com.android.launcher3.model.data.ItemInfo
+import com.android.launcher3.model.data.PackageItemInfo
+import com.android.launcher3.util.ActivityContextWrapper
+import com.android.launcher3.util.PackageUserKey
+import com.android.launcher3.util.WidgetUtils
+import com.android.launcher3.widget.LauncherAppWidgetProviderInfo
+import com.android.launcher3.widget.PendingAddWidgetInfo
+import com.android.launcher3.widget.model.WidgetsListBaseEntry
+import com.android.launcher3.widget.model.WidgetsListContentEntry
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry
+import com.android.launcher3.widget.picker.WidgetRecommendationCategory
+import com.android.launcher3.widget.picker.WidgetRecommendationCategory.DEFAULT_WIDGET_RECOMMENDATION_CATEGORY
+import com.android.launcher3.widget.picker.model.data.WidgetPickerDataUtils.findAllWidgetsForPackageUser
+import com.android.launcher3.widget.picker.model.data.WidgetPickerDataUtils.findContentEntryForPackageUser
+import com.android.launcher3.widget.picker.model.data.WidgetPickerDataUtils.withRecommendedWidgets
+import com.android.launcher3.widget.picker.model.data.WidgetPickerDataUtils.withWidgets
+import com.google.common.truth.Truth.assertThat
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.invocation.InvocationOnMock
+import org.mockito.junit.MockitoJUnit
+import org.mockito.junit.MockitoRule
+import org.mockito.kotlin.any
+import org.mockito.kotlin.doAnswer
+
+// Tests for code / classes in WidgetPickerData file.
+
+@RunWith(AndroidJUnit4::class)
+@AllowedDevices(allowed = [DeviceProduct.ROBOLECTRIC])
+class WidgetPickerDataTest {
+    @Rule @JvmField val limitDevicesRule = LimitDevicesRule()
+    @Rule @JvmField val mockitoRule: MockitoRule = MockitoJUnit.rule()
+
+    @Mock private lateinit var iconCache: IconCache
+
+    private lateinit var userHandle: UserHandle
+    private lateinit var context: Context
+    private lateinit var testInvariantProfile: InvariantDeviceProfile
+
+    private lateinit var app1PackageItemInfo: PackageItemInfo
+    private lateinit var app2PackageItemInfo: PackageItemInfo
+
+    private lateinit var app1WidgetItem1: WidgetItem
+    private lateinit var app1WidgetItem2: WidgetItem
+    private lateinit var app2WidgetItem1: WidgetItem
+
+    @Before
+    fun setUp() {
+        userHandle = UserHandle.CURRENT
+        context = ActivityContextWrapper(ApplicationProvider.getApplicationContext())
+        testInvariantProfile = LauncherAppState.getIDP(context)
+
+        doAnswer { invocation: InvocationOnMock ->
+                val componentWithLabel = invocation.getArgument<Any>(0) as ComponentWithLabel
+                componentWithLabel.getComponent().shortClassName
+            }
+            .`when`(iconCache)
+            .getTitleNoCache(any<ComponentWithLabel>())
+
+        app1PackageItemInfo = packageItemInfoWithTitle(APP_1_PACKAGE_NAME, APP_1_PACKAGE_TITLE)
+        app2PackageItemInfo = packageItemInfoWithTitle(APP_2_PACKAGE_NAME, APP_2_PACKAGE_TITLE)
+
+        app1WidgetItem1 = createWidgetItem(APP_1_PACKAGE_NAME, APP_1_PROVIDER_1_CLASS_NAME)
+        app1WidgetItem2 = createWidgetItem(APP_1_PACKAGE_NAME, APP_1_PROVIDER_2_CLASS_NAME)
+        app2WidgetItem1 = createWidgetItem(APP_2_PACKAGE_NAME, APP_2_PROVIDER_1_CLASS_NAME)
+    }
+
+    @Test
+    fun withWidgets_returnsACopyWithProvidedWidgets() {
+        // only app two
+        val widgetPickerData = WidgetPickerData(allWidgets = appTwoWidgetsListBaseEntries())
+
+        // update: only app 1 and default list set
+        val newAllWidgets: List<WidgetsListBaseEntry> =
+            appOneWidgetsListBaseEntries(includeWidgetTwo = true)
+        val newDefaultWidgets: List<WidgetsListBaseEntry> =
+            appOneWidgetsListBaseEntries(includeWidgetTwo = false)
+
+        val newWidgetData = widgetPickerData.withWidgets(newAllWidgets, newDefaultWidgets)
+
+        assertThat(newWidgetData.allWidgets).containsExactlyElementsIn(newAllWidgets)
+        assertThat(newWidgetData.defaultWidgets).containsExactlyElementsIn(newDefaultWidgets)
+    }
+
+    @Test
+    fun withWidgets_noExplicitDefaults_unsetsOld() {
+        // only app two
+        val widgetPickerData =
+            WidgetPickerData(
+                allWidgets = appTwoWidgetsListBaseEntries(),
+                defaultWidgets = appTwoWidgetsListBaseEntries()
+            )
+
+        val newWidgetData =
+            widgetPickerData.withWidgets(allWidgets = appOneWidgetsListBaseEntries())
+
+        assertThat(newWidgetData.allWidgets)
+            .containsExactlyElementsIn(appOneWidgetsListBaseEntries())
+        assertThat(newWidgetData.defaultWidgets).isEmpty() // previous values cleared.
+    }
+
+    @Test
+    fun withRecommendedWidgets_returnsACopyWithProvidedRecommendedWidgets() {
+        val widgetPickerData =
+            WidgetPickerData(
+                allWidgets =
+                    buildList {
+                        addAll(appOneWidgetsListBaseEntries())
+                        addAll(appTwoWidgetsListBaseEntries())
+                    },
+                defaultWidgets = buildList { appTwoWidgetsListBaseEntries() }
+            )
+        val recommendations: List<ItemInfo> =
+            listOf(
+                PendingAddWidgetInfo(
+                    app1WidgetItem1.widgetInfo,
+                    CONTAINER_WIDGETS_PREDICTION,
+                    CATEGORY_1
+                ),
+                PendingAddWidgetInfo(
+                    app2WidgetItem1.widgetInfo,
+                    CONTAINER_WIDGETS_PREDICTION,
+                    CATEGORY_2
+                ),
+            )
+
+        val updatedData = widgetPickerData.withRecommendedWidgets(recommendations)
+
+        assertThat(updatedData.recommendations.keys).containsExactly(CATEGORY_1, CATEGORY_2)
+        assertThat(updatedData.recommendations[CATEGORY_1]).containsExactly(app1WidgetItem1)
+        assertThat(updatedData.recommendations[CATEGORY_2]).containsExactly(app2WidgetItem1)
+    }
+
+    @Test
+    fun withRecommendedWidgets_noCategory_usesDefault() {
+        val widgetPickerData =
+            WidgetPickerData(
+                allWidgets =
+                    buildList {
+                        addAll(appOneWidgetsListBaseEntries())
+                        addAll(appTwoWidgetsListBaseEntries())
+                    },
+                defaultWidgets = buildList { appTwoWidgetsListBaseEntries() }
+            )
+        val recommendations: List<ItemInfo> =
+            listOf(
+                PendingAddWidgetInfo(app1WidgetItem1.widgetInfo, CONTAINER_WIDGETS_PREDICTION),
+                PendingAddWidgetInfo(app2WidgetItem1.widgetInfo, CONTAINER_WIDGETS_PREDICTION),
+            )
+
+        val updatedData = widgetPickerData.withRecommendedWidgets(recommendations)
+
+        assertThat(updatedData.recommendations.keys)
+            .containsExactly(DEFAULT_WIDGET_RECOMMENDATION_CATEGORY)
+        assertThat(updatedData.recommendations[DEFAULT_WIDGET_RECOMMENDATION_CATEGORY])
+            .containsExactly(app1WidgetItem1, app2WidgetItem1)
+    }
+
+    @Test
+    fun withRecommendedWidgets_emptyRecommendations_clearsOld() {
+        val widgetPickerData =
+            WidgetPickerData(
+                allWidgets =
+                    buildList {
+                        addAll(appOneWidgetsListBaseEntries())
+                        addAll(appTwoWidgetsListBaseEntries())
+                    },
+                defaultWidgets = buildList { appTwoWidgetsListBaseEntries() },
+                recommendations = mapOf(CATEGORY_1 to listOf(app1WidgetItem1))
+            )
+
+        val updatedData = widgetPickerData.withRecommendedWidgets(listOf())
+
+        assertThat(updatedData.recommendations).isEmpty()
+    }
+
+    @Test
+    fun withRecommendedWidgets_widgetNotInAllWidgets_filteredOut() {
+        val widgetPickerData =
+            WidgetPickerData(
+                allWidgets =
+                    buildList {
+                        addAll(appOneWidgetsListBaseEntries(includeWidgetTwo = false))
+                        addAll(appTwoWidgetsListBaseEntries())
+                    },
+                defaultWidgets = buildList { appTwoWidgetsListBaseEntries() },
+            )
+
+        val recommendations: List<ItemInfo> =
+            listOf(
+                PendingAddWidgetInfo(app1WidgetItem2.widgetInfo, CONTAINER_WIDGETS_PREDICTION),
+                PendingAddWidgetInfo(app2WidgetItem1.widgetInfo, CONTAINER_WIDGETS_PREDICTION),
+            )
+        val updatedData = widgetPickerData.withRecommendedWidgets(recommendations)
+
+        assertThat(updatedData.recommendations).hasSize(1)
+        // no app1widget2
+        assertThat(updatedData.recommendations.values.first()).containsExactly(app2WidgetItem1)
+    }
+
+    @Test
+    fun findContentEntryForPackageUser_returnsCorrectEntry() {
+        val widgetPickerData =
+            WidgetPickerData(
+                allWidgets =
+                    buildList {
+                        addAll(appOneWidgetsListBaseEntries())
+                        addAll(appTwoWidgetsListBaseEntries())
+                    },
+                defaultWidgets = buildList { addAll(appTwoWidgetsListBaseEntries()) }
+            )
+        val app1PackageUserKey = PackageUserKey.fromPackageItemInfo(app1PackageItemInfo)
+
+        val contentEntry = findContentEntryForPackageUser(widgetPickerData, app1PackageUserKey)
+
+        assertThat(contentEntry).isNotNull()
+        assertThat(contentEntry?.mPkgItem).isEqualTo(app1PackageItemInfo)
+        assertThat(contentEntry?.mWidgets).hasSize(2)
+    }
+
+    @Test
+    fun findContentEntryForPackageUser_fromDefaults_returnsEntryFromDefaultWidgets() {
+        val widgetPickerData =
+            WidgetPickerData(
+                allWidgets =
+                    buildList {
+                        addAll(appOneWidgetsListBaseEntries())
+                        addAll(appTwoWidgetsListBaseEntries())
+                    },
+                defaultWidgets =
+                    buildList { addAll(appOneWidgetsListBaseEntries(includeWidgetTwo = false)) }
+            )
+        val app1PackageUserKey = PackageUserKey.fromPackageItemInfo(app1PackageItemInfo)
+
+        val contentEntry =
+            findContentEntryForPackageUser(
+                widgetPickerData = widgetPickerData,
+                packageUserKey = app1PackageUserKey,
+                fromDefaultWidgets = true
+            )
+
+        assertThat(contentEntry).isNotNull()
+        assertThat(contentEntry?.mPkgItem).isEqualTo(app1PackageItemInfo)
+        // only one widget (since default widgets had only one widget for app A
+        assertThat(contentEntry?.mWidgets).hasSize(1)
+    }
+
+    @Test
+    fun findContentEntryForPackageUser_noMatch_returnsNull() {
+        val app2PackageUserKey = PackageUserKey.fromPackageItemInfo(app2PackageItemInfo)
+        val widgetPickerData =
+            WidgetPickerData(allWidgets = buildList { addAll(appOneWidgetsListBaseEntries()) })
+
+        val contentEntry = findContentEntryForPackageUser(widgetPickerData, app2PackageUserKey)
+
+        assertThat(contentEntry).isNull()
+    }
+
+    @Test
+    fun findAllWidgetsForPackageUser_returnsListOfWidgets() {
+        val app1PackageUserKey = PackageUserKey.fromPackageItemInfo(app1PackageItemInfo)
+        val widgetPickerData =
+            WidgetPickerData(
+                allWidgets =
+                    buildList {
+                        addAll(appOneWidgetsListBaseEntries())
+                        addAll(appTwoWidgetsListBaseEntries())
+                    },
+                defaultWidgets =
+                    buildList { addAll(appOneWidgetsListBaseEntries(includeWidgetTwo = false)) }
+            )
+
+        val widgets = findAllWidgetsForPackageUser(widgetPickerData, app1PackageUserKey)
+
+        // both widgets returned irrespective of default widgets list
+        assertThat(widgets).hasSize(2)
+    }
+
+    @Test
+    fun findAllWidgetsForPackageUser_noMatch_returnsEmptyList() {
+        val widgetPickerData =
+            WidgetPickerData(
+                allWidgets = buildList { addAll(appTwoWidgetsListBaseEntries()) },
+            )
+        val app1PackageUserKey = PackageUserKey.fromPackageItemInfo(app1PackageItemInfo)
+
+        val widgets = findAllWidgetsForPackageUser(widgetPickerData, app1PackageUserKey)
+
+        assertThat(widgets).isEmpty()
+    }
+
+    private fun packageItemInfoWithTitle(packageName: String, title: String): PackageItemInfo {
+        val packageItemInfo = PackageItemInfo(packageName, userHandle)
+        packageItemInfo.title = title
+        return packageItemInfo
+    }
+
+    private fun createWidgetItem(packageName: String, widgetProviderName: String): WidgetItem {
+        val providerInfo =
+            WidgetUtils.createAppWidgetProviderInfo(
+                ComponentName.createRelative(packageName, widgetProviderName)
+            )
+        val widgetInfo = LauncherAppWidgetProviderInfo.fromProviderInfo(context, providerInfo)
+        return WidgetItem(widgetInfo, testInvariantProfile, iconCache, context)
+    }
+
+    private fun appTwoWidgetsListBaseEntries(): List<WidgetsListBaseEntry> = buildList {
+        val widgets = listOf(app2WidgetItem1)
+        add(WidgetsListHeaderEntry.create(app2PackageItemInfo, APP_2_SECTION_NAME, widgets))
+        add(WidgetsListContentEntry(app2PackageItemInfo, APP_2_SECTION_NAME, widgets))
+    }
+
+    private fun appOneWidgetsListBaseEntries(
+        includeWidgetTwo: Boolean = true
+    ): List<WidgetsListBaseEntry> = buildList {
+        val widgets =
+            if (includeWidgetTwo) {
+                listOf(app1WidgetItem1, app1WidgetItem2)
+            } else {
+                listOf(app1WidgetItem1)
+            }
+
+        add(WidgetsListHeaderEntry.create(app1PackageItemInfo, APP_1_SECTION_NAME, widgets))
+        add(WidgetsListContentEntry(app1PackageItemInfo, APP_1_SECTION_NAME, widgets))
+    }
+
+    companion object {
+        private const val APP_1_PACKAGE_NAME = "com.example.app1"
+        private const val APP_1_PACKAGE_TITLE = "App1"
+        private const val APP_1_SECTION_NAME = "A" // for fast popup
+        private const val APP_1_PROVIDER_1_CLASS_NAME = "app1Provider1"
+        private const val APP_1_PROVIDER_2_CLASS_NAME = "app1Provider2"
+
+        private const val APP_2_PACKAGE_NAME = "com.example.app2"
+        private const val APP_2_PACKAGE_TITLE = "SomeApp2"
+        private const val APP_2_SECTION_NAME = "S" // for fast popup
+        private const val APP_2_PROVIDER_1_CLASS_NAME = "app2Provider1"
+
+        private val CATEGORY_1 =
+            WidgetRecommendationCategory(/* categoryTitleRes= */ 0, /* order= */ 0)
+        private val CATEGORY_2 =
+            WidgetRecommendationCategory(/* categoryTitleRes= */ 1, /* order= */ 1)
+    }
+}