Add context for sandboxing Taskbar window contexts.

Otherwise, the contexts leave the sandbox, which blocks testing.

Test: TaskbarWindowSandboxContextTest
Bug: 230027385
Flag: TEST_ONLY
Change-Id: I4c7a60858821c8959be2c82dc61a83b7e2f4d1ed
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt
index adda51c..9ecd935 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt
@@ -30,6 +30,7 @@
 import com.android.launcher3.taskbar.overlay.TaskbarOverlayController
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.InjectController
+import com.android.launcher3.taskbar.rules.TaskbarWindowSandboxContext
 import com.android.launcher3.util.LauncherMultivalentJUnit
 import com.android.launcher3.util.LauncherMultivalentJUnit.EmulatedDevices
 import com.android.launcher3.util.PackageUserKey
@@ -43,7 +44,11 @@
 class TaskbarAllAppsControllerTest {
 
     @get:Rule
-    val taskbarUnitTestRule = TaskbarUnitTestRule(this, getInstrumentation().targetContext)
+    val taskbarUnitTestRule =
+        TaskbarUnitTestRule(
+            this,
+            TaskbarWindowSandboxContext.create(getInstrumentation().targetContext),
+        )
     @get:Rule val animatorTestRule = AnimatorTestRule(this)
 
     @InjectController lateinit var allAppsController: TaskbarAllAppsController
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt
index 4f83ac0..fae5562 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt
@@ -28,6 +28,7 @@
 import com.android.launcher3.taskbar.TaskbarActivityContext
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.InjectController
+import com.android.launcher3.taskbar.rules.TaskbarWindowSandboxContext
 import com.android.launcher3.util.LauncherMultivalentJUnit
 import com.android.launcher3.util.LauncherMultivalentJUnit.EmulatedDevices
 import com.android.launcher3.views.BaseDragLayer
@@ -42,7 +43,11 @@
 class TaskbarOverlayControllerTest {
 
     @get:Rule
-    val taskbarUnitTestRule = TaskbarUnitTestRule(this, getInstrumentation().targetContext)
+    val taskbarUnitTestRule =
+        TaskbarUnitTestRule(
+            this,
+            TaskbarWindowSandboxContext.create(getInstrumentation().targetContext),
+        )
     @InjectController lateinit var overlayController: TaskbarOverlayController
 
     private val taskbarContext: TaskbarActivityContext
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRule.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRule.kt
index d3781e8..6638736 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRule.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRule.kt
@@ -20,7 +20,6 @@
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.TaskbarMode
 import com.android.launcher3.util.DisplayController
 import com.android.launcher3.util.MainThreadInitializedObject
-import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext
 import com.android.launcher3.util.NavigationMode
 import org.junit.rules.TestRule
 import org.junit.runner.Description
@@ -40,7 +39,7 @@
  * Make sure this rule precedes any rules that depend on [DisplayController], or else the instance
  * might be inconsistent across the test lifecycle.
  */
-class TaskbarModeRule(private val context: SandboxContext) : TestRule {
+class TaskbarModeRule(private val context: TaskbarWindowSandboxContext) : TestRule {
     /** The selected Taskbar mode. */
     enum class Mode {
         TRANSIENT,
@@ -60,7 +59,7 @@
             override fun evaluate() {
                 val mode = taskbarMode.mode
 
-                context.putObject(
+                context.applicationContext.putObject(
                     DisplayController.INSTANCE,
                     object : DisplayController(context) {
                         override fun getInfo(): Info {
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRuleTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRuleTest.kt
index e922bab..f75e542 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRuleTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRuleTest.kt
@@ -24,7 +24,6 @@
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.TaskbarMode
 import com.android.launcher3.util.DisplayController
 import com.android.launcher3.util.LauncherMultivalentJUnit
-import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext
 import com.android.launcher3.util.NavigationMode
 import com.google.common.truth.Truth.assertThat
 import org.junit.Rule
@@ -34,7 +33,7 @@
 @RunWith(LauncherMultivalentJUnit::class)
 class TaskbarModeRuleTest {
 
-    private val context = SandboxContext(getInstrumentation().targetContext)
+    private val context = TaskbarWindowSandboxContext.create(getInstrumentation().targetContext)
 
     @get:Rule val taskbarModeRule = TaskbarModeRule(context)
 
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 88df029..12f946e 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
@@ -18,7 +18,6 @@
 
 import android.app.Instrumentation
 import android.app.PendingIntent
-import android.content.Context
 import android.content.IIntentSender
 import android.content.Intent
 import androidx.test.platform.app.InstrumentationRegistry
@@ -64,7 +63,11 @@
  * }
  * ```
  */
-class TaskbarUnitTestRule(private val testInstance: Any, private val context: Context) : TestRule {
+class TaskbarUnitTestRule(
+    private val testInstance: Any,
+    private val context: TaskbarWindowSandboxContext,
+) : TestRule {
+
     private val instrumentation = InstrumentationRegistry.getInstrumentation()
     private val serviceTestRule = ServiceTestRule()
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRuleTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRuleTest.kt
index ef8c548..8262e0f 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRuleTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRuleTest.kt
@@ -24,7 +24,6 @@
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.InjectController
 import com.android.launcher3.util.LauncherMultivalentJUnit
 import com.android.launcher3.util.LauncherMultivalentJUnit.EmulatedDevices
-import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext
 import com.google.common.truth.Truth.assertThat
 import org.junit.Assert.assertThrows
 import org.junit.Test
@@ -36,7 +35,7 @@
 @EmulatedDevices(["pixelFoldable2023"])
 class TaskbarUnitTestRuleTest {
 
-    private val context = SandboxContext(getInstrumentation().targetContext)
+    private val context = TaskbarWindowSandboxContext.create(getInstrumentation().targetContext)
 
     @Test
     fun testSetup_taskbarInitialized() {
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContext.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContext.kt
new file mode 100644
index 0000000..321e7a9
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContext.kt
@@ -0,0 +1,61 @@
+/*
+ * 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.launcher3.taskbar.rules
+
+import android.content.Context
+import android.content.ContextWrapper
+import android.os.Bundle
+import android.view.Display
+import com.android.launcher3.util.MainThreadInitializedObject
+import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext
+
+/**
+ * Sandbox wrapper where [createWindowContext] provides contexts that are still sandboxed within
+ * [application].
+ *
+ * Taskbar can create window contexts, which need to operate under the same sandbox application, but
+ * [Context.getApplicationContext] by default returns the actual application. For this reason,
+ * [SandboxContext] overrides [getApplicationContext] to return itself, which prevents leaving the
+ * sandbox. [SandboxContext] and the real application have different sets of
+ * [MainThreadInitializedObject] instances, so overriding the application prevents the latter set
+ * from leaking into the sandbox. Similarly, this implementation overrides [getApplicationContext]
+ * to return the original sandboxed [application], and it wraps created windowed contexts to
+ * propagate this [application].
+ */
+class TaskbarWindowSandboxContext
+private constructor(private val application: SandboxContext, base: Context) : ContextWrapper(base) {
+
+    override fun createWindowContext(type: Int, options: Bundle?): Context {
+        return TaskbarWindowSandboxContext(application, super.createWindowContext(type, options))
+    }
+
+    override fun createWindowContext(display: Display, type: Int, options: Bundle?): Context {
+        return TaskbarWindowSandboxContext(
+            application,
+            super.createWindowContext(display, type, options),
+        )
+    }
+
+    override fun getApplicationContext() = application
+
+    companion object {
+        /** Creates a [TaskbarWindowSandboxContext] to sandbox [base] for Taskbar tests. */
+        fun create(base: Context): TaskbarWindowSandboxContext {
+            return SandboxContext(base).let { TaskbarWindowSandboxContext(it, it) }
+        }
+    }
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContextTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContextTest.kt
new file mode 100644
index 0000000..ad4b4de
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContextTest.kt
@@ -0,0 +1,44 @@
+/*
+ * 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.launcher3.taskbar.rules
+
+import android.view.WindowManager.LayoutParams.TYPE_APPLICATION_OVERLAY
+import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
+import com.android.launcher3.util.LauncherMultivalentJUnit
+import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(LauncherMultivalentJUnit::class)
+class TaskbarWindowSandboxContextTest {
+
+    private val context = TaskbarWindowSandboxContext.create(getInstrumentation().targetContext)
+
+    @Test
+    fun testCreateWindowContext_applicationContextSandboxed() {
+        val windowContext = context.createWindowContext(TYPE_APPLICATION_OVERLAY, null)
+        assertThat(windowContext.applicationContext).isInstanceOf(SandboxContext::class.java)
+    }
+
+    @Test
+    fun testCreateWindowContext_nested_applicationContextSandboxed() {
+        val windowContext = context.createWindowContext(TYPE_APPLICATION_OVERLAY, null)
+        val nestedContext = windowContext.createWindowContext(TYPE_APPLICATION_OVERLAY, null)
+        assertThat(nestedContext.applicationContext).isInstanceOf(SandboxContext::class.java)
+    }
+}