[flexiglass] Revamp SharedNotificationContainerViewModelTest

- Adds an extension function which allows parameterizing SceneContainerFlag
- Adds parameterization to SharedNotificationContainerViewModelTest
- Disables some tests under scene container where we have flagged off the code

Bug: 332750091
Test: atest SystemUITest
Test: atest SystemUiRoboTest
Flag: ACONFIG com.android.systemui.scene_container DEVELOPMENT
Flag: ACONFIG com.android.systemui.migrate_clocks_to_blueprint TEAMFOOD
Change-Id: Ief5202cf04d3b7692593f3e15d8893ea3dee2a18
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/shared/flag/SceneContainerFlagParameterizationTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/shared/flag/SceneContainerFlagParameterizationTest.kt
new file mode 100644
index 0000000..db31ad5
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/shared/flag/SceneContainerFlagParameterizationTest.kt
@@ -0,0 +1,70 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.scene.shared.flag
+
+import android.platform.test.flag.junit.FlagsParameterization
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.Flags.FLAG_COMPOSE_LOCKSCREEN
+import com.android.systemui.Flags.FLAG_EXAMPLE_FLAG
+import com.android.systemui.Flags.FLAG_SCENE_CONTAINER
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.flags.andSceneContainer
+import com.google.common.truth.Truth
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+internal class SceneContainerFlagParameterizationTest : SysuiTestCase() {
+
+    @Test
+    fun emptyAndSceneContainer() {
+        val result = FlagsParameterization.allCombinationsOf().andSceneContainer()
+        Truth.assertThat(result).hasSize(2)
+        Truth.assertThat(result[0].mOverrides[FLAG_SCENE_CONTAINER]).isFalse()
+        Truth.assertThat(result[1].mOverrides[FLAG_SCENE_CONTAINER]).isTrue()
+    }
+
+    @Test
+    fun oneUnrelatedAndSceneContainer() {
+        val unrelatedFlag = FLAG_EXAMPLE_FLAG
+        val result = FlagsParameterization.allCombinationsOf(unrelatedFlag).andSceneContainer()
+        Truth.assertThat(result).hasSize(4)
+        Truth.assertThat(result[0].mOverrides[unrelatedFlag]).isFalse()
+        Truth.assertThat(result[0].mOverrides[FLAG_SCENE_CONTAINER]).isFalse()
+        Truth.assertThat(result[1].mOverrides[unrelatedFlag]).isFalse()
+        Truth.assertThat(result[1].mOverrides[FLAG_SCENE_CONTAINER]).isTrue()
+        Truth.assertThat(result[2].mOverrides[unrelatedFlag]).isTrue()
+        Truth.assertThat(result[2].mOverrides[FLAG_SCENE_CONTAINER]).isFalse()
+        Truth.assertThat(result[3].mOverrides[unrelatedFlag]).isTrue()
+        Truth.assertThat(result[3].mOverrides[FLAG_SCENE_CONTAINER]).isTrue()
+    }
+
+    @Test
+    fun oneDependencyAndSceneContainer() {
+        val dependentFlag = FLAG_COMPOSE_LOCKSCREEN
+        val result = FlagsParameterization.allCombinationsOf(dependentFlag).andSceneContainer()
+        Truth.assertThat(result).hasSize(3)
+        Truth.assertThat(result[0].mOverrides[dependentFlag]).isFalse()
+        Truth.assertThat(result[0].mOverrides[FLAG_SCENE_CONTAINER]).isFalse()
+        Truth.assertThat(result[1].mOverrides[dependentFlag]).isTrue()
+        Truth.assertThat(result[1].mOverrides[FLAG_SCENE_CONTAINER]).isFalse()
+        Truth.assertThat(result[2].mOverrides[dependentFlag]).isTrue()
+        Truth.assertThat(result[2].mOverrides[FLAG_SCENE_CONTAINER]).isTrue()
+    }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelTest.kt
index ac8387f..7d69382 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelTest.kt
@@ -20,9 +20,11 @@
 package com.android.systemui.statusbar.notification.stack.ui.viewmodel
 
 import android.platform.test.annotations.DisableFlags
-import androidx.test.ext.junit.runners.AndroidJUnit4
+import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.FlagsParameterization
 import androidx.test.filters.SmallTest
 import com.android.systemui.Flags.FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX
+import com.android.systemui.Flags.FLAG_MIGRATE_CLOCKS_TO_BLUEPRINT
 import com.android.systemui.Flags.FLAG_SCENE_CONTAINER
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.common.shared.model.NotificationContainerBounds
@@ -31,6 +33,7 @@
 import com.android.systemui.coroutines.collectValues
 import com.android.systemui.flags.EnableSceneContainer
 import com.android.systemui.flags.Flags
+import com.android.systemui.flags.andSceneContainer
 import com.android.systemui.flags.fakeFeatureFlagsClassic
 import com.android.systemui.keyguard.data.repository.fakeKeyguardRepository
 import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
@@ -64,10 +67,30 @@
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.Mockito.mock
+import platform.test.runner.parameterized.ParameterizedAndroidJunit4
+import platform.test.runner.parameterized.Parameters
 
 @SmallTest
-@RunWith(AndroidJUnit4::class)
-class SharedNotificationContainerViewModelTest : SysuiTestCase() {
+@RunWith(ParameterizedAndroidJunit4::class)
+// SharedNotificationContainerViewModel is only bound when FLAG_MIGRATE_CLOCKS_TO_BLUEPRINT is on
+@EnableFlags(FLAG_MIGRATE_CLOCKS_TO_BLUEPRINT)
+class SharedNotificationContainerViewModelTest(flags: FlagsParameterization?) : SysuiTestCase() {
+
+    companion object {
+        @JvmStatic
+        @Parameters(name = "{0}")
+        fun getParams(): List<FlagsParameterization> {
+            return FlagsParameterization.allCombinationsOf(
+                    FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX,
+                )
+                .andSceneContainer()
+        }
+    }
+
+    init {
+        mSetFlagsRule.setFlagsParameterization(flags!!)
+    }
+
     val aodBurnInViewModel = mock(AodBurnInViewModel::class.java)
     lateinit var movementFlow: MutableStateFlow<BurnInModel>
 
@@ -84,14 +107,22 @@
     }
 
     val testScope = kosmos.testScope
-    val configurationRepository = kosmos.fakeConfigurationRepository
-    val keyguardRepository = kosmos.fakeKeyguardRepository
-    val keyguardInteractor = kosmos.keyguardInteractor
-    val keyguardRootViewModel = kosmos.keyguardRootViewModel
-    val keyguardTransitionRepository = kosmos.fakeKeyguardTransitionRepository
-    val shadeRepository = kosmos.shadeRepository
-    val sharedNotificationContainerInteractor = kosmos.sharedNotificationContainerInteractor
-    val largeScreenHeaderHelper = kosmos.mockLargeScreenHeaderHelper
+    val configurationRepository
+        get() = kosmos.fakeConfigurationRepository
+    val keyguardRepository
+        get() = kosmos.fakeKeyguardRepository
+    val keyguardInteractor
+        get() = kosmos.keyguardInteractor
+    val keyguardRootViewModel
+        get() = kosmos.keyguardRootViewModel
+    val keyguardTransitionRepository
+        get() = kosmos.fakeKeyguardTransitionRepository
+    val shadeRepository
+        get() = kosmos.shadeRepository
+    val sharedNotificationContainerInteractor
+        get() = kosmos.sharedNotificationContainerInteractor
+    val largeScreenHeaderHelper
+        get() = kosmos.mockLargeScreenHeaderHelper
 
     lateinit var underTest: SharedNotificationContainerViewModel
 
@@ -130,9 +161,9 @@
         }
 
     @Test
+    @DisableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
     fun validatePaddingTopInSplitShade_refactorFlagOff_usesLargeHeaderResource() =
         testScope.runTest {
-            mSetFlagsRule.disableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
             whenever(largeScreenHeaderHelper.getLargeScreenHeaderHeight()).thenReturn(5)
             overrideResource(R.bool.config_use_split_notification_shade, true)
             overrideResource(R.bool.config_use_large_screen_shade_header, true)
@@ -148,9 +179,9 @@
         }
 
     @Test
+    @EnableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
     fun validatePaddingTopInSplitShade_refactorFlagOn_usesLargeHeaderHelper() =
         testScope.runTest {
-            mSetFlagsRule.enableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
             whenever(largeScreenHeaderHelper.getLargeScreenHeaderHeight()).thenReturn(5)
             overrideResource(R.bool.config_use_split_notification_shade, true)
             overrideResource(R.bool.config_use_large_screen_shade_header, true)
@@ -243,9 +274,9 @@
 
     @Test
     @DisableFlags(FLAG_SCENE_CONTAINER)
+    @EnableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
     fun validateMarginTopWithLargeScreenHeader_refactorFlagOn_usesHelper() =
         testScope.runTest {
-            mSetFlagsRule.enableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
             val headerResourceHeight = 50
             val headerHelperHeight = 100
             whenever(largeScreenHeaderHelper.getLargeScreenHeaderHeight())
@@ -263,9 +294,9 @@
 
     @Test
     @EnableSceneContainer
+    @EnableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
     fun validateMarginTopWithLargeScreenHeader_sceneContainerFlagOn_stillZero() =
         testScope.runTest {
-            mSetFlagsRule.enableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
             val headerResourceHeight = 50
             val headerHelperHeight = 100
             whenever(largeScreenHeaderHelper.getLargeScreenHeaderHeight())
@@ -503,6 +534,7 @@
         }
 
     @Test
+    @DisableFlags(FLAG_SCENE_CONTAINER)
     fun boundsOnLockscreenNotInSplitShade() =
         testScope.runTest {
             val bounds by collectLastValue(underTest.bounds)
@@ -523,9 +555,9 @@
         }
 
     @Test
+    @DisableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX, FLAG_SCENE_CONTAINER)
     fun boundsOnLockscreenInSplitShade_refactorFlagOff_usesLargeHeaderResource() =
         testScope.runTest {
-            mSetFlagsRule.disableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
             val bounds by collectLastValue(underTest.bounds)
 
             // When in split shade
@@ -547,13 +579,20 @@
             runCurrent()
 
             // Top should be equal to bounds (1) - padding adjustment (10)
-            assertThat(bounds).isEqualTo(NotificationContainerBounds(top = -9f, bottom = 2f))
+            assertThat(bounds)
+                .isEqualTo(
+                    NotificationContainerBounds(
+                        top = -9f,
+                        bottom = 2f,
+                    )
+                )
         }
 
     @Test
+    @EnableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
+    @DisableFlags(FLAG_SCENE_CONTAINER)
     fun boundsOnLockscreenInSplitShade_refactorFlagOn_usesLargeHeaderHelper() =
         testScope.runTest {
-            mSetFlagsRule.enableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
             val bounds by collectLastValue(underTest.bounds)
 
             // When in split shade
@@ -579,6 +618,7 @@
         }
 
     @Test
+    @DisableFlags(FLAG_SCENE_CONTAINER)
     fun boundsOnShade() =
         testScope.runTest {
             val bounds by collectLastValue(underTest.bounds)
@@ -594,6 +634,7 @@
         }
 
     @Test
+    @DisableFlags(FLAG_SCENE_CONTAINER)
     fun boundsOnQS() =
         testScope.runTest {
             val bounds by collectLastValue(underTest.bounds)
@@ -693,6 +734,7 @@
         }
 
     @Test
+    @DisableFlags(FLAG_SCENE_CONTAINER)
     fun translationYUpdatesOnKeyguardForBurnIn() =
         testScope.runTest {
             val translationY by collectLastValue(underTest.translationY(BurnInParameters()))
@@ -726,6 +768,7 @@
         }
 
     @Test
+    @DisableFlags(FLAG_SCENE_CONTAINER)
     fun translationYDoesNotUpdateWhenShadeIsExpanded() =
         testScope.runTest {
             val translationY by collectLastValue(underTest.translationY(BurnInParameters()))
@@ -746,6 +789,7 @@
         }
 
     @Test
+    @DisableFlags(FLAG_SCENE_CONTAINER)
     fun updateBounds_fromKeyguardRoot() =
         testScope.runTest {
             val bounds by collectLastValue(underTest.bounds)
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt
index 0486ef5..13e36d5 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt
@@ -350,7 +350,8 @@
      *
      * When the shade is expanding, the position is controlled by... the shade.
      */
-    val bounds: StateFlow<NotificationContainerBounds> =
+    val bounds: StateFlow<NotificationContainerBounds> by lazy {
+        SceneContainerFlag.assertInLegacyMode()
         combine(
                 isOnLockscreenWithoutShade,
                 keyguardInteractor.notificationContainerBounds,
@@ -380,6 +381,7 @@
                 initialValue = NotificationContainerBounds(),
             )
             .dumpValue("bounds")
+    }
 
     /**
      * Ensure view is visible when the shade/qs are expanded. Also, as QS is expanding, fade out
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/flags/SceneContainerFlagParameterization.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/flags/SceneContainerFlagParameterization.kt
new file mode 100644
index 0000000..4e24233
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/flags/SceneContainerFlagParameterization.kt
@@ -0,0 +1,60 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.flags
+
+import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.FlagsParameterization
+import com.android.systemui.Flags.FLAG_SCENE_CONTAINER
+
+/** The name of the one flag to be disabled for OFF parameterization */
+private const val flagNameToDisable = FLAG_SCENE_CONTAINER
+
+/** Cache of the flags to be enabled for ON parameterization */
+private val flagNamesToEnable =
+    EnableSceneContainer::class.java.getAnnotation(EnableFlags::class.java)!!.value.toList()
+
+/**
+ * Provides one or two copies of this [FlagsParameterization]; one which disabled
+ * [FLAG_SCENE_CONTAINER] and if none of the dependencies of it are disabled by this, a second copy
+ * which enables [FLAG_SCENE_CONTAINER] and all the dependencies (just like [EnableSceneContainer]).
+ */
+fun FlagsParameterization.andSceneContainer(): Sequence<FlagsParameterization> = sequence {
+    check(flagNameToDisable !in mOverrides) {
+        "Can't add $flagNameToDisable to FlagsParameterization: $this"
+    }
+    yield(FlagsParameterization(mOverrides + mapOf(flagNameToDisable to false)))
+    if (flagNamesToEnable.all { mOverrides[it] != false }) {
+        // Can't add the parameterization of enabling SceneContainerFlag to a parameterization that
+        // explicitly disables one of the prerequisite flags.
+        yield(FlagsParameterization(mOverrides + flagNamesToEnable.associateWith { true }))
+    }
+}
+
+/**
+ * Doubles (roughly; see below) the given list of [FlagsParameterization] for enabling and disabling
+ * SceneContainerFlag.
+ *
+ * The input parameterization may not define [FLAG_SCENE_CONTAINER].
+ *
+ * Any [FlagsParameterization] which disables any flag that is a dependency of
+ * [FLAG_SCENE_CONTAINER], will not add a state for enabling, and the state will simply be converted
+ * to one which disables. Just like [EnableSceneContainer], enabling will also enable all the other
+ * dependencies. For any flag parameterization where a dependency is disabled, an "enabled"
+ * parameterization is inconsistent, so it will not be added.
+ */
+fun List<FlagsParameterization>.andSceneContainer(): List<FlagsParameterization> =
+    flatMap { it.andSceneContainer() }.toList()