Use ShadeDisplayRepository only when ShadeWindowGoesAround is enabled

Test: atest SystemUITests
Bug: 369337857
Flag: com.android.systemui.shared.status_bar_connected_displays
Change-Id: Id55d9b9938ef5354c3c8a2fb228ab38415155297

2 files changed, 17 insertions, 3 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/shade/ShadeHeaderController.kt b/packages/SystemUI/src/com/android/systemui/shade/ShadeHeaderController.kt
index d82f8e722744..e7a98728423e 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/ShadeHeaderController.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/ShadeHeaderController.kt
@@ -21,6 +21,7 @@ import android.animation.AnimatorListenerAdapter
 import android.annotation.IdRes
 import android.app.PendingIntent
 import android.app.StatusBarManager
+import android.content.Context
 import android.content.Intent
 import android.content.res.Configuration
 import android.graphics.Insets
@@ -58,6 +59,7 @@ import com.android.systemui.shade.ShadeViewProviderModule.Companion.SHADE_HEADER
 import com.android.systemui.shade.carrier.ShadeCarrierGroup
 import com.android.systemui.shade.carrier.ShadeCarrierGroupController
 import com.android.systemui.shade.data.repository.ShadeDisplaysRepository
+import com.android.systemui.shade.shared.flag.ShadeWindowGoesAround
 import com.android.systemui.statusbar.data.repository.StatusBarContentInsetsProviderStore
 import com.android.systemui.statusbar.phone.StatusBarLocation
 import com.android.systemui.statusbar.phone.StatusIconContainer
@@ -70,6 +72,7 @@ import com.android.systemui.statusbar.policy.NextAlarmController
 import com.android.systemui.statusbar.policy.VariableDateView
 import com.android.systemui.statusbar.policy.VariableDateViewController
 import com.android.systemui.util.ViewController
+import dagger.Lazy
 import java.io.PrintWriter
 import javax.inject.Inject
 import javax.inject.Named
@@ -93,7 +96,8 @@ constructor(
     private val privacyIconsController: HeaderPrivacyIconsController,
     private val statusBarContentInsetsProviderStore: StatusBarContentInsetsProviderStore,
     @ShadeDisplayAware private val configurationController: ConfigurationController,
-    private val shadeDisplaysRepository: ShadeDisplaysRepository,
+    @ShadeDisplayAware private val context: Context,
+    private val shadeDisplaysRepositoryLazy: Lazy<ShadeDisplaysRepository>,
     private val variableDateViewControllerFactory: VariableDateViewController.Factory,
     @Named(SHADE_HEADER) private val batteryMeterViewController: BatteryMeterViewController,
     private val dumpManager: DumpManager,
@@ -108,7 +112,15 @@ constructor(
 
     private val statusBarContentInsetsProvider
         get() =
-            statusBarContentInsetsProviderStore.forDisplay(shadeDisplaysRepository.displayId.value)
+            statusBarContentInsetsProviderStore.forDisplay(
+                if (ShadeWindowGoesAround.isEnabled) {
+                    // ShadeDisplaysRepository is the source of truth for display id when
+                    // ShadeWindowGoesAround.isEnabled
+                    shadeDisplaysRepositoryLazy.get().displayId.value
+                } else {
+                    context.displayId
+                }
+            )
 
     companion object {
         /** IDs for transitions and constraints for the [MotionLayout]. */
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/ShadeHeaderControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/shade/ShadeHeaderControllerTest.kt
index a5cd81ff3116..9a60b264014f 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shade/ShadeHeaderControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/shade/ShadeHeaderControllerTest.kt
@@ -71,6 +71,7 @@ import com.android.systemui.util.mockito.capture
 import com.android.systemui.util.mockito.eq
 import com.android.systemui.util.mockito.mock
 import com.google.common.truth.Truth.assertThat
+import dagger.Lazy
 import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
@@ -197,7 +198,8 @@ class ShadeHeaderControllerTest : SysuiTestCase() {
                 privacyIconsController,
                 insetsProviderStore,
                 configurationController,
-                kosmos.shadeDisplaysRepository,
+                viewContext,
+                Lazy { kosmos.shadeDisplaysRepository },
                 variableDateViewControllerFactory,
                 batteryMeterViewController,
                 dumpManager,