Add PiP test for WindowLayoutComponent

Add a test to check that features before entering PiP are the same as
features after exiting PiP.

Bug:239389499
Test: atest CtsWindowManagerJetpackTestCases
Change-Id: I2e9bcd7a6571beda168301c63fd45292b59961cd
Merged-In: I2e9bcd7a6571beda168301c63fd45292b59961cd
diff --git a/tests/framework/base/windowmanager/jetpack/AndroidManifest.xml b/tests/framework/base/windowmanager/jetpack/AndroidManifest.xml
index 6361c82..b692bd1 100644
--- a/tests/framework/base/windowmanager/jetpack/AndroidManifest.xml
+++ b/tests/framework/base/windowmanager/jetpack/AndroidManifest.xml
@@ -29,6 +29,7 @@
                       android:required="false" />
         <activity android:name="android.server.wm.jetpack.utils.TestActivity" />
         <activity android:name="android.server.wm.jetpack.utils.TestConfigChangeHandlingActivity"
+                  android:supportsPictureInPicture="true"
                   android:configChanges="orientation|screenSize|smallestScreenSize|screenLayout|colorMode|density|touchscreen"
         />
         <activity android:name="android.server.wm.jetpack.utils.TestGetWindowLayoutInfoActivity" />
diff --git a/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ActivityEmbeddingBoundsTests.java b/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ActivityEmbeddingBoundsTests.java
index 8caa0de..0d4a357 100644
--- a/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ActivityEmbeddingBoundsTests.java
+++ b/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ActivityEmbeddingBoundsTests.java
@@ -22,6 +22,7 @@
 import static android.server.wm.jetpack.utils.ActivityEmbeddingUtil.startActivityAndVerifySplit;
 import static android.server.wm.jetpack.utils.ActivityEmbeddingUtil.waitAndAssertNotVisible;
 import static android.server.wm.jetpack.utils.ActivityEmbeddingUtil.waitForFillsTask;
+import static android.server.wm.jetpack.utils.TestActivityLauncher.KEY_ACTIVITY_ID;
 
 import static org.junit.Assert.assertTrue;
 
@@ -225,6 +226,6 @@
         }
         return primaryActivityId.equals(((TestActivityWithId) activityIntentPair.first).getId())
                 && secondaryActivityId.equals(activityIntentPair.second.getStringExtra(
-                        ACTIVITY_ID_LABEL));
+                KEY_ACTIVITY_ID));
     }
 }
diff --git a/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ActivityEmbeddingFinishTests.java b/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ActivityEmbeddingFinishTests.java
index f4fa70c..e017d65 100644
--- a/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ActivityEmbeddingFinishTests.java
+++ b/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ActivityEmbeddingFinishTests.java
@@ -97,8 +97,7 @@
         SplitPairRule splitPairRule = createWildcardSplitPairRule();
         mActivityEmbeddingComponent.setEmbeddingRules(Collections.singleton(splitPairRule));
 
-        TestActivity primaryActivity = (TestActivityWithId)
-                startActivityNewTask(TestActivityWithId.class);
+        TestActivity primaryActivity = startActivityNewTask(TestActivityWithId.class);
         TestActivity secondaryActivity = (TestActivity) startActivityAndVerifySplit(primaryActivity,
                 TestActivityWithId.class, splitPairRule, "secondaryActivity", mSplitInfoConsumer);
 
@@ -359,8 +358,8 @@
             mActivityEmbeddingComponent.setEmbeddingRules(Collections.singleton(splitPairRule));
 
             // Launch the two activities
-            TestActivity primaryActivity = (TestActivity)
-                    startActivityNewTask(TestConfigChangeHandlingActivity.class);
+            TestActivity primaryActivity = startActivityNewTask(
+                    TestConfigChangeHandlingActivity.class);
             TestActivity secondaryActivity;
             if (mShouldPreventSideBySideActivities) {
                 secondaryActivity = startActivityAndVerifyNotSplit(primaryActivity);
diff --git a/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ActivityEmbeddingIntegrationTests.java b/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ActivityEmbeddingIntegrationTests.java
index daf48d6..ea6f21a 100644
--- a/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ActivityEmbeddingIntegrationTests.java
+++ b/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ActivityEmbeddingIntegrationTests.java
@@ -76,8 +76,8 @@
      */
     @Test
     public void testDisplayFeaturesWithEmbedding() throws Exception {
-        TestConfigChangeHandlingActivity primaryActivity = (TestConfigChangeHandlingActivity)
-                startActivityNewTask(TestConfigChangeHandlingActivity.class);
+        TestConfigChangeHandlingActivity primaryActivity = startActivityNewTask(
+                TestConfigChangeHandlingActivity.class);
         WindowLayoutInfo windowLayoutInfo = getExtensionWindowLayoutInfo(primaryActivity);
         assumeHasDisplayFeatures(windowLayoutInfo);
 
@@ -124,8 +124,8 @@
     public void testClearSplitInfoCallback() throws Exception {
         assumeVendorApiLevelAtLeast(2); // TODO(b/244450254): harden the requirement in U.
         mActivityEmbeddingComponent.clearSplitInfoCallback();
-        TestConfigChangeHandlingActivity primaryActivity = (TestConfigChangeHandlingActivity)
-                startActivityNewTask(TestConfigChangeHandlingActivity.class);
+        TestConfigChangeHandlingActivity primaryActivity = startActivityNewTask(
+                TestConfigChangeHandlingActivity.class);
 
         // Launch a second activity in a split. Use a very small split ratio, so that the secondary
         // activity occupies most of the screen.
@@ -152,8 +152,8 @@
     @Test
     public void testDisplayFeaturesWithEmbedding_differentPackage() throws Exception {
         // Start an activity to collect the window layout info.
-        TestConfigChangeHandlingActivity initialActivity = (TestConfigChangeHandlingActivity)
-                startActivityNewTask(TestConfigChangeHandlingActivity.class);
+        TestConfigChangeHandlingActivity initialActivity = startActivityNewTask(
+                TestConfigChangeHandlingActivity.class);
         WindowLayoutInfo windowLayoutInfo = getExtensionWindowLayoutInfo(initialActivity);
         assumeHasDisplayFeatures(windowLayoutInfo);
 
diff --git a/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ActivityEmbeddingPlaceholderTests.java b/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ActivityEmbeddingPlaceholderTests.java
index 0a105cf..a9e622a 100644
--- a/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ActivityEmbeddingPlaceholderTests.java
+++ b/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ActivityEmbeddingPlaceholderTests.java
@@ -22,6 +22,7 @@
 import static android.server.wm.jetpack.utils.ActivityEmbeddingUtil.waitAndAssertFinishing;
 import static android.server.wm.jetpack.utils.ActivityEmbeddingUtil.waitAndAssertNotResumed;
 import static android.server.wm.jetpack.utils.ActivityEmbeddingUtil.waitAndAssertResumed;
+import static android.server.wm.jetpack.utils.TestActivityLauncher.KEY_ACTIVITY_ID;
 
 import static androidx.window.extensions.embedding.SplitRule.FINISH_NEVER;
 
@@ -319,7 +320,7 @@
         public SplitPlaceholderRule build() {
             // Create placeholder activity intent
             Intent placeholderIntent = new Intent(mContext, TestActivityWithId.class);
-            placeholderIntent.putExtra(ACTIVITY_ID_LABEL, mPlaceholderActivityId);
+            placeholderIntent.putExtra(KEY_ACTIVITY_ID, mPlaceholderActivityId);
 
             // Create {@link SplitPlaceholderRule} that launches the placeholder in a split with the
             // target primary activity.
@@ -329,7 +330,8 @@
                                     && mPrimaryActivityId.equals(((TestActivityWithId) activity)
                                     .getId()) /* activityPredicate */,
                             intent -> mPrimaryActivityId.equals(
-                                    intent.getStringExtra(ACTIVITY_ID_LABEL)) /* intentPredicate */,
+                                    intent.getStringExtra(KEY_ACTIVITY_ID)
+                            )/* intentPredicate */,
                     mParentWindowMetricsPredicate)
                     .setSplitRatio(DEFAULT_SPLIT_RATIO);
 
diff --git a/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ActivityEmbeddingPolicyTests.java b/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ActivityEmbeddingPolicyTests.java
index a15f668c..5956402 100644
--- a/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ActivityEmbeddingPolicyTests.java
+++ b/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ActivityEmbeddingPolicyTests.java
@@ -16,6 +16,8 @@
 
 package android.server.wm.jetpack;
 
+import static android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN;
+import static android.content.Intent.FLAG_ACTIVITY_NEW_TASK;
 import static android.server.wm.jetpack.second.Components.SECOND_UNTRUSTED_EMBEDDING_ACTIVITY;
 import static android.server.wm.jetpack.signed.Components.SIGNED_EMBEDDING_ACTIVITY;
 import static android.server.wm.jetpack.utils.ActivityEmbeddingUtil.createWildcardSplitPairRule;
@@ -23,7 +25,6 @@
 import static android.server.wm.jetpack.utils.ExtensionUtil.getWindowExtensions;
 import static android.server.wm.jetpack.utils.WindowManagerJetpackTestBase.EXTRA_EMBED_ACTIVITY;
 import static android.server.wm.jetpack.utils.WindowManagerJetpackTestBase.startActivityFromActivity;
-import static android.server.wm.jetpack.utils.WindowManagerJetpackTestBase.startActivityNewTask;
 import static android.server.wm.jetpack.utils.WindowManagerJetpackTestBase.startActivityOnDisplaySingleTop;
 import static android.view.Display.DEFAULT_DISPLAY;
 
@@ -43,6 +44,7 @@
 import android.server.wm.NestedShellPermission;
 import android.server.wm.WindowManagerState;
 import android.server.wm.jetpack.utils.TestActivityKnownEmbeddingCerts;
+import android.server.wm.jetpack.utils.TestActivityLauncher;
 import android.server.wm.jetpack.utils.TestConfigChangeHandlingActivity;
 
 import androidx.annotation.NonNull;
@@ -105,9 +107,11 @@
         // to app.
         assumeFalse(ENABLE_SHELL_TRANSITIONS);
 
-        Activity primaryActivity = startActivityNewTask(mContext, mInstrumentation,
-                TestConfigChangeHandlingActivity.class, null /* activityId */,
-                true /* isFullScreen */);
+        Activity primaryActivity = new TestActivityLauncher<>(mContext,
+                TestConfigChangeHandlingActivity.class)
+                .addIntentFlag(FLAG_ACTIVITY_NEW_TASK)
+                .setWindowingMode(WINDOWING_MODE_FULLSCREEN)
+                .launch(mInstrumentation);
 
         SplitPairRule splitPairRule = createWildcardSplitPairRule(true /* shouldClearTop */);
         mActivityEmbeddingComponent.setEmbeddingRules(Collections.singleton(splitPairRule));
diff --git a/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ExtensionRearDisplayTest.java b/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ExtensionRearDisplayTest.java
index 0806cd3a..d479e6e 100644
--- a/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ExtensionRearDisplayTest.java
+++ b/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ExtensionRearDisplayTest.java
@@ -118,7 +118,7 @@
         mDeviceStateManager.registerCallback(Runnable::run, this);
         mWindowAreaComponent.addRearDisplayStatusListener(mStatusListener);
         unlockDeviceIfNeeded();
-        mActivity = (TestRearDisplayActivity) startActivityNewTask(TestRearDisplayActivity.class);
+        mActivity = startActivityNewTask(TestRearDisplayActivity.class);
         waitAndAssert(() -> mWindowAreaStatus != null);
     }
 
diff --git a/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ExtensionWindowLayoutComponentTest.java b/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ExtensionWindowLayoutComponentTest.java
index ef3ffb1..337749c 100644
--- a/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ExtensionWindowLayoutComponentTest.java
+++ b/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/ExtensionWindowLayoutComponentTest.java
@@ -108,7 +108,7 @@
         mWindowLayoutComponent =
                 (WindowLayoutComponent) mWindowExtensionTestRule.getExtensionComponent();
         assumeNotNull(mWindowLayoutComponent);
-        mActivity = (TestActivity) startActivityNewTask(TestActivity.class);
+        mActivity = startActivityNewTask(TestActivity.class);
     }
 
     private Context createContextWithNonActivityWindow() {
@@ -249,8 +249,7 @@
 
         WindowLayoutInfo windowLayoutInfoFromContext = getExtensionWindowLayoutInfo(windowContext);
 
-        TestConfigChangeHandlingActivity configHandlingActivity =
-                (TestConfigChangeHandlingActivity) startFullScreenActivityNewTask(
+        TestConfigChangeHandlingActivity configHandlingActivity = startFullScreenActivityNewTask(
                         TestConfigChangeHandlingActivity.class, null);
         WindowLayoutInfo windowLayoutInfoFromActivity = getExtensionWindowLayoutInfo(
                 configHandlingActivity);
@@ -266,8 +265,7 @@
         mWindowLayoutInfo = getExtensionWindowLayoutInfo(mActivity);
         assumeHasDisplayFeatures(mWindowLayoutInfo);
 
-        TestConfigChangeHandlingActivity configHandlingActivity =
-                (TestConfigChangeHandlingActivity) startFullScreenActivityNewTask(
+        TestConfigChangeHandlingActivity configHandlingActivity = startFullScreenActivityNewTask(
                         TestConfigChangeHandlingActivity.class, null);
 
         setActivityOrientationActivityHandlesOrientationChanges(configHandlingActivity,
@@ -290,6 +288,28 @@
                 portraitBounds, landscapeBounds, doesDisplayRotateForOrientation);
     }
 
+    @ApiTest(apis = {"androidx.window.extensions.layout.WindowLayoutInfo#getDisplayFeatures"})
+    @Test
+    public void testGetWindowLayoutInfo_enterExitPip_windowLayoutInfoMatches()
+            throws InterruptedException {
+        mWindowLayoutInfo = getExtensionWindowLayoutInfo(mActivity);
+        assumeHasDisplayFeatures(mWindowLayoutInfo);
+
+        TestConfigChangeHandlingActivity configHandlingActivity = startActivityNewTask(
+                        TestConfigChangeHandlingActivity.class, null);
+
+        final WindowLayoutInfo initialInfo = getExtensionWindowLayoutInfo(
+                configHandlingActivity);
+
+        enterPipActivityHandlesConfigChanges(configHandlingActivity);
+        exitPipActivityHandlesConfigChanges(configHandlingActivity);
+
+        final WindowLayoutInfo updatedInfo = getExtensionWindowLayoutInfo(
+                configHandlingActivity);
+
+        assertEquals(initialInfo, updatedInfo);
+    }
+
     /*
      * Similar to #testGetWindowLayoutInfo_configChanged_windowLayoutUpdates, here we trigger
      * rotations with a full screen activity on one Display Area, verify that WindowLayoutInfo
@@ -308,8 +328,7 @@
 
         Context windowContext = createContextWithNonActivityWindow();
 
-        TestConfigChangeHandlingActivity configHandlingActivity =
-                (TestConfigChangeHandlingActivity) startFullScreenActivityNewTask(
+        TestConfigChangeHandlingActivity configHandlingActivity = startFullScreenActivityNewTask(
                         TestConfigChangeHandlingActivity.class, null);
 
         setActivityOrientationActivityHandlesOrientationChanges(configHandlingActivity,
diff --git a/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/SidecarTest.java b/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/SidecarTest.java
index ddd63a7..f30a7d4 100644
--- a/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/SidecarTest.java
+++ b/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/SidecarTest.java
@@ -174,8 +174,7 @@
 
         // The value is verified inside TestGetWindowLayoutInfoActivity
         TestGetWindowLayoutInfoActivity.resetResumeCounter();
-        TestGetWindowLayoutInfoActivity testGetWindowLayoutInfoActivity
-                = (TestGetWindowLayoutInfoActivity) startActivityNewTask(
+        TestGetWindowLayoutInfoActivity testGetWindowLayoutInfoActivity = startActivityNewTask(
                         TestGetWindowLayoutInfoActivity.class);
 
         // Make sure the activity has gone through all states.
@@ -187,8 +186,7 @@
     public void testGetWindowLayoutInfo_configChanged_windowLayoutUpdates() {
         assumeHasDisplayFeatures(mSidecarInterface, mWindowToken);
 
-        TestConfigChangeHandlingActivity configHandlingActivity
-                = (TestConfigChangeHandlingActivity) startActivityNewTask(
+        TestConfigChangeHandlingActivity configHandlingActivity = startActivityNewTask(
                 TestConfigChangeHandlingActivity.class);
         SidecarInterface sidecar = getSidecarInterface(configHandlingActivity);
         assertThat(sidecar).isNotNull();
diff --git a/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/utils/TestActivityLauncher.java b/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/utils/TestActivityLauncher.java
new file mode 100644
index 0000000..ee2fcba
--- /dev/null
+++ b/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/utils/TestActivityLauncher.java
@@ -0,0 +1,72 @@
+/*
+ * Copyright (C) 2022 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 android.server.wm.jetpack.utils;
+
+import android.app.Activity;
+import android.app.ActivityOptions;
+import android.app.Instrumentation;
+import android.content.Context;
+import android.content.Intent;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
+public class TestActivityLauncher<T extends Activity> {
+
+    /** Key for string extra, ID to track an Activity that is launched. */
+    public static final String KEY_ACTIVITY_ID = "ActivityID";
+    /**
+     * Options that will be passed to the instrumentation.
+     * @see TestActivityLauncher#launch(Instrumentation)
+     */
+    private final ActivityOptions mOptions = ActivityOptions.makeBasic();
+
+    /**
+     * The class for the {@link Activity} that you are launching.
+     */
+    private final Class<T> mActivityClass;
+
+    /**
+     * The intent that will be used to launch the {@link Activity}.
+     */
+    private final Intent mIntent;
+
+    public TestActivityLauncher(@NonNull Context context, @NonNull Class<T> activityClass) {
+        mActivityClass = activityClass;
+        mIntent = new Intent(context, activityClass);
+    }
+
+    public TestActivityLauncher<T> addIntentFlag(int flag) {
+        mIntent.addFlags(flag);
+        return this;
+    }
+
+    public TestActivityLauncher<T> setActivityId(@Nullable String id) {
+        mIntent.putExtra(KEY_ACTIVITY_ID, id);
+        return this;
+    }
+
+    public TestActivityLauncher<T> setWindowingMode(int windowingMode) {
+        mOptions.setLaunchWindowingMode(windowingMode);
+        return this;
+    }
+
+    public T launch(@NonNull Instrumentation instrumentation) {
+        return mActivityClass.cast(instrumentation.startActivitySync(mIntent, mOptions.toBundle()));
+    }
+
+}
diff --git a/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/utils/TestActivityWithId.java b/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/utils/TestActivityWithId.java
index 1987020..455db39 100644
--- a/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/utils/TestActivityWithId.java
+++ b/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/utils/TestActivityWithId.java
@@ -16,7 +16,7 @@
 
 package android.server.wm.jetpack.utils;
 
-import static android.server.wm.jetpack.utils.WindowManagerJetpackTestBase.ACTIVITY_ID_LABEL;
+import static android.server.wm.jetpack.utils.TestActivityLauncher.KEY_ACTIVITY_ID;
 
 import android.content.Intent;
 import android.os.Bundle;
@@ -37,8 +37,8 @@
 
         // Get ID
         Intent intent = getIntent();
-        if (intent != null && intent.hasExtra(ACTIVITY_ID_LABEL)) {
-            mId = intent.getStringExtra(ACTIVITY_ID_LABEL);
+        if (intent != null && intent.hasExtra(KEY_ACTIVITY_ID)) {
+            mId = intent.getStringExtra(KEY_ACTIVITY_ID);
         }
     }
 
diff --git a/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/utils/WindowManagerJetpackTestBase.java b/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/utils/WindowManagerJetpackTestBase.java
index 76492ea..5eb33bb 100644
--- a/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/utils/WindowManagerJetpackTestBase.java
+++ b/tests/framework/base/windowmanager/jetpack/src/android/server/wm/jetpack/utils/WindowManagerJetpackTestBase.java
@@ -17,6 +17,7 @@
 package android.server.wm.jetpack.utils;
 
 import static android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN;
+import static android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED;
 import static android.content.Intent.FLAG_ACTIVITY_NEW_TASK;
 import static android.content.Intent.FLAG_ACTIVITY_SINGLE_TOP;
 import static android.content.pm.ActivityInfo.SCREEN_ORIENTATION_LANDSCAPE;
@@ -25,6 +26,7 @@
 import static android.content.pm.PackageManager.FEATURE_SCREEN_PORTRAIT;
 import static android.content.res.Configuration.ORIENTATION_LANDSCAPE;
 import static android.content.res.Configuration.ORIENTATION_PORTRAIT;
+import static android.server.wm.jetpack.utils.TestActivityLauncher.KEY_ACTIVITY_ID;
 
 import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
 import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
@@ -39,6 +41,7 @@
 import android.app.ActivityOptions;
 import android.app.Application;
 import android.app.Instrumentation;
+import android.app.PictureInPictureParams;
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
@@ -61,7 +64,6 @@
 /** Base class for all tests in the module. */
 public class WindowManagerJetpackTestBase {
 
-    public static final String ACTIVITY_ID_LABEL = "ActivityID";
     public static final String EXTRA_EMBED_ACTIVITY = "EmbedActivity";
     public static final String EXTRA_SPLIT_RATIO = "SplitRatio";
 
@@ -105,36 +107,30 @@
                 || (!supportsLandscape && !supportsPortrait));
     }
 
-
-    public Activity startActivityNewTask(@NonNull Class activityClass) {
+    public <T extends Activity> T startActivityNewTask(@NonNull Class<T> activityClass) {
         return startActivityNewTask(activityClass, null /* activityId */);
     }
 
-    public Activity startActivityNewTask(@NonNull Class activityClass,
+    public <T extends Activity> T startActivityNewTask(@NonNull Class<T> activityClass,
             @Nullable String activityId) {
-        return startActivityNewTask(mContext, mInstrumentation, activityClass, activityId,
-                false /* isFullScreen */);
+        return launcherForActivityNewTask(activityClass, activityId, false /* isFullScreen */)
+                .launch(mInstrumentation);
     }
 
-    public Activity startFullScreenActivityNewTask(@NonNull Class activityClass,
+    public <T extends  Activity> T startFullScreenActivityNewTask(@NonNull Class<T> activityClass,
             @Nullable String activityId) {
-        return startActivityNewTask(mContext, mInstrumentation, activityClass, activityId,
-                true/* isFullScreen */);
+        return launcherForActivityNewTask(activityClass, activityId, true/* isFullScreen */)
+                .launch(mInstrumentation);
     }
 
-    public static Activity startActivityNewTask(@NonNull Context context,
-            @NonNull Instrumentation instrumentation, @NonNull Class activityClass,
-            @Nullable String activityId, boolean isFullScreen) {
-        final Intent intent = new Intent(context, activityClass);
-        intent.addFlags(FLAG_ACTIVITY_NEW_TASK);
-        if (activityId != null) {
-            intent.putExtra(ACTIVITY_ID_LABEL, activityId);
-        }
-        final ActivityOptions options = ActivityOptions.makeBasic();
-        if (isFullScreen) {
-            options.setLaunchWindowingMode(WINDOWING_MODE_FULLSCREEN);
-        }
-        return instrumentation.startActivitySync(intent, options.toBundle());
+    private <T extends Activity> TestActivityLauncher<T> launcherForActivityNewTask(
+            @NonNull Class<T> activityClass, @Nullable String activityId, boolean isFullScreen) {
+        final int windowingMode = isFullScreen ? WINDOWING_MODE_FULLSCREEN :
+                WINDOWING_MODE_UNDEFINED;
+        return new TestActivityLauncher<>(mContext, activityClass)
+                .addIntentFlag(FLAG_ACTIVITY_NEW_TASK)
+                .setActivityId(activityId)
+                .setWindowingMode(windowingMode);
     }
 
     /**
@@ -172,7 +168,7 @@
     public static <T extends Activity> void startActivityFromActivity(Activity activityToLaunchFrom,
             Class<T> activityToLaunchClass, String newActivityId) {
         Intent intent = new Intent(activityToLaunchFrom, activityToLaunchClass);
-        intent.putExtra(ACTIVITY_ID_LABEL, newActivityId);
+        intent.putExtra(KEY_ACTIVITY_ID, newActivityId);
         activityToLaunchFrom.startActivity(intent);
     }
 
@@ -185,7 +181,7 @@
         Intent intent = new Intent();
         intent.setClassName(activityToLaunchComponent.getPackageName(),
                 activityToLaunchComponent.getClassName());
-        intent.putExtra(ACTIVITY_ID_LABEL, newActivityId);
+        intent.putExtra(KEY_ACTIVITY_ID, newActivityId);
         intent.putExtras(extras);
         activityToLaunchFrom.startActivity(intent);
     }
@@ -241,6 +237,34 @@
         assertEquals(orientation, activity.getResources().getConfiguration().orientation);
     }
 
+    public static void enterPipActivityHandlesConfigChanges(TestActivity activity) {
+        if (activity.isInPictureInPictureMode()) {
+            throw new IllegalStateException("Activity must not be in PiP");
+        }
+        activity.resetLayoutCounter();
+        // Change the orientation
+        PictureInPictureParams params = (new PictureInPictureParams.Builder()).build();
+        activity.enterPictureInPictureMode(params);
+        // Wait for the activity to layout, which will happen after the orientation change
+        assertTrue(activity.waitForLayout());
+        // Check that orientation matches
+        assertTrue(activity.isInPictureInPictureMode());
+    }
+
+    public static void exitPipActivityHandlesConfigChanges(TestActivity activity) {
+        if (!activity.isInPictureInPictureMode()) {
+            throw new IllegalStateException("Activity must be in PiP");
+        }
+        activity.resetLayoutCounter();
+        Intent intent = new Intent(activity, activity.getClass());
+        intent.addFlags(FLAG_ACTIVITY_SINGLE_TOP);
+        activity.startActivity(intent);
+        // Wait for the activity to layout, which will happen after the orientation change
+        assertTrue(activity.waitForLayout());
+        // Check that orientation matches
+        assertFalse(activity.isInPictureInPictureMode());
+    }
+
     public static void setActivityOrientationActivityDoesNotHandleOrientationChanges(
             TestActivity activity, int orientation) {
         // Make sure that the provided orientation is a fixed orientation