diff --git a/packages/SystemUI/src/com/android/systemui/wmshell/WMShell.java b/packages/SystemUI/src/com/android/systemui/wmshell/WMShell.java
index e1e18a22af51b..27af5f91c9e47 100644
--- a/packages/SystemUI/src/com/android/systemui/wmshell/WMShell.java
+++ b/packages/SystemUI/src/com/android/systemui/wmshell/WMShell.java
@@ -21,6 +21,7 @@ import static com.android.systemui.shared.system.WindowManagerWrapper.WINDOWING_
 import android.app.ActivityManager;
 import android.content.Context;
 
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.keyguard.KeyguardUpdateMonitor;
 import com.android.keyguard.KeyguardUpdateMonitorCallback;
 import com.android.systemui.SystemUI;
@@ -48,11 +49,11 @@ import javax.inject.Inject;
  * Proxy in SysUiScope to delegate events to controllers in WM Shell library.
  */
 @SysUISingleton
-public final class WMShell extends SystemUI implements ProtoTraceable<SystemUiTraceProto>,
-        CommandQueue.Callbacks {
+public final class WMShell extends SystemUI implements ProtoTraceable<SystemUiTraceProto> {
     private final CommandQueue mCommandQueue;
     private final DisplayImeController mDisplayImeController;
     private final KeyguardUpdateMonitor mKeyguardUpdateMonitor;
+    private final ActivityManagerWrapper mActivityManagerWrapper;
     private final Optional<Pip> mPipOptional;
     private final Optional<SplitScreen> mSplitScreenOptional;
     private final ProtoTracer mProtoTracer;
@@ -60,6 +61,7 @@ public final class WMShell extends SystemUI implements ProtoTraceable<SystemUiTr
     @Inject
     public WMShell(Context context, CommandQueue commandQueue,
             KeyguardUpdateMonitor keyguardUpdateMonitor,
+            ActivityManagerWrapper activityManagerWrapper,
             DisplayImeController displayImeController,
             Optional<Pip> pipOptional,
             Optional<SplitScreen> splitScreenOptional,
@@ -67,6 +69,7 @@ public final class WMShell extends SystemUI implements ProtoTraceable<SystemUiTr
         super(context);
         mCommandQueue = commandQueue;
         mKeyguardUpdateMonitor = keyguardUpdateMonitor;
+        mActivityManagerWrapper = activityManagerWrapper;
         mDisplayImeController = displayImeController;
         mPipOptional = pipOptional;
         mSplitScreenOptional = splitScreenOptional;
@@ -79,14 +82,14 @@ public final class WMShell extends SystemUI implements ProtoTraceable<SystemUiTr
         // This is to prevent circular init problem by separating registration step out of its
         // constructor. And make sure the initialization of DisplayImeController won't depend on
         // specific feature anymore.
-        mCommandQueue.addCallback(this);
         mDisplayImeController.startMonitorDisplays();
 
         mPipOptional.ifPresent(this::initPip);
         mSplitScreenOptional.ifPresent(this::initSplitScreen);
     }
 
-    private void initPip(Pip pip) {
+    @VisibleForTesting
+    void initPip(Pip pip) {
         mCommandQueue.addCallback(new CommandQueue.Callbacks() {
             @Override
             public void showPictureInPictureMenu() {
@@ -95,7 +98,8 @@ public final class WMShell extends SystemUI implements ProtoTraceable<SystemUiTr
         });
     }
 
-    private void initSplitScreen(SplitScreen splitScreen) {
+    @VisibleForTesting
+    void initSplitScreen(SplitScreen splitScreen) {
         mKeyguardUpdateMonitor.registerCallback(new KeyguardUpdateMonitorCallback() {
             @Override
             public void onKeyguardVisibilityChanged(boolean showing) {
@@ -107,7 +111,7 @@ public final class WMShell extends SystemUI implements ProtoTraceable<SystemUiTr
             }
         });
 
-        ActivityManagerWrapper.getInstance().registerTaskStackListener(
+        mActivityManagerWrapper.registerTaskStackListener(
                 new TaskStackChangeListener() {
                     @Override
                     public void onActivityRestartAttempt(ActivityManager.RunningTaskInfo task,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/wmshell/WMShellTest.java b/packages/SystemUI/tests/src/com/android/systemui/wmshell/WMShellTest.java
index 2f3c86dfc8ae5..ac9346f675be7 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/wmshell/WMShellTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/wmshell/WMShellTest.java
@@ -16,21 +16,19 @@
 
 package com.android.systemui.wmshell;
 
-import static org.mockito.Mockito.times;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
 
-import android.app.Instrumentation;
-import android.content.Context;
 import android.test.suitebuilder.annotation.SmallTest;
 
-import androidx.test.InstrumentationRegistry;
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.keyguard.KeyguardUpdateMonitor;
+import com.android.keyguard.KeyguardUpdateMonitorCallback;
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.pip.Pip;
-import com.android.systemui.pip.tv.PipController;
+import com.android.systemui.shared.system.ActivityManagerWrapper;
+import com.android.systemui.shared.system.TaskStackChangeListener;
 import com.android.systemui.stackdivider.SplitScreen;
 import com.android.systemui.statusbar.CommandQueue;
 import com.android.systemui.tracing.ProtoTracer;
@@ -48,29 +46,43 @@ import java.util.Optional;
 @RunWith(AndroidJUnit4.class)
 public class WMShellTest extends SysuiTestCase {
 
-    Instrumentation mInstrumentation;
     WMShell mWMShell;
-    @Mock Context mContext;
     @Mock CommandQueue mCommandQueue;
     @Mock KeyguardUpdateMonitor mKeyguardUpdateMonitor;
+    @Mock ActivityManagerWrapper mActivityManagerWrapper;
     @Mock DisplayImeController mDisplayImeController;
-    @Mock Optional<Pip> mPipOptional;
-    @Mock Optional<SplitScreen> mSplitScreenOptional;
-    @Mock PipController mPipController;
+    @Mock Pip mPip;
+    @Mock SplitScreen mSplitScreen;
     @Mock ProtoTracer mProtoTracer;
 
     @Before
-    public void setUp() throws Exception {
-        mInstrumentation = InstrumentationRegistry.getInstrumentation();
+    public void setUp() {
         MockitoAnnotations.initMocks(this);
         mWMShell = new WMShell(mContext, mCommandQueue, mKeyguardUpdateMonitor,
-                mDisplayImeController, mPipOptional, mSplitScreenOptional, mProtoTracer);
-        mWMShell.start();
-        when(mPipOptional.get()).thenReturn(mPipController);
+                mActivityManagerWrapper, mDisplayImeController, Optional.of(mPip),
+                Optional.of(mSplitScreen), mProtoTracer);
     }
 
     @Test
-    public void testWMShellRegisterCommandQueue() {
-        verify(mCommandQueue, times(1)).addCallback(mWMShell);
+    public void start_startsMonitorDisplays() {
+        mWMShell.start();
+
+        verify(mDisplayImeController).startMonitorDisplays();
+    }
+
+    @Test
+    public void initPip_registersCommandQueueCallback() {
+        mWMShell.initPip(mPip);
+
+        verify(mCommandQueue).addCallback(any(CommandQueue.Callbacks.class));
+    }
+
+    @Test
+    public void initSplitScreen_registersCallbacks() {
+        mWMShell.initSplitScreen(mSplitScreen);
+
+        verify(mKeyguardUpdateMonitor).registerCallback(any(KeyguardUpdateMonitorCallback.class));
+        verify(mActivityManagerWrapper).registerTaskStackListener(
+                any(TaskStackChangeListener.class));
     }
 }