Add @BrokenWithSceneContainer annotation to mark tests that require remedy

Test: atest SystemUITests
Flag: NA
Change-Id: I80dc27b1ac1bedf75ed71c5094cb5d03c50be6f3
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/flags/BrokenWithSceneContainer.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/flags/BrokenWithSceneContainer.kt
new file mode 100644
index 0000000..29b088b
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/flags/BrokenWithSceneContainer.kt
@@ -0,0 +1,27 @@
+/*
+ * 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 com.android.systemui.Flags.FLAG_SCENE_CONTAINER
+
+/**
+ * This is used by [SceneContainerRule] to assert that the test is broken when
+ * [FLAG_SCENE_CONTAINER] is enabled.
+ */
+@Retention(AnnotationRetention.RUNTIME)
+@Target(AnnotationTarget.FUNCTION, AnnotationTarget.CLASS)
+annotation class BrokenWithSceneContainer(val bugId: Int)
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/flags/SceneContainerRule.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/flags/SceneContainerRule.kt
index 775ad14..9ec1481 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/flags/SceneContainerRule.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/flags/SceneContainerRule.kt
@@ -18,6 +18,7 @@
 
 import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import org.junit.Assert
+import org.junit.AssumptionViolatedException
 import org.junit.rules.TestRule
 import org.junit.runner.Description
 import org.junit.runners.model.Statement
@@ -33,10 +34,7 @@
         return object : Statement() {
             @Throws(Throwable::class)
             override fun evaluate() {
-                val hasAnnotation =
-                    description?.testClass?.getAnnotation(EnableSceneContainer::class.java) !=
-                        null || description?.getAnnotation(EnableSceneContainer::class.java) != null
-                if (hasAnnotation) {
+                if (description.hasAnnotation<EnableSceneContainer>()) {
                     Assert.assertTrue(
                         "SceneContainerFlag.isEnabled is false:" +
                             "\n * Did you forget to add a new aconfig flag dependency in" +
@@ -45,8 +43,31 @@
                         SceneContainerFlag.isEnabled
                     )
                 }
+                if (
+                    description.hasAnnotation<BrokenWithSceneContainer>() &&
+                        SceneContainerFlag.isEnabled
+                ) {
+                    runCatching { base?.evaluate() }
+                        .onFailure { exception ->
+                            if (exception is AssumptionViolatedException) {
+                                throw AssertionError(
+                                    "This is marked @BrokenWithSceneContainer, but was skipped.",
+                                    exception
+                                )
+                            }
+                            throw AssumptionViolatedException("Test is still broken", exception)
+                        }
+                    throw AssertionError(
+                        "HOORAY! You fixed a test that was marked @BrokenWithSceneContainer. " +
+                            "Remove the obsolete annotation to fix this failure."
+                    )
+                }
                 base?.evaluate()
             }
         }
     }
+
+    inline fun <reified T : Annotation> Description?.hasAnnotation(): Boolean =
+        this?.testClass?.getAnnotation(T::class.java) != null ||
+            this?.getAnnotation(T::class.java) != null
 }