Make SysUI state instances display aware

With this cl each instance of SysUIState will be associated to a
specific display id.
The default instance provided in the dagger graph is for the default
display. In a follow up, support for multiple instance will be added,
and propagated to launcher.

In this cl DisplayTracker is removed, as the default display is always
Display.DEFAULT_DISPLAY (and will always be that). The old
"onSystemUiStateChanged" is only propagated for the default display also
 after this change.

Bug: 362719719
Bug: 398011576
Test: SysUiStateTest
Flag: com.android.systemui.shade_window_goes_around
Change-Id: Ie9b3b9f57848a92e25db389ee4cb18ae308d679e

10 files changed, 161 insertions, 85 deletions

diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/model/SysUiStateExtTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/model/SysUiStateExtTest.kt
index a3be9e35b912..09588f9f3751 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/model/SysUiStateExtTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/model/SysUiStateExtTest.kt
@@ -20,7 +20,6 @@ import android.view.Display
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.settings.FakeDisplayTracker
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
 import org.junit.Test
@@ -32,20 +31,11 @@ class SysUiStateExtTest : SysuiTestCase() {
 
     private val kosmos = testKosmos()
 
-    private val underTest =
-        SysUiState(
-            FakeDisplayTracker(context),
-            kosmos.sceneContainerPlugin,
-        )
+    private val underTest = kosmos.sysUiState
 
     @Test
     fun updateFlags() {
-        underTest.updateFlags(
-            Display.DEFAULT_DISPLAY,
-            1L to true,
-            2L to false,
-            3L to true,
-        )
+        underTest.updateFlags(Display.DEFAULT_DISPLAY, 1L to true, 2L to false, 3L to true)
 
         assertThat(underTest.flags and 1L).isNotEqualTo(0L)
         assertThat(underTest.flags and 2L).isEqualTo(0L)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/model/SysUiStateTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/model/SysUiStateTest.java
index 9a78bd93f424..eb4277a66be8 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/model/SysUiStateTest.java
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/model/SysUiStateTest.java
@@ -20,6 +20,8 @@ package com.android.systemui.model;
 import static android.view.Display.DEFAULT_DISPLAY;
 
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
@@ -28,7 +30,7 @@ import androidx.test.filters.SmallTest;
 
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.kosmos.KosmosJavaAdapter;
-import com.android.systemui.settings.FakeDisplayTracker;
+import com.android.systemui.settings.DisplayTracker;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -46,14 +48,23 @@ public class SysUiStateTest extends SysuiTestCase {
     private KosmosJavaAdapter mKosmos;
     private SysUiState.SysUiStateCallback mCallback;
     private SysUiState mFlagsContainer;
+    private DisplayTracker mDisplayTracker;
+    private SceneContainerPlugin mSceneContainerPlugin;
+
+    private SysUiState createInstance(int displayId) {
+        var sysuiState = new SysUiStateImpl(displayId, mSceneContainerPlugin);
+        sysuiState.addCallback(mCallback);
+        return sysuiState;
+    }
 
     @Before
     public void setup() {
-        FakeDisplayTracker displayTracker = new FakeDisplayTracker(mContext);
         mKosmos = new KosmosJavaAdapter(this);
-        mFlagsContainer = new SysUiState(displayTracker, mKosmos.getSceneContainerPlugin());
+        mDisplayTracker = mKosmos.getDisplayTracker();
+        mFlagsContainer = mKosmos.getSysuiState();
+        mSceneContainerPlugin = mKosmos.getSceneContainerPlugin();
         mCallback = mock(SysUiState.SysUiStateCallback.class);
-        mFlagsContainer.addCallback(mCallback);
+        mFlagsContainer = createInstance(DEFAULT_DISPLAY);
     }
 
     @Test
@@ -69,20 +80,17 @@ public class SysUiStateTest extends SysuiTestCase {
         setFlags(FLAG_2);
 
         verify(mCallback, times(1)).onSystemUiStateChanged(FLAG_1);
-        verify(mCallback, times(1))
-                .onSystemUiStateChanged(FLAG_1 | FLAG_2);
+        verify(mCallback, times(1)).onSystemUiStateChanged(FLAG_1 | FLAG_2);
     }
 
     @Test
     public void addMultipleRemoveOne_setFlag() {
         setFlags(FLAG_1);
         setFlags(FLAG_2);
-        mFlagsContainer.setFlag(FLAG_1, false)
-                .commitUpdate(DISPLAY_ID);
+        mFlagsContainer.setFlag(FLAG_1, false).commitUpdate(DISPLAY_ID);
 
         verify(mCallback, times(1)).onSystemUiStateChanged(FLAG_1);
-        verify(mCallback, times(1))
-                .onSystemUiStateChanged(FLAG_1 | FLAG_2);
+        verify(mCallback, times(1)).onSystemUiStateChanged(FLAG_1 | FLAG_2);
         verify(mCallback, times(1)).onSystemUiStateChanged(FLAG_2);
     }
 
@@ -97,8 +105,7 @@ public class SysUiStateTest extends SysuiTestCase {
     @Test
     public void addMultipleRemoveOne_setFlags() {
         setFlags(FLAG_1, FLAG_2, FLAG_3, FLAG_4);
-        mFlagsContainer.setFlag(FLAG_2, false)
-                .commitUpdate(DISPLAY_ID);
+        mFlagsContainer.setFlag(FLAG_2, false).commitUpdate(DISPLAY_ID);
 
         int expected1 = FLAG_1 | FLAG_2 | FLAG_3 | FLAG_4;
         verify(mCallback, times(1)).onSystemUiStateChanged(expected1);
@@ -115,10 +122,31 @@ public class SysUiStateTest extends SysuiTestCase {
         verify(mCallback, times(0)).onSystemUiStateChanged(expected);
     }
 
+    @Test
+    public void setFlag_receivedForDefaultDisplay() {
+        setFlags(FLAG_1);
+
+        verify(mCallback, times(1)).onSystemUiStateChangedForDisplay(FLAG_1, DEFAULT_DISPLAY);
+    }
+
+    @Test
+    public void setFlag_externalDisplayInstance_defaultDisplayCallbackNotPropagated() {
+        var instance = createInstance(/* displayId = */ 2);
+        reset(mCallback);
+        setFlags(instance, FLAG_1);
+
+        verify(mCallback, times(1)).onSystemUiStateChangedForDisplay(FLAG_1, /* displayId= */ 2);
+        verify(mCallback, never()).onSystemUiStateChanged(FLAG_1);
+    }
+
     private void setFlags(int... flags) {
-        for (int i = 0; i < flags.length; i++) {
-            mFlagsContainer.setFlag(flags[i], true);
+        setFlags(mFlagsContainer, flags);
+    }
+
+    private void setFlags(SysUiState instance, int... flags) {
+        for (int flag : flags) {
+            instance.setFlag(flag, true);
         }
-        mFlagsContainer.commitUpdate(DISPLAY_ID);
+        instance.commitUpdate();
     }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/SystemUIDialogTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/SystemUIDialogTest.java
index 58856d970711..ffde34efa7ad 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/SystemUIDialogTest.java
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/SystemUIDialogTest.java
@@ -48,7 +48,6 @@ import com.android.systemui.animation.back.BackAnimationSpec;
 import com.android.systemui.broadcast.BroadcastDispatcher;
 import com.android.systemui.kosmos.KosmosJavaAdapter;
 import com.android.systemui.model.SysUiState;
-import com.android.systemui.settings.FakeDisplayTracker;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -84,8 +83,7 @@ public class SystemUIDialogTest extends SysuiTestCase {
     public void setup() {
         MockitoAnnotations.initMocks(this);
         KosmosJavaAdapter kosmos = new KosmosJavaAdapter(this);
-        FakeDisplayTracker displayTracker = new FakeDisplayTracker(mContext);
-        mSysUiState = new SysUiState(displayTracker, kosmos.getSceneContainerPlugin());
+        mSysUiState = kosmos.getSysuiState();
         mDependency.injectTestDependency(BroadcastDispatcher.class, mBroadcastDispatcher);
         when(mDelegate.getBackAnimationSpec(ArgumentMatchers.any()))
                 .thenReturn(mock(BackAnimationSpec.class));
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java b/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
index f8cf6b007041..bbc470c77ee4 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
+++ b/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
@@ -21,6 +21,7 @@ import android.app.Service;
 import android.app.backup.BackupManager;
 import android.content.Context;
 import android.service.dreams.IDreamManager;
+import android.view.Display;
 
 import androidx.annotation.Nullable;
 
@@ -86,8 +87,8 @@ import com.android.systemui.mediaprojection.MediaProjectionModule;
 import com.android.systemui.mediaprojection.appselector.MediaProjectionActivitiesModule;
 import com.android.systemui.mediaprojection.taskswitcher.MediaProjectionTaskSwitcherModule;
 import com.android.systemui.mediarouter.MediaRouterModule;
-import com.android.systemui.model.SceneContainerPlugin;
 import com.android.systemui.model.SysUiState;
+import com.android.systemui.model.SysUiStateImpl;
 import com.android.systemui.motiontool.MotionToolModule;
 import com.android.systemui.navigationbar.NavigationBarComponent;
 import com.android.systemui.navigationbar.gestural.dagger.GestureModule;
@@ -113,7 +114,6 @@ import com.android.systemui.scene.ui.view.WindowRootViewComponent;
 import com.android.systemui.screenrecord.ScreenRecordModule;
 import com.android.systemui.screenshot.dagger.ScreenshotModule;
 import com.android.systemui.security.data.repository.SecurityRepositoryModule;
-import com.android.systemui.settings.DisplayTracker;
 import com.android.systemui.settings.UserTracker;
 import com.android.systemui.shade.ShadeController;
 import com.android.systemui.shade.ShadeDisplayAwareModule;
@@ -326,10 +326,9 @@ public abstract class SystemUIModule {
     @SysUISingleton
     @Provides
     static SysUiState provideSysUiState(
-            DisplayTracker displayTracker,
             DumpManager dumpManager,
-            SceneContainerPlugin sceneContainerPlugin) {
-        final SysUiState state = new SysUiState(displayTracker, sceneContainerPlugin);
+            SysUiStateImpl.Factory sysUiStateFactory) {
+        final SysUiState state = sysUiStateFactory.create(Display.DEFAULT_DISPLAY);
         dumpManager.registerDumpable(state);
         return state;
     }
diff --git a/packages/SystemUI/src/com/android/systemui/model/SysUiState.kt b/packages/SystemUI/src/com/android/systemui/model/SysUiState.kt
index ed190a14e680..432dadbb2136 100644
--- a/packages/SystemUI/src/com/android/systemui/model/SysUiState.kt
+++ b/packages/SystemUI/src/com/android/systemui/model/SysUiState.kt
@@ -16,26 +16,78 @@
 package com.android.systemui.model
 
 import android.util.Log
+import android.view.Display
 import com.android.systemui.Dumpable
-import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.settings.DisplayTracker
+import com.android.systemui.model.SysUiState.SysUiStateCallback
 import com.android.systemui.shared.system.QuickStepContract
 import com.android.systemui.shared.system.QuickStepContract.SystemUiStateFlags
+import dagger.assisted.Assisted
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
 import dalvik.annotation.optimization.NeverCompile
 import java.io.PrintWriter
 
 /** Contains sysUi state flags and notifies registered listeners whenever changes happen. */
-@SysUISingleton
-class SysUiState(
-    private val displayTracker: DisplayTracker,
+interface SysUiState : Dumpable {
+    /**
+     * Add listener to be notified of changes made to SysUI state.
+     *
+     * The callback will also be called as part of this function.
+     */
+    fun addCallback(callback: SysUiStateCallback)
+
+    /** Removes a callback for state changes. */
+    fun removeCallback(callback: SysUiStateCallback)
+
+    /** Returns whether a flag is enabled in this state. */
+    fun isFlagEnabled(@SystemUiStateFlags flag: Long): Boolean
+
+    /** Returns the current sysui state flags. */
+    val flags: Long
+
+    /** Methods to this call can be chained together before calling [commitUpdate]. */
+    fun setFlag(@SystemUiStateFlags flag: Long, enabled: Boolean): SysUiState
+
+    /** Call to save all the flags updated from [setFlag]. */
+    @Deprecated("Each SysUIState instance is now display specific. Just use commitUpdate()")
+    fun commitUpdate(displayId: Int)
+
+    /** Call to save all the flags updated from [setFlag]. */
+    fun commitUpdate()
+
+    /** Notify all those who are registered that the state has changed. */
+    fun notifyAndSetSystemUiStateChanged(newFlags: Long, oldFlags: Long)
+
+    /** Callback to be notified whenever system UI state flags are changed. */
+    interface SysUiStateCallback {
+        /** To be called when any SysUiStateFlag gets updated **for the default display** */
+        fun onSystemUiStateChanged(@SystemUiStateFlags sysUiFlags: Long)
+
+        /** To be called when any SysUiStateFlag gets updated for a specific [displayId]. */
+        fun onSystemUiStateChangedForDisplay(
+            @SystemUiStateFlags sysUiFlags: Long,
+            displayId: Int,
+        ) {}
+    }
+
+    companion object {
+        const val DEBUG: Boolean = false
+    }
+}
+
+class SysUiStateImpl
+@AssistedInject
+constructor(
+    @Assisted private val displayId: Int,
     private val sceneContainerPlugin: SceneContainerPlugin?,
-) : Dumpable {
+) : SysUiState {
     /** Returns the current sysui state flags. */
     @get:SystemUiStateFlags
     @SystemUiStateFlags
-    var flags: Long = 0
-        private set
+    override val flags: Long
+        get() = _flags
 
+    private var _flags: Long = 0
     private val callbacks: MutableList<SysUiStateCallback> = ArrayList()
     private var flagsToSet: Long = 0
     private var flagsToClear: Long = 0
@@ -44,26 +96,26 @@ class SysUiState(
      * Add listener to be notified of changes made to SysUI state. The callback will also be called
      * as part of this function.
      */
-    fun addCallback(callback: SysUiStateCallback) {
+    override fun addCallback(callback: SysUiStateCallback) {
         callbacks.add(callback)
         callback.onSystemUiStateChanged(flags)
     }
 
     /** Callback will no longer receive events on state change */
-    fun removeCallback(callback: SysUiStateCallback) {
+    override fun removeCallback(callback: SysUiStateCallback) {
         callbacks.remove(callback)
     }
 
-    fun isFlagEnabled(@SystemUiStateFlags flag: Long): Boolean {
+    override fun isFlagEnabled(@SystemUiStateFlags flag: Long): Boolean {
         return (flags and flag) != 0L
     }
 
     /** Methods to this call can be chained together before calling [.commitUpdate]. */
-    fun setFlag(@SystemUiStateFlags flag: Long, enabled: Boolean): SysUiState {
+    override fun setFlag(@SystemUiStateFlags flag: Long, enabled: Boolean): SysUiState {
         var enabled = enabled
         val overrideOrNull = sceneContainerPlugin?.flagValueOverride(flag)
         if (overrideOrNull != null && enabled != overrideOrNull) {
-            if (DEBUG) {
+            if (SysUiState.DEBUG) {
                 Log.d(
                     TAG,
                     "setFlag for flag $flag and value $enabled overridden to $overrideOrNull by scene container plugin",
@@ -81,20 +133,22 @@ class SysUiState(
         return this
     }
 
-    /** Call to save all the flags updated from [.setFlag]. */
-    fun commitUpdate(displayId: Int) {
-        updateFlags(displayId)
+    @Deprecated(
+        "Each SysUIState instance is now display specific. Just use commitUpdate.",
+        ReplaceWith("commitUpdate()"),
+    )
+    override fun commitUpdate(displayId: Int) {
+        // TODO b/398011576 - handle updates for different displays.
+        commitUpdate()
+    }
+
+    override fun commitUpdate() {
+        updateFlags()
         flagsToSet = 0
         flagsToClear = 0
     }
 
-    private fun updateFlags(displayId: Int) {
-        if (displayId != displayTracker.defaultDisplayId) {
-            // Ignore non-default displays for now
-            Log.w(TAG, "Ignoring flag update for display: $displayId", Throwable())
-            return
-        }
-
+    private fun updateFlags() {
         var newState = flags
         newState = newState or flagsToSet
         newState = newState and flagsToClear.inv()
@@ -102,16 +156,22 @@ class SysUiState(
     }
 
     /** Notify all those who are registered that the state has changed. */
-    private fun notifyAndSetSystemUiStateChanged(newFlags: Long, oldFlags: Long) {
-        if (DEBUG) {
+    override fun notifyAndSetSystemUiStateChanged(newFlags: Long, oldFlags: Long) {
+        if (SysUiState.DEBUG) {
             Log.d(TAG, "SysUiState changed: old=$oldFlags new=$newFlags")
         }
         if (newFlags != oldFlags) {
             callbacks.forEach { callback: SysUiStateCallback ->
-                callback.onSystemUiStateChanged(newFlags)
+                if (displayId == Display.DEFAULT_DISPLAY) {
+                    callback.onSystemUiStateChanged(newFlags)
+                }
+                callback.onSystemUiStateChangedForDisplay(
+                    sysUiFlags = newFlags,
+                    displayId = displayId,
+                )
             }
 
-            flags = newFlags
+            _flags = newFlags
         }
     }
 
@@ -127,14 +187,13 @@ class SysUiState(
         pw.println(QuickStepContract.isAssistantGestureDisabled(flags))
     }
 
-    /** Callback to be notified whenever system UI state flags are changed. */
-    interface SysUiStateCallback {
-        /** To be called when any SysUiStateFlag gets updated */
-        fun onSystemUiStateChanged(@SystemUiStateFlags sysUiFlags: Long)
+    @AssistedFactory
+    interface Factory {
+        /** Creates a new instance of [SysUiStateImpl] for a given [displayId]. */
+        fun create(displayId: Int): SysUiStateImpl
     }
 
     companion object {
         private val TAG: String = SysUiState::class.java.simpleName
-        const val DEBUG: Boolean = false
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/accessibility/WindowMagnificationControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/accessibility/WindowMagnificationControllerTest.java
index 7cf93277bb5b..5b32b922d377 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/accessibility/WindowMagnificationControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/accessibility/WindowMagnificationControllerTest.java
@@ -191,7 +191,7 @@ public class WindowMagnificationControllerTest extends SysuiTestCase {
         mWindowManager = spy(new TestableWindowManager(wm));
 
         mContext.addMockSystemService(Context.WINDOW_SERVICE, mWindowManager);
-        mSysUiState = new SysUiState(mDisplayTracker, mKosmos.getSceneContainerPlugin());
+        mSysUiState = mKosmos.getSysuiState();
         mSysUiState.addCallback(Mockito.mock(SysUiState.SysUiStateCallback.class));
         when(mSecureSettings.getIntForUser(anyString(), anyInt(), anyInt())).then(
                 returnsSecondArg());
diff --git a/packages/SystemUI/tests/src/com/android/systemui/recents/LauncherProxyServiceTest.kt b/packages/SystemUI/tests/src/com/android/systemui/recents/LauncherProxyServiceTest.kt
index 40547c2787ac..ec37b7592650 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/recents/LauncherProxyServiceTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/recents/LauncherProxyServiceTest.kt
@@ -35,8 +35,7 @@ import com.android.systemui.dump.DumpManager
 import com.android.systemui.keyguard.KeyguardUnlockAnimationController
 import com.android.systemui.keyguard.WakefulnessLifecycle
 import com.android.systemui.keyguard.ui.view.InWindowLauncherUnlockAnimationManager
-import com.android.systemui.model.SysUiState
-import com.android.systemui.model.sceneContainerPlugin
+import com.android.systemui.model.sysUiState
 import com.android.systemui.navigationbar.NavigationBarController
 import com.android.systemui.navigationbar.NavigationModeController
 import com.android.systemui.process.ProcessWrapper
@@ -93,7 +92,7 @@ class LauncherProxyServiceTest : SysuiTestCase() {
     @Mock private lateinit var processWrapper: ProcessWrapper
     private val displayTracker = FakeDisplayTracker(mContext)
     private val fakeSystemClock = FakeSystemClock()
-    private val sysUiState = SysUiState(displayTracker, kosmos.sceneContainerPlugin)
+    private val sysUiState = kosmos.sysUiState
     private val wakefulnessLifecycle =
         WakefulnessLifecycle(mContext, null, fakeSystemClock, dumpManager)
 
@@ -161,7 +160,8 @@ class LauncherProxyServiceTest : SysuiTestCase() {
         wakefulnessLifecycle.dispatchFinishedGoingToSleep()
         clearInvocations(launcherProxy)
 
-        wakefulnessLifecycle.dispatchFinishedWakingUp()
+        wakefulnessLifecycle
+            .dispatchFinishedWakingUp()
 
         verify(launcherProxy)
             .onSystemUiStateChanged(
diff --git a/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java b/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
index 68d84ecaf4b1..59c6c9bb17fd 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
@@ -33,8 +33,6 @@ import static com.android.wm.shell.Flags.FLAG_ENABLE_BUBBLE_BAR;
 
 import static com.google.common.truth.Truth.assertThat;
 
-import static kotlinx.coroutines.flow.StateFlowKt.MutableStateFlow;
-
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
@@ -57,6 +55,8 @@ import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
 
+import static kotlinx.coroutines.flow.StateFlowKt.MutableStateFlow;
+
 import android.app.ActivityManager;
 import android.app.IActivityManager;
 import android.app.INotificationManager;
@@ -460,7 +460,7 @@ public class BubblesTest extends SysuiTestCase {
         mZenModeConfig.suppressedVisualEffects = 0;
         when(mZenModeController.getConfig()).thenReturn(mZenModeConfig);
 
-        mSysUiState = new SysUiState(mDisplayTracker, mKosmos.getSceneContainerPlugin());
+        mSysUiState = mKosmos.getSysuiState();
         mSysUiState.addCallback(sysUiFlags -> {
             mSysUiStateBubblesManageMenuExpanded =
                     (sysUiFlags
@@ -620,7 +620,8 @@ public class BubblesTest extends SysuiTestCase {
                     TAG,
                     String.format("waiting for animations to complete. attempt %d", retryCount));
             // post a message to the looper and wait for it to be processed
-            mTestableLooper.runWithLooper(() -> {});
+            mTestableLooper.runWithLooper(() -> {
+            });
             retryCount++;
         }
         mTestableLooper.processAllMessages();
@@ -2895,9 +2896,11 @@ public class BubblesTest extends SysuiTestCase {
         }
 
         @Override
-        public void onDragItemOverBubbleBarDragZone(@NonNull BubbleBarLocation location) {}
+        public void onDragItemOverBubbleBarDragZone(@NonNull BubbleBarLocation location) {
+        }
 
         @Override
-        public void onItemDraggedOutsideBubbleBarDropZone() {}
+        public void onItemDraggedOutsideBubbleBarDropZone() {
+        }
     }
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/kosmos/KosmosJavaAdapter.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/kosmos/KosmosJavaAdapter.kt
index 02cf1f5a7214..f40dd50e347d 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/kosmos/KosmosJavaAdapter.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/kosmos/KosmosJavaAdapter.kt
@@ -52,6 +52,7 @@ import com.android.systemui.keyguard.domain.interactor.pulseExpansionInteractor
 import com.android.systemui.keyguard.ui.viewmodel.glanceableHubToLockscreenTransitionViewModel
 import com.android.systemui.keyguard.ui.viewmodel.lockscreenToGlanceableHubTransitionViewModel
 import com.android.systemui.model.sceneContainerPlugin
+import com.android.systemui.model.sysUiState
 import com.android.systemui.plugins.statusbar.statusBarStateController
 import com.android.systemui.power.data.repository.fakePowerRepository
 import com.android.systemui.power.domain.interactor.powerInteractor
@@ -63,6 +64,7 @@ import com.android.systemui.scene.sceneContainerConfig
 import com.android.systemui.scene.shared.model.sceneDataSource
 import com.android.systemui.scene.ui.view.mockWindowRootViewProvider
 import com.android.systemui.settings.brightness.data.repository.brightnessMirrorShowingRepository
+import com.android.systemui.settings.displayTracker
 import com.android.systemui.shade.data.repository.shadeRepository
 import com.android.systemui.shade.domain.interactor.shadeInteractor
 import com.android.systemui.shade.domain.interactor.shadeLayoutParams
@@ -194,4 +196,6 @@ class KosmosJavaAdapter() {
     val fakeDisableFlagsRepository by lazy { kosmos.fakeDisableFlagsRepository }
     val mockWindowRootViewProvider by lazy { kosmos.mockWindowRootViewProvider }
     val windowRootViewBlurInteractor by lazy { kosmos.windowRootViewBlurInteractor }
+    val sysuiState by lazy { kosmos.sysUiState }
+    val displayTracker by lazy { kosmos.displayTracker }
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/model/SysUiStateKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/model/SysUiStateKosmos.kt
index 6ddf633e58e5..8aecf886b578 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/model/SysUiStateKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/model/SysUiStateKosmos.kt
@@ -16,16 +16,11 @@
 
 package com.android.systemui.model
 
+import android.view.Display
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
-import com.android.systemui.settings.displayTracker
 import org.mockito.Mockito.spy
 
 val Kosmos.sysUiState by Fixture {
-    spy(
-        SysUiState(
-            displayTracker,
-            sceneContainerPlugin,
-        )
-    )
+    spy(SysUiStateImpl(Display.DEFAULT_DISPLAY, sceneContainerPlugin))
 }