Remove expensive queries for lock task mode (screen pinning)

- Instead use the TaskStackListener callbacks to update the state
  in each component as needed

Bug: 224544468
Test: Enable/disable lock task mode and
      dumpsys activity service SystemUIService to confirm the sysui state
      is still updated
Test: atest SystemUITests
Change-Id: I15ec6251c75a29c563820ea6424768fa33245497
diff --git a/packages/SystemUI/shared/src/com/android/systemui/shared/system/ActivityManagerWrapper.java b/packages/SystemUI/shared/src/com/android/systemui/shared/system/ActivityManagerWrapper.java
index fd41cb06..6bfaf5e 100644
--- a/packages/SystemUI/shared/src/com/android/systemui/shared/system/ActivityManagerWrapper.java
+++ b/packages/SystemUI/shared/src/com/android/systemui/shared/system/ActivityManagerWrapper.java
@@ -283,17 +283,6 @@
     }
 
     /**
-     * @return whether screen pinning is active.
-     */
-    public boolean isScreenPinningActive() {
-        try {
-            return getService().getLockTaskModeState() == LOCK_TASK_MODE_PINNED;
-        } catch (RemoteException e) {
-            return false;
-        }
-    }
-
-    /**
      * @return whether screen pinning is enabled.
      */
     public boolean isScreenPinningEnabled() {
diff --git a/packages/SystemUI/shared/src/com/android/systemui/shared/system/TaskStackChangeListeners.java b/packages/SystemUI/shared/src/com/android/systemui/shared/system/TaskStackChangeListeners.java
index 8af934f..dd52cfb 100644
--- a/packages/SystemUI/shared/src/com/android/systemui/shared/system/TaskStackChangeListeners.java
+++ b/packages/SystemUI/shared/src/com/android/systemui/shared/system/TaskStackChangeListeners.java
@@ -16,6 +16,7 @@
 
 package com.android.systemui.shared.system;
 
+import android.annotation.NonNull;
 import android.app.ActivityManager.RunningTaskInfo;
 import android.app.ActivityTaskManager;
 import android.app.TaskStackListener;
@@ -27,6 +28,8 @@
 import android.util.Log;
 import android.window.TaskSnapshot;
 
+import androidx.annotation.VisibleForTesting;
+
 import com.android.internal.os.SomeArgs;
 import com.android.systemui.shared.recents.model.ThumbnailData;
 
@@ -43,15 +46,51 @@
 
     private final Impl mImpl;
 
+    /**
+     * Proxies calls to the given handler callback synchronously for testing purposes.
+     */
+    private static class TestSyncHandler extends Handler {
+        private Handler.Callback mCb;
+
+        public TestSyncHandler() {
+            super(Looper.getMainLooper());
+        }
+
+        public void setCallback(Handler.Callback cb) {
+            mCb = cb;
+        }
+
+        @Override
+        public boolean sendMessageAtTime(@NonNull Message msg, long uptimeMillis) {
+            return mCb.handleMessage(msg);
+        }
+    }
+
     private TaskStackChangeListeners() {
         mImpl = new Impl(Looper.getMainLooper());
     }
 
+    private TaskStackChangeListeners(Handler h) {
+        mImpl = new Impl(h);
+    }
+
     public static TaskStackChangeListeners getInstance() {
         return INSTANCE;
     }
 
     /**
+     * Returns an instance of the listeners that can be called upon synchronously for testsing
+     * purposes.
+     */
+    @VisibleForTesting
+    public static TaskStackChangeListeners getTestInstance() {
+        TestSyncHandler h = new TestSyncHandler();
+        TaskStackChangeListeners l = new TaskStackChangeListeners(h);
+        h.setCallback(l.mImpl);
+        return l;
+    }
+
+    /**
      * Registers a task stack listener with the system.
      * This should be called on the main thread.
      */
@@ -71,7 +110,15 @@
         }
     }
 
-    private static class Impl extends TaskStackListener implements Handler.Callback {
+    /**
+     * Returns an instance of the listener to call upon from tests.
+     */
+    @VisibleForTesting
+    public TaskStackListener getListenerImpl() {
+        return mImpl;
+    }
+
+    private class Impl extends TaskStackListener implements Handler.Callback {
 
         private static final int ON_TASK_STACK_CHANGED = 1;
         private static final int ON_TASK_SNAPSHOT_CHANGED = 2;
@@ -104,10 +151,14 @@
         private final Handler mHandler;
         private boolean mRegistered;
 
-        Impl(Looper looper) {
+        private Impl(Looper looper) {
             mHandler = new Handler(looper, this);
         }
 
+        private Impl(Handler handler) {
+            mHandler = handler;
+        }
+
         public void addListener(TaskStackChangeListener listener) {
             synchronized (mTaskStackListeners) {
                 mTaskStackListeners.add(listener);
diff --git a/packages/SystemUI/src/com/android/systemui/navigationbar/NavigationBar.java b/packages/SystemUI/src/com/android/systemui/navigationbar/NavigationBar.java
index 5b306c9..2647600 100644
--- a/packages/SystemUI/src/com/android/systemui/navigationbar/NavigationBar.java
+++ b/packages/SystemUI/src/com/android/systemui/navigationbar/NavigationBar.java
@@ -16,6 +16,7 @@
 
 package com.android.systemui.navigationbar;
 
+import static android.app.ActivityManager.LOCK_TASK_MODE_PINNED;
 import static android.app.StatusBarManager.NAVIGATION_HINT_BACK_ALT;
 import static android.app.StatusBarManager.NAVIGATION_HINT_IME_SWITCHER_SHOWN;
 import static android.app.StatusBarManager.WINDOW_STATE_HIDDEN;
@@ -44,6 +45,7 @@
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SHOWING;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SWITCHER_SHOWING;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_NAV_BAR_HIDDEN;
+import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_SCREEN_PINNING;
 import static com.android.systemui.shared.system.QuickStepContract.isGesturalMode;
 import static com.android.systemui.statusbar.phone.BarTransitions.MODE_OPAQUE;
 import static com.android.systemui.statusbar.phone.BarTransitions.TransitionMode;
@@ -136,9 +138,10 @@
 import com.android.systemui.shared.recents.utilities.Utilities;
 import com.android.systemui.shared.rotation.RotationButton;
 import com.android.systemui.shared.rotation.RotationButtonController;
-import com.android.systemui.shared.system.ActivityManagerWrapper;
 import com.android.systemui.shared.system.QuickStepContract;
 import com.android.systemui.shared.system.SysUiStatsLog;
+import com.android.systemui.shared.system.TaskStackChangeListener;
+import com.android.systemui.shared.system.TaskStackChangeListeners;
 import com.android.systemui.statusbar.AutoHideUiElement;
 import com.android.systemui.statusbar.CommandQueue;
 import com.android.systemui.statusbar.CommandQueue.Callbacks;
@@ -252,6 +255,7 @@
     private final AutoHideController.Factory mAutoHideControllerFactory;
     private final Optional<TelecomManager> mTelecomManagerOptional;
     private final InputMethodManager mInputMethodManager;
+    private final TaskStackChangeListeners mTaskStackChangeListeners;
 
     @VisibleForTesting
     public int mDisplayId;
@@ -488,6 +492,18 @@
             }
     };
 
+    private boolean mScreenPinningActive = false;
+    private final TaskStackChangeListener mTaskStackListener = new TaskStackChangeListener() {
+        @Override
+        public void onLockTaskModeChanged(int mode) {
+            mScreenPinningActive = (mode == LOCK_TASK_MODE_PINNED);
+            mSysUiFlagsContainer.setFlag(SYSUI_STATE_SCREEN_PINNING, mScreenPinningActive)
+                    .commitUpdate(mDisplayId);
+            mView.setInScreenPinning(mScreenPinningActive);
+            updateScreenPinningGestures();
+        }
+    };
+
     @Inject
     NavigationBar(
             NavigationBarView navigationBarView,
@@ -529,7 +545,8 @@
             EdgeBackGestureHandler edgeBackGestureHandler,
             Optional<BackAnimation> backAnimation,
             UserContextProvider userContextProvider,
-            WakefulnessLifecycle wakefulnessLifecycle) {
+            WakefulnessLifecycle wakefulnessLifecycle,
+            TaskStackChangeListeners taskStackChangeListeners) {
         super(navigationBarView);
         mFrame = navigationBarFrame;
         mContext = context;
@@ -568,6 +585,7 @@
         mInputMethodManager = inputMethodManager;
         mUserContextProvider = userContextProvider;
         mWakefulnessLifecycle = wakefulnessLifecycle;
+        mTaskStackChangeListeners = taskStackChangeListeners;
 
         mNavColorSampleMargin = getResources()
                 .getDimensionPixelSize(R.dimen.navigation_handle_sample_horizontal_margin);
@@ -676,6 +694,7 @@
         mCommandQueue.recomputeDisableFlags(mDisplayId, false);
 
         mNotificationShadeDepthController.addListener(mDepthListener);
+        mTaskStackChangeListeners.registerTaskStackListener(mTaskStackListener);
     }
 
     public void destroyView() {
@@ -689,6 +708,7 @@
         mNotificationShadeDepthController.removeListener(mDepthListener);
 
         mDeviceConfigProxy.removeOnPropertiesChangedListener(mOnPropertiesChangedListener);
+        mTaskStackChangeListeners.unregisterTaskStackListener(mTaskStackListener);
     }
 
     @Override
@@ -990,6 +1010,7 @@
         pw.println("  mTransientShown=" + mTransientShown);
         pw.println("  mTransientShownFromGestureOnSystemBar="
                 + mTransientShownFromGestureOnSystemBar);
+        pw.println("  mScreenPinningActive=" + mScreenPinningActive);
         dumpBarTransitions(pw, "mNavigationBarView", getBarTransitions());
 
         pw.println("  mOrientedHandleSamplingRegion: " + mOrientedHandleSamplingRegion);
@@ -1213,10 +1234,9 @@
 
     private void updateScreenPinningGestures() {
         // Change the cancel pin gesture to home and back if recents button is invisible
-        boolean pinningActive = ActivityManagerWrapper.getInstance().isScreenPinningActive();
         ButtonDispatcher backButton = mView.getBackButton();
         ButtonDispatcher recentsButton = mView.getRecentsButton();
-        if (pinningActive) {
+        if (mScreenPinningActive) {
             boolean recentsVisible = mView.isRecentsButtonVisible();
             backButton.setOnLongClickListener(recentsVisible
                     ? this::onLongPressBackRecents
@@ -1227,8 +1247,8 @@
             recentsButton.setOnLongClickListener(null);
         }
         // Note, this needs to be set after even if we're setting the listener to null
-        backButton.setLongClickable(pinningActive);
-        recentsButton.setLongClickable(pinningActive);
+        backButton.setLongClickable(mScreenPinningActive);
+        recentsButton.setLongClickable(mScreenPinningActive);
     }
 
     private void notifyNavigationBarScreenOn() {
@@ -1311,8 +1331,7 @@
 
     @VisibleForTesting
     boolean onHomeLongClick(View v) {
-        if (!mView.isRecentsButtonVisible()
-                && ActivityManagerWrapper.getInstance().isScreenPinningActive()) {
+        if (!mView.isRecentsButtonVisible() && mScreenPinningActive) {
             return onLongPressBackHome(v);
         }
         if (shouldDisableNavbarGestures()) {
diff --git a/packages/SystemUI/src/com/android/systemui/navigationbar/NavigationBarController.java b/packages/SystemUI/src/com/android/systemui/navigationbar/NavigationBarController.java
index 8914552..5d43c5d 100644
--- a/packages/SystemUI/src/com/android/systemui/navigationbar/NavigationBarController.java
+++ b/packages/SystemUI/src/com/android/systemui/navigationbar/NavigationBarController.java
@@ -57,6 +57,7 @@
 import com.android.systemui.model.SysUiState;
 import com.android.systemui.recents.OverviewProxyService;
 import com.android.systemui.shared.system.QuickStepContract;
+import com.android.systemui.shared.system.TaskStackChangeListeners;
 import com.android.systemui.statusbar.CommandQueue;
 import com.android.systemui.statusbar.CommandQueue.Callbacks;
 import com.android.systemui.statusbar.phone.AutoHideController;
@@ -88,7 +89,6 @@
     private FeatureFlags mFeatureFlags;
     private final DisplayManager mDisplayManager;
     private final TaskbarDelegate mTaskbarDelegate;
-    private final StatusBarKeyguardViewManager mStatusBarKeyguardViewManager;
     private int mNavMode;
     @VisibleForTesting boolean mIsTablet;
 
@@ -112,10 +112,10 @@
             NavBarHelper navBarHelper,
             TaskbarDelegate taskbarDelegate,
             NavigationBarComponent.Factory navigationBarComponentFactory,
-            StatusBarKeyguardViewManager statusBarKeyguardViewManager,
             DumpManager dumpManager,
             AutoHideController autoHideController,
             LightBarController lightBarController,
+            TaskStackChangeListeners taskStackChangeListeners,
             Optional<Pip> pipOptional,
             Optional<BackAnimation> backAnimation,
             FeatureFlags featureFlags) {
@@ -129,11 +129,10 @@
         mConfigChanges.applyNewConfig(mContext.getResources());
         mNavMode = navigationModeController.addListener(this);
         mTaskbarDelegate = taskbarDelegate;
-        mStatusBarKeyguardViewManager = statusBarKeyguardViewManager;
         mTaskbarDelegate.setDependencies(commandQueue, overviewProxyService,
                 navBarHelper, navigationModeController, sysUiFlagsContainer,
                 dumpManager, autoHideController, lightBarController, pipOptional,
-                backAnimation.orElse(null));
+                backAnimation.orElse(null), taskStackChangeListeners);
         mIsTablet = isTablet(mContext);
         dumpManager.registerDumpable(this);
     }
diff --git a/packages/SystemUI/src/com/android/systemui/navigationbar/NavigationBarView.java b/packages/SystemUI/src/com/android/systemui/navigationbar/NavigationBarView.java
index 403d276..88c4fd5 100644
--- a/packages/SystemUI/src/com/android/systemui/navigationbar/NavigationBarView.java
+++ b/packages/SystemUI/src/com/android/systemui/navigationbar/NavigationBarView.java
@@ -16,6 +16,7 @@
 
 package com.android.systemui.navigationbar;
 
+import static android.app.ActivityManager.LOCK_TASK_MODE_PINNED;
 import static android.inputmethodservice.InputMethodService.canImeRenderGesturalNavButtons;
 import static android.view.WindowManagerPolicyConstants.NAV_BAR_MODE_GESTURAL;
 
@@ -80,6 +81,7 @@
 import com.android.systemui.shared.rotation.RotationButtonController;
 import com.android.systemui.shared.system.ActivityManagerWrapper;
 import com.android.systemui.shared.system.QuickStepContract;
+import com.android.systemui.shared.system.TaskStackChangeListener;
 import com.android.systemui.statusbar.phone.AutoHideController;
 import com.android.systemui.statusbar.phone.CentralSurfaces;
 import com.android.systemui.statusbar.phone.LightBarTransitionsController;
@@ -160,6 +162,7 @@
      * fully locked mode we only show that unlocking is blocked.
      */
     private ScreenPinningNotify mScreenPinningNotify;
+    private boolean mScreenPinningActive = false;
 
     /**
      * {@code true} if the IME can render the back button and the IME switcher button.
@@ -636,14 +639,13 @@
         // When screen pinning, don't hide back and home when connected service or back and
         // recents buttons when disconnected from launcher service in screen pinning mode,
         // as they are used for exiting.
-        final boolean pinningActive = ActivityManagerWrapper.getInstance().isScreenPinningActive();
         if (mOverviewProxyEnabled) {
             // Force disable recents when not in legacy mode
             disableRecent |= !QuickStepContract.isLegacyMode(mNavBarMode);
-            if (pinningActive && !QuickStepContract.isGesturalMode(mNavBarMode)) {
+            if (mScreenPinningActive && !QuickStepContract.isGesturalMode(mNavBarMode)) {
                 disableBack = disableHome = false;
             }
-        } else if (pinningActive) {
+        } else if (mScreenPinningActive) {
             disableBack = disableRecent = false;
         }
 
@@ -738,9 +740,7 @@
     public void updateDisabledSystemUiStateFlags(SysUiState sysUiState) {
         int displayId = mContext.getDisplayId();
 
-        sysUiState.setFlag(SYSUI_STATE_SCREEN_PINNING,
-                        ActivityManagerWrapper.getInstance().isScreenPinningActive())
-                .setFlag(SYSUI_STATE_OVERVIEW_DISABLED,
+        sysUiState.setFlag(SYSUI_STATE_OVERVIEW_DISABLED,
                         (mDisabledFlags & View.STATUS_BAR_DISABLE_RECENT) != 0)
                 .setFlag(SYSUI_STATE_HOME_DISABLED,
                         (mDisabledFlags & View.STATUS_BAR_DISABLE_HOME) != 0)
@@ -749,6 +749,10 @@
                 .commitUpdate(displayId);
     }
 
+    public void setInScreenPinning(boolean active) {
+        mScreenPinningActive = active;
+    }
+
     private void updatePanelSystemUiStateFlags() {
         if (SysUiState.DEBUG) {
             Log.d(TAG, "Updating panel sysui state flags: panelView=" + mPanelView);
diff --git a/packages/SystemUI/src/com/android/systemui/navigationbar/TaskbarDelegate.java b/packages/SystemUI/src/com/android/systemui/navigationbar/TaskbarDelegate.java
index 5e26e60..6ee86aa 100644
--- a/packages/SystemUI/src/com/android/systemui/navigationbar/TaskbarDelegate.java
+++ b/packages/SystemUI/src/com/android/systemui/navigationbar/TaskbarDelegate.java
@@ -16,6 +16,7 @@
 
 package com.android.systemui.navigationbar;
 
+import static android.app.ActivityManager.LOCK_TASK_MODE_PINNED;
 import static android.app.StatusBarManager.NAVIGATION_HINT_BACK_ALT;
 import static android.app.StatusBarManager.NAVIGATION_HINT_IME_SWITCHER_SHOWN;
 import static android.app.StatusBarManager.WINDOW_STATE_SHOWING;
@@ -40,6 +41,7 @@
 
 import android.app.StatusBarManager;
 import android.app.StatusBarManager.WindowVisibleState;
+import android.content.ComponentName;
 import android.content.Context;
 import android.content.res.Configuration;
 import android.graphics.Rect;
@@ -68,6 +70,8 @@
 import com.android.systemui.shared.recents.utilities.Utilities;
 import com.android.systemui.shared.system.ActivityManagerWrapper;
 import com.android.systemui.shared.system.QuickStepContract;
+import com.android.systemui.shared.system.TaskStackChangeListener;
+import com.android.systemui.shared.system.TaskStackChangeListeners;
 import com.android.systemui.statusbar.AutoHideUiElement;
 import com.android.systemui.statusbar.CommandQueue;
 import com.android.systemui.statusbar.phone.AutoHideController;
@@ -101,6 +105,7 @@
     private AutoHideController mAutoHideController;
     private LightBarController mLightBarController;
     private LightBarTransitionsController mLightBarTransitionsController;
+    private TaskStackChangeListeners mTaskStackChangeListeners;
     private Optional<Pip> mPipOptional;
     private int mDisplayId;
     private int mNavigationIconHints;
@@ -127,6 +132,14 @@
     private final DisplayManager mDisplayManager;
     private Context mWindowContext;
     private ScreenPinningNotify mScreenPinningNotify;
+    private final TaskStackChangeListener mTaskStackListener = new TaskStackChangeListener() {
+        @Override
+        public void onLockTaskModeChanged(int mode) {
+            mSysUiState.setFlag(SYSUI_STATE_SCREEN_PINNING, mode == LOCK_TASK_MODE_PINNED)
+                    .commitUpdate(mDisplayId);
+        }
+    };
+
     private int mNavigationMode = -1;
     private final Consumer<Rect> mPipListener;
 
@@ -176,7 +189,8 @@
             AutoHideController autoHideController,
             LightBarController lightBarController,
             Optional<Pip> pipOptional,
-            BackAnimation backAnimation) {
+            BackAnimation backAnimation,
+            TaskStackChangeListeners taskStackChangeListeners) {
         // TODO: adding this in the ctor results in a dagger dependency cycle :(
         mCommandQueue = commandQueue;
         mOverviewProxyService = overviewProxyService;
@@ -189,6 +203,7 @@
         mPipOptional = pipOptional;
         mBackAnimation = backAnimation;
         mLightBarTransitionsController = createLightBarTransitionsController();
+        mTaskStackChangeListeners = taskStackChangeListeners;
     }
 
     // Separated into a method to keep setDependencies() clean/readable.
@@ -234,6 +249,7 @@
         mPipOptional.ifPresent(this::addPipExclusionBoundsChangeListener);
         mEdgeBackGestureHandler.setBackAnimation(mBackAnimation);
         mEdgeBackGestureHandler.onConfigurationChanged(mContext.getResources().getConfiguration());
+        mTaskStackChangeListeners.registerTaskStackListener(mTaskStackListener);
         mInitialized = true;
     }
 
@@ -253,6 +269,7 @@
         mLightBarTransitionsController.destroy();
         mLightBarController.setNavigationBar(null);
         mPipOptional.ifPresent(this::removePipExclusionBoundsChangeListener);
+        mTaskStackChangeListeners.unregisterTaskStackListener(mTaskStackListener);
         mInitialized = false;
     }
 
@@ -300,8 +317,6 @@
                 .setFlag(SYSUI_STATE_NAV_BAR_HIDDEN, !isWindowVisible())
                 .setFlag(SYSUI_STATE_ALLOW_GESTURE_IGNORING_BAR_VISIBILITY,
                         allowSystemGestureIgnoringBarVisibility())
-                .setFlag(SYSUI_STATE_SCREEN_PINNING,
-                        ActivityManagerWrapper.getInstance().isScreenPinningActive())
                 .setFlag(SYSUI_STATE_IMMERSIVE_MODE, isImmersiveMode())
                 .commitUpdate(mDisplayId);
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/navigationbar/NavigationBarControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/navigationbar/NavigationBarControllerTest.java
index 9bf27a2..8b0342e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/navigationbar/NavigationBarControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/navigationbar/NavigationBarControllerTest.java
@@ -51,6 +51,7 @@
 import com.android.systemui.model.SysUiState;
 import com.android.systemui.recents.OverviewProxyService;
 import com.android.systemui.shared.recents.utilities.Utilities;
+import com.android.systemui.shared.system.TaskStackChangeListeners;
 import com.android.systemui.statusbar.CommandQueue;
 import com.android.systemui.statusbar.phone.AutoHideController;
 import com.android.systemui.statusbar.phone.LightBarController;
@@ -102,10 +103,10 @@
                         mock(NavBarHelper.class),
                         mTaskbarDelegate,
                         mNavigationBarFactory,
-                        mock(StatusBarKeyguardViewManager.class),
                         mock(DumpManager.class),
                         mock(AutoHideController.class),
                         mock(LightBarController.class),
+                        TaskStackChangeListeners.getTestInstance(),
                         Optional.of(mock(Pip.class)),
                         Optional.of(mock(BackAnimation.class)),
                         mock(FeatureFlags.class)));
diff --git a/packages/SystemUI/tests/src/com/android/systemui/navigationbar/NavigationBarTest.java b/packages/SystemUI/tests/src/com/android/systemui/navigationbar/NavigationBarTest.java
index 80adbf0..2ad865e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/navigationbar/NavigationBarTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/navigationbar/NavigationBarTest.java
@@ -28,6 +28,7 @@
 
 import static com.android.internal.config.sysui.SystemUiDeviceConfigFlags.HOME_BUTTON_LONG_PRESS_DURATION_MS;
 import static com.android.systemui.navigationbar.NavigationBar.NavBarActionEvent.NAVBAR_ASSIST_LONGPRESS;
+import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_SCREEN_PINNING;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -44,6 +45,7 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.app.ActivityManager;
 import android.content.Context;
 import android.content.res.Resources;
 import android.hardware.display.DisplayManagerGlobal;
@@ -90,6 +92,7 @@
 import com.android.systemui.shade.NotificationShadeWindowView;
 import com.android.systemui.shade.ShadeController;
 import com.android.systemui.shared.rotation.RotationButtonController;
+import com.android.systemui.shared.system.TaskStackChangeListeners;
 import com.android.systemui.statusbar.CommandQueue;
 import com.android.systemui.statusbar.NotificationRemoteInputManager;
 import com.android.systemui.statusbar.NotificationShadeDepthController;
@@ -203,6 +206,8 @@
     private ViewRootImpl mViewRootImpl;
     private FakeExecutor mFakeExecutor = new FakeExecutor(new FakeSystemClock());
     private DeviceConfigProxyFake mDeviceConfigProxyFake = new DeviceConfigProxyFake();
+    private TaskStackChangeListeners mTaskStackChangeListeners =
+            TaskStackChangeListeners.getTestInstance();
 
     @Rule
     public final LeakCheckedTest.SysuiLeakCheck mLeakCheck = new LeakCheckedTest.SysuiLeakCheck();
@@ -437,6 +442,14 @@
         verify(mNavBarHelper, times(1)).getCurrentSysuiState();
     }
 
+    @Test
+    public void testScreenPinningEnabled_updatesSysuiState() {
+        mNavigationBar.init();
+        mTaskStackChangeListeners.getListenerImpl().onLockTaskModeChanged(
+                ActivityManager.LOCK_TASK_MODE_PINNED);
+        verify(mMockSysUiState).setFlag(eq(SYSUI_STATE_SCREEN_PINNING), eq(true));
+    }
+
     private NavigationBar createNavBar(Context context) {
         DeviceProvisionedController deviceProvisionedController =
                 mock(DeviceProvisionedController.class);
@@ -481,7 +494,8 @@
                 mEdgeBackGestureHandler,
                 Optional.of(mock(BackAnimation.class)),
                 mUserContextProvider,
-                mWakefulnessLifecycle));
+                mWakefulnessLifecycle,
+                mTaskStackChangeListeners));
     }
 
     private void processAllMessages() {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/navigationbar/TaskbarDelegateTest.kt b/packages/SystemUI/tests/src/com/android/systemui/navigationbar/TaskbarDelegateTest.kt
index 1742c69..537dfb8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/navigationbar/TaskbarDelegateTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/navigationbar/TaskbarDelegateTest.kt
@@ -1,11 +1,14 @@
 package com.android.systemui.navigationbar
 
+import android.app.ActivityManager
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.dump.DumpManager
 import com.android.systemui.model.SysUiState
 import com.android.systemui.navigationbar.gestural.EdgeBackGestureHandler
 import com.android.systemui.recents.OverviewProxyService
+import com.android.systemui.shared.system.QuickStepContract
+import com.android.systemui.shared.system.TaskStackChangeListeners
 import com.android.systemui.statusbar.CommandQueue
 import com.android.systemui.statusbar.phone.AutoHideController
 import com.android.systemui.statusbar.phone.LightBarController
@@ -14,6 +17,7 @@
 import com.android.wm.shell.pip.Pip
 import org.junit.Before
 import org.junit.Test
+import org.mockito.ArgumentMatchers
 import org.mockito.Mock
 import org.mockito.Mockito.`when`
 import org.mockito.Mockito.any
@@ -30,6 +34,7 @@
     val MODE_GESTURE = 0;
     val MODE_THREE_BUTTON = 1;
 
+    private lateinit var mTaskStackChangeListeners: TaskStackChangeListeners
     private lateinit var mTaskbarDelegate: TaskbarDelegate
     @Mock
     lateinit var mEdgeBackGestureHandlerFactory : EdgeBackGestureHandler.Factory
@@ -69,11 +74,12 @@
         `when`(mLightBarControllerFactory.create(any())).thenReturn(mLightBarTransitionController)
         `when`(mNavBarHelper.currentSysuiState).thenReturn(mCurrentSysUiState)
         `when`(mSysUiState.setFlag(anyInt(), anyBoolean())).thenReturn(mSysUiState)
+        mTaskStackChangeListeners = TaskStackChangeListeners.getTestInstance()
         mTaskbarDelegate = TaskbarDelegate(context, mEdgeBackGestureHandlerFactory,
                 mLightBarControllerFactory)
         mTaskbarDelegate.setDependencies(mCommandQueue, mOverviewProxyService, mNavBarHelper,
         mNavigationModeController, mSysUiState, mDumpManager, mAutoHideController,
-                mLightBarController, mOptionalPip, mBackAnimation)
+                mLightBarController, mOptionalPip, mBackAnimation, mTaskStackChangeListeners)
     }
 
     @Test
@@ -90,4 +96,15 @@
         mTaskbarDelegate.init(DISPLAY_ID)
         verify(mEdgeBackGestureHandler, times(1)).onNavigationModeChanged(MODE_GESTURE)
     }
+
+    @Test
+    fun screenPinningEnabled_updatesSysuiState() {
+        mTaskbarDelegate.init(DISPLAY_ID)
+        mTaskStackChangeListeners.listenerImpl.onLockTaskModeChanged(
+            ActivityManager.LOCK_TASK_MODE_PINNED)
+        verify(mSysUiState, times(1)).setFlag(
+            ArgumentMatchers.eq(QuickStepContract.SYSUI_STATE_SCREEN_PINNING),
+            ArgumentMatchers.eq(true)
+        )
+    }
 }
\ No newline at end of file