[Partial Screensharing] Add a temporary entrypoint in SysUI screen recorder

Adds a button that allows to start app selection flow in
SystemUI screen recorder dialog (see screenshot in the bug).

This button is needed for teamfood testing of partial
screensharing and will be shown only when the feature flag
is enabled.

It will be removed when redesigned dialog with a dropdown
will be implemented.

Bug: 239422418
Test: turn on feature flag and start partial screen recording
Test: atest com.android.systemui.screenrecord.ScreenRecordDialogTest
Change-Id: I9caa913ef52eeb2203dca32e94e0b10d71fc4b83
diff --git a/packages/SystemUI/res/layout/screen_record_dialog.xml b/packages/SystemUI/res/layout/screen_record_dialog.xml
index f57d65a..ab38dd2 100644
--- a/packages/SystemUI/res/layout/screen_record_dialog.xml
+++ b/packages/SystemUI/res/layout/screen_record_dialog.xml
@@ -148,6 +148,20 @@
                     android:layout_width="0dp"
                     android:layout_height="match_parent"
                     android:layout_weight="1"/>
+
+                <!-- Temporary entrypoint for the partial screensharing used for teamfooding -->
+                <!-- TODO(b/236838395) remove this and use redesigned dialog -->
+                <TextView
+                    android:id="@+id/button_app"
+                    android:visibility="gone"
+                    android:layout_width="wrap_content"
+                    android:layout_height="wrap_content"
+                    android:layout_weight="0"
+                    android:layout_gravity="end"
+                    android:layout_marginEnd="8dp"
+                    android:text="App"
+                    style="@style/Widget.Dialog.Button.BorderButton" />
+
                 <TextView
                     android:id="@+id/button_start"
                     android:layout_width="wrap_content"
diff --git a/packages/SystemUI/src/com/android/systemui/media/MediaProjectionAppSelectorActivity.kt b/packages/SystemUI/src/com/android/systemui/media/MediaProjectionAppSelectorActivity.kt
index 6802da3..53abd99 100644
--- a/packages/SystemUI/src/com/android/systemui/media/MediaProjectionAppSelectorActivity.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/MediaProjectionAppSelectorActivity.kt
@@ -22,6 +22,7 @@
 import android.os.Binder
 import android.os.Bundle
 import android.os.IBinder
+import android.os.ResultReceiver
 import android.view.View
 import com.android.internal.app.ChooserActivity
 import com.android.internal.app.chooser.NotSelectableTargetInfo
@@ -103,17 +104,42 @@
     }
 
     private fun onTargetActivityLaunched(launchToken: IBinder) {
-        val mediaProjectionBinder = intent.getIBinderExtra(EXTRA_MEDIA_PROJECTION)
-        val projection = IMediaProjection.Stub.asInterface(mediaProjectionBinder)
+        if (intent.hasExtra(EXTRA_CAPTURE_REGION_RESULT_RECEIVER)) {
+            // The client requested to return the result in the result receiver instead of
+            // activity result, let's send the media projection to the result receiver
+            val resultReceiver = intent
+                .getParcelableExtra(EXTRA_CAPTURE_REGION_RESULT_RECEIVER,
+                    ResultReceiver::class.java) as ResultReceiver
+            val captureRegion = MediaProjectionCaptureTarget(launchToken)
+            val data = Bundle().apply {
+                putParcelable(KEY_CAPTURE_TARGET, captureRegion)
+            }
+            resultReceiver.send(RESULT_OK, data)
+        } else {
+            // Return the media projection instance as activity result
+            val mediaProjectionBinder = intent.getIBinderExtra(EXTRA_MEDIA_PROJECTION)
+            val projection = IMediaProjection.Stub.asInterface(mediaProjectionBinder)
 
-        projection.launchCookie = launchToken
+            projection.launchCookie = launchToken
 
-        val intent = Intent()
-        intent.putExtra(EXTRA_MEDIA_PROJECTION, projection.asBinder())
-        setResult(RESULT_OK, intent)
-        setForceSendResultForMediaProjection()
+            val intent = Intent()
+            intent.putExtra(EXTRA_MEDIA_PROJECTION, projection.asBinder())
+            setResult(RESULT_OK, intent)
+            setForceSendResultForMediaProjection()
+        }
+
         finish()
     }
 
     override fun shouldGetOnlyDefaultActivities() = false
+
+    companion object {
+        /**
+         * When EXTRA_CAPTURE_REGION_RESULT_RECEIVER is passed as intent extra
+         * the activity will send the [CaptureRegion] to the result receiver
+         * instead of returning media projection instance through activity result.
+         */
+        const val EXTRA_CAPTURE_REGION_RESULT_RECEIVER = "capture_region_result_receiver"
+        const val KEY_CAPTURE_TARGET = "capture_region"
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/media/MediaProjectionCaptureTarget.kt b/packages/SystemUI/src/com/android/systemui/media/MediaProjectionCaptureTarget.kt
new file mode 100644
index 0000000..fbf9294
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/media/MediaProjectionCaptureTarget.kt
@@ -0,0 +1,49 @@
+/*
+ * 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 com.android.systemui.media
+
+import android.os.IBinder
+import android.os.Parcel
+import android.os.Parcelable
+
+/**
+ * Class that represents an area that should be captured.
+ * Currently it has only a launch cookie that represents a task but
+ * we potentially could add more identifiers e.g. for a pair of tasks.
+ */
+data class MediaProjectionCaptureTarget(
+    val launchCookie: IBinder?
+): Parcelable {
+
+    constructor(parcel: Parcel) : this(parcel.readStrongBinder())
+
+    override fun writeToParcel(dest: Parcel, flags: Int) {
+        dest.writeStrongBinder(launchCookie)
+    }
+
+    override fun describeContents(): Int = 0
+
+    companion object CREATOR : Parcelable.Creator<MediaProjectionCaptureTarget> {
+        override fun createFromParcel(parcel: Parcel): MediaProjectionCaptureTarget {
+            return MediaProjectionCaptureTarget(parcel)
+        }
+
+        override fun newArray(size: Int): Array<MediaProjectionCaptureTarget?> {
+            return arrayOfNulls(size)
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/ScreenRecordTile.java b/packages/SystemUI/src/com/android/systemui/qs/tiles/ScreenRecordTile.java
index 02d30c5..75fb393 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/ScreenRecordTile.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/ScreenRecordTile.java
@@ -34,6 +34,7 @@
 import com.android.systemui.animation.DialogLaunchAnimator;
 import com.android.systemui.dagger.qualifiers.Background;
 import com.android.systemui.dagger.qualifiers.Main;
+import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.plugins.ActivityStarter;
 import com.android.systemui.plugins.FalsingManager;
 import com.android.systemui.plugins.qs.QSTile;
@@ -61,6 +62,7 @@
     private final KeyguardStateController mKeyguardStateController;
     private final Callback mCallback = new Callback();
     private final DialogLaunchAnimator mDialogLaunchAnimator;
+    private final FeatureFlags mFlags;
 
     private long mMillisUntilFinished = 0;
 
@@ -71,6 +73,7 @@
             @Main Handler mainHandler,
             FalsingManager falsingManager,
             MetricsLogger metricsLogger,
+            FeatureFlags flags,
             StatusBarStateController statusBarStateController,
             ActivityStarter activityStarter,
             QSLogger qsLogger,
@@ -83,6 +86,7 @@
                 statusBarStateController, activityStarter, qsLogger);
         mController = controller;
         mController.observe(this, mCallback);
+        mFlags = flags;
         mKeyguardDismissUtil = keyguardDismissUtil;
         mKeyguardStateController = keyguardStateController;
         mDialogLaunchAnimator = dialogLaunchAnimator;
@@ -164,8 +168,8 @@
             mDialogLaunchAnimator.disableAllCurrentDialogsExitAnimations();
             getHost().collapsePanels();
         };
-        ScreenRecordDialog dialog = mController.createScreenRecordDialog(mContext,
-                onStartRecordingClicked);
+        ScreenRecordDialog dialog = mController.createScreenRecordDialog(mContext, mFlags,
+                mDialogLaunchAnimator, mActivityStarter, onStartRecordingClicked);
 
         ActivityStarter.OnDismissAction dismissAction = () -> {
             if (shouldAnimateFromView) {
diff --git a/packages/SystemUI/src/com/android/systemui/screenrecord/RecordingController.java b/packages/SystemUI/src/com/android/systemui/screenrecord/RecordingController.java
index 1a08878..1083f22 100644
--- a/packages/SystemUI/src/com/android/systemui/screenrecord/RecordingController.java
+++ b/packages/SystemUI/src/com/android/systemui/screenrecord/RecordingController.java
@@ -29,8 +29,11 @@
 import androidx.annotation.Nullable;
 
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.systemui.animation.DialogLaunchAnimator;
 import com.android.systemui.broadcast.BroadcastDispatcher;
 import com.android.systemui.dagger.SysUISingleton;
+import com.android.systemui.flags.FeatureFlags;
+import com.android.systemui.plugins.ActivityStarter;
 import com.android.systemui.settings.UserContextProvider;
 import com.android.systemui.statusbar.policy.CallbackController;
 
@@ -94,9 +97,11 @@
     }
 
     /** Create a dialog to show screen recording options to the user. */
-    public ScreenRecordDialog createScreenRecordDialog(Context context,
+    public ScreenRecordDialog createScreenRecordDialog(Context context, FeatureFlags flags,
+            DialogLaunchAnimator dialogLaunchAnimator, ActivityStarter activityStarter,
             @Nullable Runnable onStartRecordingClicked) {
-        return new ScreenRecordDialog(context, this, mUserContextProvider, onStartRecordingClicked);
+        return new ScreenRecordDialog(context, this, activityStarter, mUserContextProvider,
+                flags, dialogLaunchAnimator, onStartRecordingClicked);
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/screenrecord/RecordingService.java b/packages/SystemUI/src/com/android/systemui/screenrecord/RecordingService.java
index a837cbb..0477626 100644
--- a/packages/SystemUI/src/com/android/systemui/screenrecord/RecordingService.java
+++ b/packages/SystemUI/src/com/android/systemui/screenrecord/RecordingService.java
@@ -16,6 +16,7 @@
 
 package com.android.systemui.screenrecord;
 
+import android.annotation.Nullable;
 import android.app.Notification;
 import android.app.NotificationChannel;
 import android.app.NotificationManager;
@@ -42,6 +43,7 @@
 import com.android.systemui.R;
 import com.android.systemui.dagger.qualifiers.LongRunning;
 import com.android.systemui.dagger.qualifiers.Main;
+import com.android.systemui.media.MediaProjectionCaptureTarget;
 import com.android.systemui.screenrecord.ScreenMediaRecorder.ScreenMediaRecorderListener;
 import com.android.systemui.settings.UserContextProvider;
 import com.android.systemui.statusbar.phone.KeyguardDismissUtil;
@@ -67,6 +69,7 @@
     private static final String EXTRA_PATH = "extra_path";
     private static final String EXTRA_AUDIO_SOURCE = "extra_useAudio";
     private static final String EXTRA_SHOW_TAPS = "extra_showTaps";
+    private static final String EXTRA_CAPTURE_TARGET = "extra_captureTarget";
 
     private static final String ACTION_START = "com.android.systemui.screenrecord.START";
     private static final String ACTION_STOP = "com.android.systemui.screenrecord.STOP";
@@ -110,14 +113,18 @@
      * @param audioSource   The ordinal value of the audio source
      *                      {@link com.android.systemui.screenrecord.ScreenRecordingAudioSource}
      * @param showTaps   True to make touches visible while recording
+     * @param captureTarget   pass this parameter to capture a specific part instead
+     *                        of the full screen
      */
     public static Intent getStartIntent(Context context, int resultCode,
-            int audioSource, boolean showTaps) {
+            int audioSource, boolean showTaps,
+            @Nullable MediaProjectionCaptureTarget captureTarget) {
         return new Intent(context, RecordingService.class)
                 .setAction(ACTION_START)
                 .putExtra(EXTRA_RESULT_CODE, resultCode)
                 .putExtra(EXTRA_AUDIO_SOURCE, audioSource)
-                .putExtra(EXTRA_SHOW_TAPS, showTaps);
+                .putExtra(EXTRA_SHOW_TAPS, showTaps)
+                .putExtra(EXTRA_CAPTURE_TARGET, captureTarget);
     }
 
     @Override
@@ -136,6 +143,9 @@
                         .values()[intent.getIntExtra(EXTRA_AUDIO_SOURCE, 0)];
                 Log.d(TAG, "recording with audio source" + mAudioSource);
                 mShowTaps = intent.getBooleanExtra(EXTRA_SHOW_TAPS, false);
+                MediaProjectionCaptureTarget captureTarget =
+                        intent.getParcelableExtra(EXTRA_CAPTURE_TARGET,
+                                MediaProjectionCaptureTarget.class);
 
                 mOriginalShowTaps = Settings.System.getInt(
                         getApplicationContext().getContentResolver(),
@@ -148,6 +158,7 @@
                         mMainHandler,
                         currentUserId,
                         mAudioSource,
+                        captureTarget,
                         this
                 );
 
diff --git a/packages/SystemUI/src/com/android/systemui/screenrecord/ScreenMediaRecorder.java b/packages/SystemUI/src/com/android/systemui/screenrecord/ScreenMediaRecorder.java
index d098b4b..b8d96f7 100644
--- a/packages/SystemUI/src/com/android/systemui/screenrecord/ScreenMediaRecorder.java
+++ b/packages/SystemUI/src/com/android/systemui/screenrecord/ScreenMediaRecorder.java
@@ -51,6 +51,7 @@
 import android.view.Surface;
 import android.view.WindowManager;
 
+import com.android.systemui.media.MediaProjectionCaptureTarget;
 import java.io.File;
 import java.io.Closeable;
 import java.io.IOException;
@@ -85,6 +86,7 @@
     private ScreenRecordingMuxer mMuxer;
     private ScreenInternalAudioRecorder mAudio;
     private ScreenRecordingAudioSource mAudioSource;
+    private final MediaProjectionCaptureTarget mCaptureRegion;
     private final Handler mHandler;
 
     private Context mContext;
@@ -92,10 +94,12 @@
 
     public ScreenMediaRecorder(Context context, Handler handler,
             int user, ScreenRecordingAudioSource audioSource,
+            MediaProjectionCaptureTarget captureRegion,
             ScreenMediaRecorderListener listener) {
         mContext = context;
         mHandler = handler;
         mUser = user;
+        mCaptureRegion = captureRegion;
         mListener = listener;
         mAudioSource = audioSource;
     }
@@ -108,9 +112,11 @@
         IMediaProjection proj = null;
         proj = mediaService.createProjection(mUser, mContext.getPackageName(),
                     MediaProjectionManager.TYPE_SCREEN_CAPTURE, false);
-        IBinder projection = proj.asBinder();
-        mMediaProjection = new MediaProjection(mContext,
-                IMediaProjection.Stub.asInterface(projection));
+        IMediaProjection projection = IMediaProjection.Stub.asInterface(proj.asBinder());
+        if (mCaptureRegion != null) {
+            projection.setLaunchCookie(mCaptureRegion.getLaunchCookie());
+        }
+        mMediaProjection = new MediaProjection(mContext, projection);
         mMediaProjection.registerCallback(this, mHandler);
 
         File cacheDir = mContext.getCacheDir();
diff --git a/packages/SystemUI/src/com/android/systemui/screenrecord/ScreenRecordDialog.java b/packages/SystemUI/src/com/android/systemui/screenrecord/ScreenRecordDialog.java
index 1fb88df..efa45a4 100644
--- a/packages/SystemUI/src/com/android/systemui/screenrecord/ScreenRecordDialog.java
+++ b/packages/SystemUI/src/com/android/systemui/screenrecord/ScreenRecordDialog.java
@@ -16,6 +16,10 @@
 
 package com.android.systemui.screenrecord;
 
+import static android.app.Activity.RESULT_OK;
+
+import static com.android.systemui.media.MediaProjectionAppSelectorActivity.EXTRA_CAPTURE_REGION_RESULT_RECEIVER;
+import static com.android.systemui.media.MediaProjectionAppSelectorActivity.KEY_CAPTURE_TARGET;
 import static com.android.systemui.screenrecord.ScreenRecordingAudioSource.INTERNAL;
 import static com.android.systemui.screenrecord.ScreenRecordingAudioSource.MIC;
 import static com.android.systemui.screenrecord.ScreenRecordingAudioSource.MIC_AND_INTERNAL;
@@ -24,8 +28,13 @@
 import android.app.Activity;
 import android.app.PendingIntent;
 import android.content.Context;
+import android.content.Intent;
 import android.os.Bundle;
+import android.os.Handler;
+import android.os.Looper;
+import android.os.ResultReceiver;
 import android.view.Gravity;
+import android.view.View;
 import android.view.Window;
 import android.view.WindowManager;
 import android.widget.ArrayAdapter;
@@ -36,6 +45,13 @@
 import androidx.annotation.Nullable;
 
 import com.android.systemui.R;
+import com.android.systemui.animation.ActivityLaunchAnimator;
+import com.android.systemui.animation.DialogLaunchAnimator;
+import com.android.systemui.flags.FeatureFlags;
+import com.android.systemui.flags.Flags;
+import com.android.systemui.media.MediaProjectionAppSelectorActivity;
+import com.android.systemui.media.MediaProjectionCaptureTarget;
+import com.android.systemui.plugins.ActivityStarter;
 import com.android.systemui.settings.UserContextProvider;
 import com.android.systemui.statusbar.phone.SystemUIDialog;
 
@@ -55,15 +71,23 @@
     private final UserContextProvider mUserContextProvider;
     @Nullable
     private final Runnable mOnStartRecordingClicked;
+    private final ActivityStarter mActivityStarter;
+    private final FeatureFlags mFlags;
+    private final DialogLaunchAnimator mDialogLaunchAnimator;
     private Switch mTapsSwitch;
     private Switch mAudioSwitch;
     private Spinner mOptions;
 
     public ScreenRecordDialog(Context context, RecordingController controller,
-            UserContextProvider userContextProvider, @Nullable Runnable onStartRecordingClicked) {
+            ActivityStarter activityStarter, UserContextProvider userContextProvider,
+            FeatureFlags flags, DialogLaunchAnimator dialogLaunchAnimator,
+            @Nullable Runnable onStartRecordingClicked) {
         super(context);
         mController = controller;
         mUserContextProvider = userContextProvider;
+        mActivityStarter = activityStarter;
+        mDialogLaunchAnimator = dialogLaunchAnimator;
+        mFlags = flags;
         mOnStartRecordingClicked = onStartRecordingClicked;
     }
 
@@ -91,10 +115,36 @@
                 mOnStartRecordingClicked.run();
             }
 
-            requestScreenCapture();
+            // Start full-screen recording
+            requestScreenCapture(/* captureTarget= */ null);
             dismiss();
         });
 
+        if (mFlags.isEnabled(Flags.WM_ENABLE_PARTIAL_SCREEN_SHARING)) {
+            TextView appBtn = findViewById(R.id.button_app);
+
+            appBtn.setVisibility(View.VISIBLE);
+            appBtn.setOnClickListener(v -> {
+                Intent intent = new Intent(getContext(), MediaProjectionAppSelectorActivity.class);
+                intent.addFlags(Intent.FLAG_ACTIVITY_NEW_TASK);
+
+                // We can't start activity for result here so we use result receiver to get
+                // the selected target to capture
+                intent.putExtra(EXTRA_CAPTURE_REGION_RESULT_RECEIVER,
+                        new CaptureTargetResultReceiver());
+
+                ActivityLaunchAnimator.Controller animationController =
+                        mDialogLaunchAnimator.createActivityLaunchController(appBtn);
+
+                if (animationController == null) {
+                    dismiss();
+                }
+
+                mActivityStarter.startActivity(intent, /* dismissShade= */ true,
+                        animationController);
+            });
+        }
+
         mAudioSwitch = findViewById(R.id.screenrecord_audio_switch);
         mTapsSwitch = findViewById(R.id.screenrecord_taps_switch);
         mOptions = findViewById(R.id.screen_recording_options);
@@ -108,7 +158,12 @@
         });
     }
 
-    private void requestScreenCapture() {
+    /**
+     * Starts screen capture after some countdown
+     * @param captureTarget target to capture (could be e.g. a task) or
+     *                      null to record the whole screen
+     */
+    private void requestScreenCapture(@Nullable MediaProjectionCaptureTarget captureTarget) {
         Context userContext = mUserContextProvider.getUserContext();
         boolean showTaps = mTapsSwitch.isChecked();
         ScreenRecordingAudioSource audioMode = mAudioSwitch.isChecked()
@@ -118,7 +173,7 @@
                 RecordingService.REQUEST_CODE,
                 RecordingService.getStartIntent(
                         userContext, Activity.RESULT_OK,
-                        audioMode.ordinal(), showTaps),
+                        audioMode.ordinal(), showTaps, captureTarget),
                 PendingIntent.FLAG_UPDATE_CURRENT | PendingIntent.FLAG_IMMUTABLE);
         PendingIntent stopIntent = PendingIntent.getService(userContext,
                 RecordingService.REQUEST_CODE,
@@ -126,4 +181,22 @@
                 PendingIntent.FLAG_UPDATE_CURRENT | PendingIntent.FLAG_IMMUTABLE);
         mController.startCountdown(DELAY_MS, INTERVAL_MS, startIntent, stopIntent);
     }
+
+    private class CaptureTargetResultReceiver extends ResultReceiver {
+
+        CaptureTargetResultReceiver() {
+            super(new Handler(Looper.getMainLooper()));
+        }
+
+        @Override
+        protected void onReceiveResult(int resultCode, Bundle resultData) {
+            if (resultCode == RESULT_OK) {
+                MediaProjectionCaptureTarget captureTarget = resultData
+                        .getParcelable(KEY_CAPTURE_TARGET, MediaProjectionCaptureTarget.class);
+
+                // Start recording of the selected target
+                requestScreenCapture(captureTarget);
+            }
+        }
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/ScreenRecordTileTest.java b/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/ScreenRecordTileTest.java
index e4c5299..e6bd396 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/ScreenRecordTileTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/ScreenRecordTileTest.java
@@ -22,6 +22,7 @@
 
 import static org.junit.Assert.assertFalse;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -38,6 +39,7 @@
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.animation.DialogLaunchAnimator;
 import com.android.systemui.classifier.FalsingManagerFake;
+import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.plugins.ActivityStarter;
 import com.android.systemui.plugins.statusbar.StatusBarStateController;
 import com.android.systemui.qs.QSTileHost;
@@ -73,6 +75,8 @@
     @Mock
     private QSLogger mQSLogger;
     @Mock
+    private FeatureFlags mFeatureFlags;
+    @Mock
     private KeyguardStateController mKeyguardStateController;
     @Mock
     private DialogLaunchAnimator mDialogLaunchAnimator;
@@ -94,6 +98,7 @@
                 new Handler(mTestableLooper.getLooper()),
                 new FalsingManagerFake(),
                 mMetricsLogger,
+                mFeatureFlags,
                 mStatusBarStateController,
                 mActivityStarter,
                 mQSLogger,
@@ -125,7 +130,8 @@
         mTestableLooper.processAllMessages();
 
         ArgumentCaptor<Runnable> onStartRecordingClicked = ArgumentCaptor.forClass(Runnable.class);
-        verify(mController).createScreenRecordDialog(any(), onStartRecordingClicked.capture());
+        verify(mController).createScreenRecordDialog(any(), eq(mFeatureFlags),
+                eq(mDialogLaunchAnimator), eq(mActivityStarter), onStartRecordingClicked.capture());
 
         // When starting the recording, we collapse the shade and disable the dialog animation.
         assertNotNull(onStartRecordingClicked.getValue());
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenrecord/RecordingServiceTest.java b/packages/SystemUI/tests/src/com/android/systemui/screenrecord/RecordingServiceTest.java
index b05d9a3..a1d78cb 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenrecord/RecordingServiceTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenrecord/RecordingServiceTest.java
@@ -30,6 +30,7 @@
 import android.app.Notification;
 import android.app.NotificationManager;
 import android.content.Intent;
+import android.os.Binder;
 import android.os.Handler;
 import android.os.RemoteException;
 import android.testing.AndroidTestingRunner;
@@ -38,6 +39,7 @@
 
 import com.android.internal.logging.UiEventLogger;
 import com.android.systemui.SysuiTestCase;
+import com.android.systemui.media.MediaProjectionCaptureTarget;
 import com.android.systemui.plugins.ActivityStarter;
 import com.android.systemui.settings.UserContextProvider;
 import com.android.systemui.statusbar.phone.KeyguardDismissUtil;
@@ -108,8 +110,17 @@
     }
 
     @Test
-    public void testLogStartRecording() {
-        Intent startIntent = RecordingService.getStartIntent(mContext, 0, 0, false);
+    public void testLogStartFullScreenRecording() {
+        Intent startIntent = RecordingService.getStartIntent(mContext, 0, 0, false, null);
+        mRecordingService.onStartCommand(startIntent, 0, 0);
+
+        verify(mUiEventLogger, times(1)).log(Events.ScreenRecordEvent.SCREEN_RECORD_START);
+    }
+
+    @Test
+    public void testLogStartPartialRecording() {
+        MediaProjectionCaptureTarget target = new MediaProjectionCaptureTarget(new Binder());
+        Intent startIntent = RecordingService.getStartIntent(mContext, 0, 0, false, target);
         mRecordingService.onStartCommand(startIntent, 0, 0);
 
         verify(mUiEventLogger, times(1)).log(Events.ScreenRecordEvent.SCREEN_RECORD_START);
@@ -142,7 +153,7 @@
         // When the screen recording does not start properly
         doThrow(new RuntimeException("fail")).when(mScreenMediaRecorder).start();
 
-        Intent startIntent = RecordingService.getStartIntent(mContext, 0, 0, false);
+        Intent startIntent = RecordingService.getStartIntent(mContext, 0, 0, false, null);
         mRecordingService.onStartCommand(startIntent, 0, 0);
 
         // Then the state is set to not recording
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenrecord/ScreenRecordDialogTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenrecord/ScreenRecordDialogTest.kt
new file mode 100644
index 0000000..03d9444
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenrecord/ScreenRecordDialogTest.kt
@@ -0,0 +1,111 @@
+/*
+ * 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.
+ */
+
+/*
+ * 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 com.android.systemui.screenrecord
+
+import android.testing.AndroidTestingRunner
+import android.testing.TestableLooper
+import android.view.View
+import androidx.test.filters.SmallTest
+import com.android.systemui.R
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.animation.DialogLaunchAnimator
+import com.android.systemui.flags.FeatureFlags
+import com.android.systemui.flags.Flags
+import com.android.systemui.plugins.ActivityStarter
+import com.android.systemui.settings.UserContextProvider
+import com.google.common.truth.Truth.assertThat
+import org.junit.After
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.MockitoAnnotations
+import org.mockito.Mockito.`when` as whenever
+
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+@TestableLooper.RunWithLooper(setAsMainLooper = true)
+class ScreenRecordDialogTest : SysuiTestCase() {
+
+    @Mock
+    private lateinit var starter: ActivityStarter
+    @Mock
+    private lateinit var controller: RecordingController
+    @Mock
+    private lateinit var userContextProvider: UserContextProvider
+    @Mock
+    private lateinit var flags: FeatureFlags
+    @Mock
+    private lateinit var dialogLaunchAnimator: DialogLaunchAnimator
+    @Mock
+    private lateinit var onStartRecordingClicked: Runnable
+
+    private lateinit var dialog: ScreenRecordDialog
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+
+        dialog = ScreenRecordDialog(
+            context, controller, starter, userContextProvider, flags, dialogLaunchAnimator,
+            onStartRecordingClicked
+        )
+    }
+
+    @After
+    fun teardown() {
+        if (::dialog.isInitialized) {
+            dialog.dismiss()
+        }
+    }
+
+    @Test
+    fun testShowDialog_partialScreenSharingDisabled_appButtonIsNotVisible() {
+        whenever(flags.isEnabled(Flags.WM_ENABLE_PARTIAL_SCREEN_SHARING)).thenReturn(false)
+
+        dialog.show()
+
+        val visibility = dialog.requireViewById<View>(R.id.button_app).visibility
+        assertThat(visibility).isEqualTo(View.GONE)
+    }
+
+    @Test
+    fun testShowDialog_partialScreenSharingEnabled_appButtonIsVisible() {
+        whenever(flags.isEnabled(Flags.WM_ENABLE_PARTIAL_SCREEN_SHARING)).thenReturn(true)
+
+        dialog.show()
+
+        val visibility = dialog.requireViewById<View>(R.id.button_app).visibility
+        assertThat(visibility).isEqualTo(View.VISIBLE)
+    }
+}