Recreate hotseat predictor whenever we query it due to workspace change

Fix: b/289013842
Test: unit test, also verified moving icons will recreate hotseat predictor
Change-Id: I1f19b17654b87156132a4e4dee26e12312589dba
diff --git a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
index bc97df1..32361a8 100644
--- a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
+++ b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
@@ -52,6 +52,7 @@
 import androidx.annotation.CallSuper;
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
 import androidx.annotation.WorkerThread;
 
 import com.android.launcher3.InvariantDeviceProfile;
@@ -93,11 +94,14 @@
     private static final boolean IS_DEBUG = false;
     private static final String TAG = "QuickstepModelDelegate";
 
-    private final PredictorState mAllAppsState =
+    @VisibleForTesting
+    final PredictorState mAllAppsState =
             new PredictorState(CONTAINER_PREDICTION, "all_apps_predictions");
-    private final PredictorState mHotseatState =
+    @VisibleForTesting
+    final PredictorState mHotseatState =
             new PredictorState(CONTAINER_HOTSEAT_PREDICTION, "hotseat_predictions");
-    private final PredictorState mWidgetsRecommendationState =
+    @VisibleForTesting
+    final PredictorState mWidgetsRecommendationState =
             new PredictorState(CONTAINER_WIDGETS_PREDICTION, "widgets_prediction");
 
     private final InvariantDeviceProfile mIDP;
@@ -348,12 +352,7 @@
                         .build()));
 
         // TODO: get bundle
-        registerPredictor(mHotseatState, apm.createAppPredictionSession(
-                new AppPredictionContext.Builder(context)
-                        .setUiSurface("hotseat")
-                        .setPredictedTargetCount(mIDP.numDatabaseHotseatIcons)
-                        .setExtras(convertDataModelToAppTargetBundle(context, mDataModel))
-                        .build()));
+        registerHotseatPredictor(apm, context);
 
         registerWidgetsPredictor(apm.createAppPredictionSession(
                 new AppPredictionContext.Builder(context)
@@ -363,6 +362,29 @@
                         .build()));
     }
 
+    @WorkerThread
+    private void recreateHotseatPredictor() {
+        mHotseatState.destroyPredictor();
+        if (!mActive) {
+            return;
+        }
+        Context context = mApp.getContext();
+        AppPredictionManager apm = context.getSystemService(AppPredictionManager.class);
+        if (apm == null) {
+            return;
+        }
+        registerHotseatPredictor(apm, context);
+    }
+
+    private void registerHotseatPredictor(AppPredictionManager apm, Context context) {
+        registerPredictor(mHotseatState, apm.createAppPredictionSession(
+                new AppPredictionContext.Builder(context)
+                        .setUiSurface("hotseat")
+                        .setPredictedTargetCount(mIDP.numDatabaseHotseatIcons)
+                        .setExtras(convertDataModelToAppTargetBundle(context, mDataModel))
+                        .build()));
+    }
+
     private void registerPredictor(PredictorState state, AppPredictor predictor) {
         state.setTargets(Collections.emptyList());
         state.predictor = predictor;
@@ -393,7 +415,8 @@
         mWidgetsRecommendationState.predictor.requestPredictionUpdate();
     }
 
-    private void onAppTargetEvent(AppTargetEvent event, int client) {
+    @VisibleForTesting
+    void onAppTargetEvent(AppTargetEvent event, int client) {
         PredictorState state;
         switch(client) {
             case CONTAINER_PREDICTION:
@@ -411,6 +434,13 @@
             state.predictor.notifyAppTargetEvent(event);
             Log.d(TAG, "notifyAppTargetEvent action=" + event.getAction()
                     + " launchLocation=" + event.getLaunchLocation());
+            if (state == mHotseatState
+                    && (event.getAction() == AppTargetEvent.ACTION_PIN
+                            || event.getAction() == AppTargetEvent.ACTION_UNPIN)) {
+                // Recreate hot seat predictor when we need to query for hot seat due to pin or
+                // unpin app icons.
+                recreateHotseatPredictor();
+            }
         }
     }
 
diff --git a/quickstep/tests/src/com/android/launcher3/model/QuickstepModelDelegateTest.kt b/quickstep/tests/src/com/android/launcher3/model/QuickstepModelDelegateTest.kt
new file mode 100644
index 0000000..a532762
--- /dev/null
+++ b/quickstep/tests/src/com/android/launcher3/model/QuickstepModelDelegateTest.kt
@@ -0,0 +1,127 @@
+package com.android.launcher3.model
+
+import android.app.prediction.AppPredictor
+import android.app.prediction.AppTarget
+import android.app.prediction.AppTargetEvent
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.launcher3.LauncherAppState
+import com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT_PREDICTION
+import com.android.launcher3.LauncherSettings.Favorites.CONTAINER_PREDICTION
+import com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WALLPAPERS
+import com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION
+import com.android.launcher3.util.LauncherModelHelper
+import org.junit.After
+import org.junit.Assert.assertNotSame
+import org.junit.Assert.assertSame
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.Mockito.never
+import org.mockito.Mockito.verify
+import org.mockito.Mockito.verifyZeroInteractions
+import org.mockito.MockitoAnnotations
+
+/** Unit tests for [QuickstepModelDelegate]. */
+@RunWith(AndroidJUnit4::class)
+class QuickstepModelDelegateTest {
+
+    private lateinit var underTest: QuickstepModelDelegate
+    private lateinit var modelHelper: LauncherModelHelper
+
+    @Mock private lateinit var target: AppTarget
+    @Mock private lateinit var mockedAppTargetEvent: AppTargetEvent
+    @Mock private lateinit var allAppsPredictor: AppPredictor
+    @Mock private lateinit var hotseatPredictor: AppPredictor
+    @Mock private lateinit var widgetRecommendationPredictor: AppPredictor
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+        modelHelper = LauncherModelHelper()
+        underTest = QuickstepModelDelegate(modelHelper.sandboxContext)
+        underTest.mAllAppsState.predictor = allAppsPredictor
+        underTest.mHotseatState.predictor = hotseatPredictor
+        underTest.mWidgetsRecommendationState.predictor = widgetRecommendationPredictor
+        underTest.mApp = LauncherAppState.getInstance(modelHelper.sandboxContext)
+        underTest.mDataModel = BgDataModel()
+    }
+
+    @After
+    fun tearDown() {
+        modelHelper.destroy()
+    }
+
+    @Test
+    fun onAppTargetEvent_notifyTarget() {
+        underTest.onAppTargetEvent(mockedAppTargetEvent, CONTAINER_PREDICTION)
+
+        verify(allAppsPredictor).notifyAppTargetEvent(mockedAppTargetEvent)
+        verifyZeroInteractions(hotseatPredictor)
+        verifyZeroInteractions(widgetRecommendationPredictor)
+    }
+
+    @Test
+    fun onWidgetPrediction_notifyWidgetRecommendationPredictor() {
+        underTest.onAppTargetEvent(mockedAppTargetEvent, CONTAINER_WIDGETS_PREDICTION)
+
+        verifyZeroInteractions(allAppsPredictor)
+        verify(widgetRecommendationPredictor).notifyAppTargetEvent(mockedAppTargetEvent)
+        verifyZeroInteractions(hotseatPredictor)
+    }
+
+    @Test
+    fun onHotseatPrediction_notifyHotseatPredictor() {
+        underTest.onAppTargetEvent(mockedAppTargetEvent, CONTAINER_HOTSEAT_PREDICTION)
+
+        verifyZeroInteractions(allAppsPredictor)
+        verifyZeroInteractions(widgetRecommendationPredictor)
+        verify(hotseatPredictor).notifyAppTargetEvent(mockedAppTargetEvent)
+    }
+
+    @Test
+    fun onOtherClient_notifyHotseatPredictor() {
+        underTest.onAppTargetEvent(mockedAppTargetEvent, CONTAINER_WALLPAPERS)
+
+        verifyZeroInteractions(allAppsPredictor)
+        verifyZeroInteractions(widgetRecommendationPredictor)
+        verify(hotseatPredictor).notifyAppTargetEvent(mockedAppTargetEvent)
+    }
+
+    @Test
+    fun hotseatActionPin_recreateHotSeat() {
+        assertSame(underTest.mHotseatState.predictor, hotseatPredictor)
+        val appTargetEvent = AppTargetEvent.Builder(target, AppTargetEvent.ACTION_PIN).build()
+        underTest.markActive()
+
+        underTest.onAppTargetEvent(appTargetEvent, CONTAINER_HOTSEAT_PREDICTION)
+
+        verify(hotseatPredictor).destroy()
+        assertNotSame(underTest.mHotseatState.predictor, hotseatPredictor)
+    }
+
+    @Test
+    fun hotseatActionUnpin_recreateHotSeat() {
+        assertSame(underTest.mHotseatState.predictor, hotseatPredictor)
+        underTest.markActive()
+        val appTargetEvent = AppTargetEvent.Builder(target, AppTargetEvent.ACTION_UNPIN).build()
+
+        underTest.onAppTargetEvent(appTargetEvent, CONTAINER_HOTSEAT_PREDICTION)
+
+        verify(hotseatPredictor).destroy()
+        assertNotSame(underTest.mHotseatState.predictor, hotseatPredictor)
+    }
+
+    @Test
+    fun container_actionPin_notRecreateHotSeat() {
+        assertSame(underTest.mHotseatState.predictor, hotseatPredictor)
+        val appTargetEvent = AppTargetEvent.Builder(target, AppTargetEvent.ACTION_UNPIN).build()
+        underTest.markActive()
+
+        underTest.onAppTargetEvent(appTargetEvent, CONTAINER_PREDICTION)
+
+        verify(allAppsPredictor, never()).destroy()
+        verify(hotseatPredictor, never()).destroy()
+        assertSame(underTest.mHotseatState.predictor, hotseatPredictor)
+    }
+}