From d2ea106305f6b8d17adccce6b99ca9ab10781771 Mon Sep 17 00:00:00 2001
From: Ajinkya Chalke <achalke@google.com>
Date: Wed, 28 May 2025 14:24:20 +0000
Subject: [PATCH] Update DisplayControllerSpy for overview on CD

- Use both infoModifier and infoModifierForDisplay to support overview
  on CD changes in DisplayController.

Test: TaskbarModeRuleTest TaskbarInsetsControllerTes
Bug: 363251602
Flag: EXEMPT test only change
Change-Id: I18f2dd7f6c486a269769d93f5ad7623bcba5fe3c
---
 .../taskbar/rules/TaskbarSandboxComponent.kt    | 17 +++++++++++++++--
 .../taskbar/rules/TaskbarUnitTestRule.kt        |  2 +-
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarSandboxComponent.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarSandboxComponent.kt
index 6371f9ad28..1f9cb1150f 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarSandboxComponent.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarSandboxComponent.kt
@@ -107,7 +107,18 @@ constructor(
 ) : DisplayController(context, wmProxy, prefs, lifecycle) {
 
     var infoModifier: ((Info) -> Info)? = null
-    var infoModifierForDisplay: ((Info?) -> Info?)? = null
+
+    // When overview on CD is enabled, DisplayController queries getInfoForDisplay instead of
+    // getInfo for the primary (virtual) display used in tests. So, override it to get info from the
+    // default display.
+    private val defaultInfoModifierForDisplay: ((Info?) -> Info?)? =
+        if (Flags.enableOverviewOnConnectedDisplays()) {
+            { _ -> info }
+        } else {
+            null
+        }
+
+    var infoModifierForDisplay: ((Info?) -> Info?)? = defaultInfoModifierForDisplay
 
     private var prefListener: LauncherPrefChangeListener? = null
 
@@ -141,8 +152,10 @@ constructor(
                 .also { prefs.addListener(it, TASKBAR_PINNING) }
     }
 
-    fun removeTaskbarPinningPrefListener() {
+    fun cleanup() {
         prefListener?.let { prefs.removeListener(it, TASKBAR_PINNING) }
+        infoModifier = null
+        infoModifierForDisplay = defaultInfoModifierForDisplay
     }
 }
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt
index 21501562da..210f03c386 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt
@@ -156,7 +156,7 @@ class TaskbarUnitTestRule(
                     base.evaluate()
                 } finally {
                     instrumentation.runOnMainSync { taskbarManager.destroy() }
-                    context.displayControllerSpy?.removeTaskbarPinningPrefListener()
+                    context.displayControllerSpy?.cleanup()
                 }
             }
         }