Update shade transition on full expand or collapse

On the lockscreen, this is normally done by a call in
LockscreenShadeTransitionController when updating dragDownAmount.
However it seems that sometimes this update doesn't happen, leaving
MediaHierarchyManager in a bad state. This change listens to the
ShadeInteractor as a backup to ensure the state is updated when the
transition has finished.

Bug: 319244625
Flag: com.android.systemui.media_controls_lockscreen_shade_bug_fix
Test: manual - swipe shade while locked and unlocked and landscape mode
Test: manual - same on tablet device
Change-Id: I5c310ab038603c33f0a6a202ebcee9110e0612d9
diff --git a/packages/SystemUI/aconfig/systemui.aconfig b/packages/SystemUI/aconfig/systemui.aconfig
index 4311e79..6b3ddc7 100644
--- a/packages/SystemUI/aconfig/systemui.aconfig
+++ b/packages/SystemUI/aconfig/systemui.aconfig
@@ -981,6 +981,16 @@
 }
 
 flag {
+  name: "media_controls_lockscreen_shade_bug_fix"
+  namespace: "systemui"
+  description: "Use ShadeInteractor for media location changes"
+  bug: "319244625"
+  metadata {
+    purpose: PURPOSE_BUGFIX
+  }
+}
+
+flag {
   namespace: "systemui"
   name: "enable_view_capture_tracing"
   description: "Enables view capture tracing in System UI."
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/ui/controller/MediaHierarchyManager.kt b/packages/SystemUI/src/com/android/systemui/media/controls/ui/controller/MediaHierarchyManager.kt
index 601d563..88a28bf 100644
--- a/packages/SystemUI/src/com/android/systemui/media/controls/ui/controller/MediaHierarchyManager.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/ui/controller/MediaHierarchyManager.kt
@@ -36,6 +36,7 @@
 import com.android.app.animation.Interpolators
 import com.android.app.tracing.traceSection
 import com.android.keyguard.KeyguardViewController
+import com.android.systemui.Flags.mediaControlsLockscreenShadeBugFix
 import com.android.systemui.communal.ui.viewmodel.CommunalTransitionViewModel
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
@@ -483,8 +484,7 @@
             object : StatusBarStateController.StateListener {
                 override fun onStatePreChange(oldState: Int, newState: Int) {
                     // We're updating the location before the state change happens, since we want
-                    // the
-                    // location of the previous state to still be up to date when the animation
+                    // the location of the previous state to still be up to date when the animation
                     // starts
                     if (
                         newState == StatusBarState.SHADE_LOCKED &&
@@ -588,6 +588,17 @@
             }
         }
 
+        if (mediaControlsLockscreenShadeBugFix()) {
+            coroutineScope.launch {
+                shadeInteractor.shadeExpansion.collect { expansion ->
+                    if (expansion >= 1f || expansion <= 0f) {
+                        // Shade has fully expanded or collapsed: force transition amount update
+                        setTransitionToFullShadeAmount(expansion)
+                    }
+                }
+            }
+        }
+
         val settingsObserver: ContentObserver =
             object : ContentObserver(handler) {
                 override fun onChange(selfChange: Boolean, uri: Uri?) {