Merge "Structure for screenshot detection notification." into tm-qpr-dev

packages/SystemUI/res/layout/screenshot_detection_notice.xml

@@ -0,0 +1,16 @@
+<?xml version="1.0" encoding="utf-8"?>
+<FrameLayout
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:id="@+id/screenshot_detection_notice"
+    android:layout_height="wrap_content"
+    android:layout_width="match_parent"
+    android:padding="12dp"
+    android:visibility="gone">
+
+    <TextView
+        android:id="@+id/screenshot_detection_notice_text"
+        android:layout_width="match_parent"
+        android:layout_height="wrap_content"
+        android:lineHeight="24sp"
+        android:textSize="18sp" />
+</FrameLayout>

packages/SystemUI/res/layout/screenshot_static.xml

@@ -134,7 +134,7 @@
         android:orientation="horizontal"
         app:layout_constraintGuide_end="0dp" />
 
-    <androidx.constraintlayout.widget.ConstraintLayout
+    <FrameLayout
         android:id="@+id/screenshot_message_container"
         android:layout_width="0dp"
         android:layout_height="wrap_content"
@@ -145,10 +145,12 @@
         android:paddingVertical="@dimen/overlay_action_container_padding_vertical"
         android:elevation="4dp"
         android:background="@drawable/action_chip_container_background"
-        android:visibility="invisible"
+        android:visibility="gone"
         app:layout_constraintTop_toBottomOf="@id/guideline"
         app:layout_constraintStart_toStartOf="parent"
         app:layout_constraintEnd_toEndOf="parent"
         >
-    </androidx.constraintlayout.widget.ConstraintLayout>
+        <include layout="@layout/screenshot_work_profile_first_run" />
+        <include layout="@layout/screenshot_detection_notice" />
+    </FrameLayout>
 </com.android.systemui.screenshot.DraggableConstraintLayout>

packages/SystemUI/res/layout/screenshot_work_profile_first_run.xml

@@ -1,36 +1,28 @@
 <?xml version="1.0" encoding="utf-8"?>
-<merge
+<LinearLayout
     xmlns:android="http://schemas.android.com/apk/res/android"
-    xmlns:app="http://schemas.android.com/apk/res-auto">
+    android:id="@+id/work_profile_first_run"
+    android:layout_height="wrap_content"
+    android:layout_width="match_parent"
+    android:visibility="gone">
     <ImageView
         android:id="@+id/screenshot_message_icon"
         android:layout_width="48dp"
         android:layout_height="48dp"
         android:paddingEnd="4dp"
-        android:src="@drawable/ic_work_app_badge"
+        android:src="@drawable/ic_work_app_badge"/>
-        app:layout_constraintStart_toStartOf="parent"
-        app:layout_constraintEnd_toStartOf="@id/screenshot_message_content"
-        app:layout_constraintTop_toTopOf="parent"
-        app:layout_constraintBottom_toBottomOf="parent"/>
 
     <TextView
         android:id="@+id/screenshot_message_content"
         android:layout_width="0dp"
         android:layout_height="wrap_content"
-        android:layout_gravity="start"
+        android:layout_weight="1"
-        app:layout_constraintTop_toTopOf="parent"
+        android:layout_gravity="start"/>
-        app:layout_constraintBottom_toBottomOf="parent"
-        app:layout_constraintStart_toEndOf="@id/screenshot_message_icon"
-        app:layout_constraintEnd_toStartOf="@id/message_dismiss_button"/>
 
     <FrameLayout
         android:id="@+id/message_dismiss_button"
         android:layout_width="@dimen/overlay_dismiss_button_tappable_size"
         android:layout_height="@dimen/overlay_dismiss_button_tappable_size"
-        app:layout_constraintStart_toEndOf="@id/screenshot_message_content"
-        app:layout_constraintEnd_toEndOf="parent"
-        app:layout_constraintTop_toTopOf="parent"
-        app:layout_constraintBottom_toBottomOf="parent"
         android:contentDescription="@string/screenshot_dismiss_work_profile">
         <ImageView
             android:layout_width="match_parent"
@@ -38,4 +30,4 @@
             android:layout_margin="@dimen/overlay_dismiss_button_margin"
             android:src="@drawable/overlay_cancel"/>
     </FrameLayout>
-</merge>
+</LinearLayout>

packages/SystemUI/res/values/strings.xml

@@ -243,6 +243,10 @@
     <string name="screenshot_work_profile_notification">Work screenshots are saved in the <xliff:g id="app" example="Work Files">%1$s</xliff:g> app</string>
     <!-- Default name referring to the app on the device that lets the user browse stored files. [CHAR LIMIT=NONE] -->
     <string name="screenshot_default_files_app_name">Files</string>
+    <!-- A notice shown to the user to indicate that an app has detected the screenshot that the user has just taken. [CHAR LIMIT=75] -->
+    <string name="screenshot_detected_template"><xliff:g id="appName" example="Google Chrome">%1$s</xliff:g> detected this screenshot.</string>
+    <!-- A notice shown to the user to indicate that multiple apps have detected the screenshot that the user has just taken. [CHAR LIMIT=75] -->
+    <string name="screenshot_detected_multiple_template"><xliff:g id="appName" example="Google Chrome">%1$s</xliff:g> and other open apps detected this screenshot.</string>
 
     <!-- Notification title displayed for screen recording [CHAR LIMIT=50]-->
     <string name="screenrecord_name">Screen Recorder</string>

packages/SystemUI/src/com/android/systemui/flags/Flags.kt

@@ -512,6 +512,9 @@ object Flags {
     // TODO(b/264916608): Tracking Bug
     @JvmField val SCREENSHOT_METADATA = unreleasedFlag(1302, "screenshot_metadata")
 
+    // TODO(b/266955521): Tracking bug
+    @JvmField val SCREENSHOT_DETECTION = unreleasedFlag(1303, "screenshot_detection")
+
     // 1400 - columbus
     // TODO(b/254512756): Tracking Bug
     val QUICK_TAP_IN_PCC = releasedFlag(1400, "quick_tap_in_pcc")

packages/SystemUI/src/com/android/systemui/screenshot/MessageContainerController.kt

@@ -3,101 +3,133 @@ package com.android.systemui.screenshot
 import android.animation.Animator
 import android.animation.AnimatorListenerAdapter
 import android.animation.ValueAnimator
-import android.graphics.drawable.Drawable
+import android.os.UserHandle
 import android.view.View
 import android.view.ViewGroup
 import android.view.ViewGroup.MarginLayoutParams
 import android.view.ViewTreeObserver
 import android.view.animation.AccelerateDecelerateInterpolator
-import android.widget.ImageView
-import android.widget.TextView
 import androidx.constraintlayout.widget.Guideline
 import com.android.systemui.R
+import com.android.systemui.flags.FeatureFlags
+import com.android.systemui.flags.Flags
+import javax.inject.Inject
 
 /**
  * MessageContainerController controls the display of content in the screenshot message container.
  */
 class MessageContainerController
+@Inject
 constructor(
-    parent: ViewGroup,
+    private val workProfileMessageController: WorkProfileMessageController,
+    private val screenshotDetectionController: ScreenshotDetectionController,
+    private val featureFlags: FeatureFlags,
 ) {
-    private val guideline: Guideline = parent.requireViewById(R.id.guideline)
+    private lateinit var container: ViewGroup
-    private val messageContainer: ViewGroup =
+    private lateinit var guideline: Guideline
-        parent.requireViewById(R.id.screenshot_message_container)
+    private lateinit var workProfileFirstRunView: ViewGroup
+    private lateinit var detectionNoticeView: ViewGroup
+    private var animateOut: Animator? = null
 
-    /**
+    fun setView(screenshotView: ViewGroup) {
-     * Show a notification under the screenshot view indicating that a work profile screenshot has
+        container = screenshotView.requireViewById(R.id.screenshot_message_container)
-     * been taken and which app can be used to view it.
+        guideline = screenshotView.requireViewById(R.id.guideline)
-     *
+
-     * @param appName The name of the app to use to view screenshots
+        workProfileFirstRunView = container.requireViewById(R.id.work_profile_first_run)
-     * @param appIcon Optional icon for the relevant files app
+        detectionNoticeView = container.requireViewById(R.id.screenshot_detection_notice)
-     * @param onDismiss Runnable to be run when the user dismisses this message
+
-     */
+        // Restore to starting state.
-    fun showWorkProfileMessage(appName: CharSequence, appIcon: Drawable?, onDismiss: Runnable) {
+        container.visibility = View.GONE
-        // Eventually this container will support multiple notification types, but for now just make
+        guideline.setGuidelineEnd(0)
-        // sure we don't double inflate.
+        workProfileFirstRunView.visibility = View.GONE
-        if (messageContainer.childCount == 0) {
+        detectionNoticeView.visibility = View.GONE
-            View.inflate(
+    }
-                messageContainer.context,
+
-                R.layout.screenshot_work_profile_first_run,
+    // Minimal implementation for use when Flags.SCREENSHOT_METADATA isn't turned on.
-                messageContainer
+    fun onScreenshotTaken(userHandle: UserHandle) {
-            )
+        if (featureFlags.isEnabled(Flags.SCREENSHOT_WORK_PROFILE_POLICY)) {
+            val workProfileData = workProfileMessageController.onScreenshotTaken(userHandle)
+            if (workProfileData != null) {
+                workProfileFirstRunView.visibility = View.VISIBLE
+                detectionNoticeView.visibility = View.GONE
+
+                workProfileMessageController.populateView(
+                    workProfileFirstRunView,
+                    workProfileData,
+                    this::animateOutMessageContainer
+                )
+                animateInMessageContainer()
+            }
         }
-        if (appIcon != null) {
+    }
-            // Replace the default icon if one is provided.
+
-            val imageView: ImageView =
+    fun onScreenshotTaken(screenshot: ScreenshotData) {
-                messageContainer.requireViewById<ImageView>(R.id.screenshot_message_icon)
+        if (featureFlags.isEnabled(Flags.SCREENSHOT_WORK_PROFILE_POLICY)) {
-            imageView.setImageDrawable(appIcon)
+            val workProfileData =
-        }
+                workProfileMessageController.onScreenshotTaken(screenshot.userHandle)
-        val messageContent =
+            var notifiedApps: List<CharSequence> = listOf()
-            messageContainer.requireViewById<TextView>(R.id.screenshot_message_content)
+            if (featureFlags.isEnabled(Flags.SCREENSHOT_DETECTION)) {
-        messageContent.text =
+                notifiedApps = screenshotDetectionController.maybeNotifyOfScreenshot(screenshot)
-            messageContainer.context.getString(
+            }
-                R.string.screenshot_work_profile_notification,
+
-                appName
+            // If work profile first run needs to show, bias towards that, otherwise show screenshot
-            )
+            // detection notification if needed.
-        messageContainer.requireViewById<View>(R.id.message_dismiss_button).setOnClickListener {
+            if (workProfileData != null) {
-            animateOutMessageContainer()
+                workProfileFirstRunView.visibility = View.VISIBLE
-            onDismiss.run()
+                detectionNoticeView.visibility = View.GONE
+                workProfileMessageController.populateView(
+                    workProfileFirstRunView,
+                    workProfileData,
+                    this::animateOutMessageContainer
+                )
+                animateInMessageContainer()
+            } else if (notifiedApps.isNotEmpty()) {
+                detectionNoticeView.visibility = View.VISIBLE
+                workProfileFirstRunView.visibility = View.GONE
+                screenshotDetectionController.populateView(detectionNoticeView, notifiedApps)
+                animateInMessageContainer()
+            }
         }
+    }
+
+    private fun animateInMessageContainer() {
+        if (container.visibility == View.VISIBLE) return
 
         // Need the container to be fully measured before animating in (to know animation offset
         // destination)
-        messageContainer.viewTreeObserver.addOnPreDrawListener(
+        container.visibility = View.VISIBLE
+        container.viewTreeObserver.addOnPreDrawListener(
             object : ViewTreeObserver.OnPreDrawListener {
                 override fun onPreDraw(): Boolean {
-                    messageContainer.viewTreeObserver.removeOnPreDrawListener(this)
+                    container.viewTreeObserver.removeOnPreDrawListener(this)
-                    animateInMessageContainer()
+                    getAnimator(true).start()
                     return false
                 }
             }
         )
     }
 
-    private fun animateInMessageContainer() {
-        if (messageContainer.visibility == View.VISIBLE) return
-
-        messageContainer.visibility = View.VISIBLE
-        getAnimator(true).start()
-    }
-
     private fun animateOutMessageContainer() {
-        getAnimator(false).apply {
+        if (animateOut != null) return
-            addListener(
+
-                object : AnimatorListenerAdapter() {
+        animateOut =
-                    override fun onAnimationEnd(animation: Animator) {
+            getAnimator(false).apply {
-                        super.onAnimationEnd(animation)
+                addListener(
-                        messageContainer.visibility = View.INVISIBLE
+                    object : AnimatorListenerAdapter() {
+                        override fun onAnimationEnd(animation: Animator) {
+                            super.onAnimationEnd(animation)
+                            container.visibility = View.GONE
+                            animateOut = null
+                        }
                     }
-                }
+                )
-            )
+                start()
-            start()
+            }
-        }
     }
 
     private fun getAnimator(animateIn: Boolean): Animator {
-        val params = messageContainer.layoutParams as MarginLayoutParams
+        val params = container.layoutParams as MarginLayoutParams
-        val offset = messageContainer.height + params.topMargin + params.bottomMargin
+        val offset = container.height + params.topMargin + params.bottomMargin
         val anim = if (animateIn) ValueAnimator.ofFloat(0f, 1f) else ValueAnimator.ofFloat(1f, 0f)
         with(anim) {
             duration = ScreenshotView.SCREENSHOT_ACTIONS_EXPANSION_DURATION_MS
@@ -105,7 +137,7 @@ constructor(
             addUpdateListener { valueAnimator: ValueAnimator ->
                 val interpolation = valueAnimator.animatedValue as Float
                 guideline.setGuidelineEnd((interpolation * offset).toInt())
-                messageContainer.alpha = interpolation
+                container.alpha = interpolation
             }
         }
         return anim

packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java

@@ -282,7 +282,6 @@ public class ScreenshotController {
     private final TimeoutHandler mScreenshotHandler;
     private final ActionIntentExecutor mActionExecutor;
     private final UserManager mUserManager;
-    private final WorkProfileMessageController mWorkProfileMessageController;
     private final AssistContentRequester mAssistContentRequester;
 
     private final OnBackInvokedCallback mOnBackInvokedCallback = () -> {
@@ -293,7 +292,7 @@ public class ScreenshotController {
     };
 
     private ScreenshotView mScreenshotView;
-    private MessageContainerController mMessageContainerController;
+    private final MessageContainerController mMessageContainerController;
     private Bitmap mScreenBitmap;
     private SaveImageInBackgroundTask mSaveInBgTask;
     private boolean mScreenshotTakenInPortrait;
@@ -332,8 +331,8 @@ public class ScreenshotController {
             ScreenshotNotificationSmartActionsProvider screenshotNotificationSmartActionsProvider,
             ActionIntentExecutor actionExecutor,
             UserManager userManager,
-            WorkProfileMessageController workProfileMessageController,
             AssistContentRequester assistContentRequester,
+            MessageContainerController messageContainerController,
             DisplayTracker displayTracker
     ) {
         mScreenshotSmartActions = screenshotSmartActions;
@@ -367,7 +366,7 @@ public class ScreenshotController {
         mFlags = flags;
         mActionExecutor = actionExecutor;
         mUserManager = userManager;
-        mWorkProfileMessageController = workProfileMessageController;
+        mMessageContainerController = messageContainerController;
         mAssistContentRequester = assistContentRequester;
 
         mAccessibilityManager = AccessibilityManager.getInstance(mContext);
@@ -468,7 +467,11 @@ public class ScreenshotController {
             }
         }
 
-        prepareAnimation(screenshot.getScreenBounds(), showFlash);
+        prepareAnimation(screenshot.getScreenBounds(), showFlash, () -> {
+            if (mFlags.isEnabled(SCREENSHOT_WORK_PROFILE_POLICY)) {
+                mMessageContainerController.onScreenshotTaken(screenshot);
+            }
+        });
 
         if (mFlags.isEnabled(SCREENSHOT_WORK_PROFILE_POLICY)) {
             mScreenshotView.badgeScreenshot(mContext.getPackageManager().getUserBadgedIcon(
@@ -632,7 +635,9 @@ public class ScreenshotController {
         // Inflate the screenshot layout
         mScreenshotView = (ScreenshotView)
                 LayoutInflater.from(mContext).inflate(R.layout.screenshot, null);
-        mMessageContainerController = new MessageContainerController(mScreenshotView);
+        if (mFlags.isEnabled(SCREENSHOT_WORK_PROFILE_POLICY)) {
+            mMessageContainerController.setView(mScreenshotView);
+        }
         mScreenshotView.addOnAttachStateChangeListener(
                 new View.OnAttachStateChangeListener() {
                     @Override
@@ -782,7 +787,11 @@ public class ScreenshotController {
         enqueueScrollCaptureRequest(owner);
 
         attachWindow();
-        prepareAnimation(screenRect, showFlash);
+        prepareAnimation(screenRect, showFlash, () -> {
+            if (mFlags.isEnabled(SCREENSHOT_WORK_PROFILE_POLICY)) {
+                mMessageContainerController.onScreenshotTaken(owner);
+            }
+        });
 
         if (mFlags.isEnabled(SCREENSHOT_WORK_PROFILE_POLICY)) {
             mScreenshotView.badgeScreenshot(mContext.getPackageManager().getUserBadgedIcon(
@@ -799,7 +808,8 @@ public class ScreenshotController {
         mScreenshotHandler.cancelTimeout(); // restarted after animation
     }
 
-    private void prepareAnimation(Rect screenRect, boolean showFlash) {
+    private void prepareAnimation(Rect screenRect, boolean showFlash,
+            Runnable onAnimationComplete) {
         mScreenshotView.getViewTreeObserver().addOnPreDrawListener(
                 new ViewTreeObserver.OnPreDrawListener() {
                     @Override
@@ -808,7 +818,7 @@ public class ScreenshotController {
                             Log.d(TAG, "onPreDraw: startAnimation");
                         }
                         mScreenshotView.getViewTreeObserver().removeOnPreDrawListener(this);
-                        startAnimation(screenRect, showFlash);
+                        startAnimation(screenRect, showFlash, onAnimationComplete);
                         return true;
                     }
                 });
@@ -1089,13 +1099,22 @@ public class ScreenshotController {
     /**
      * Starts the animation after taking the screenshot
      */
-    private void startAnimation(Rect screenRect, boolean showFlash) {
+    private void startAnimation(Rect screenRect, boolean showFlash, Runnable onAnimationComplete) {
         if (mScreenshotAnimation != null && mScreenshotAnimation.isRunning()) {
             mScreenshotAnimation.cancel();
         }
 
         mScreenshotAnimation =
                 mScreenshotView.createScreenshotDropInAnimation(screenRect, showFlash);
+        if (onAnimationComplete != null) {
+            mScreenshotAnimation.addListener(new AnimatorListenerAdapter() {
+                @Override
+                public void onAnimationEnd(Animator animation) {
+                    super.onAnimationEnd(animation);
+                    onAnimationComplete.run();
+                }
+            });
+        }
 
         // Play the shutter sound to notify that we've taken a screenshot
         playCameraSound();
@@ -1194,10 +1213,6 @@ public class ScreenshotController {
 
     private void doPostAnimation(ScreenshotController.SavedImageData imageData) {
         mScreenshotView.setChipIntents(imageData);
-        if (mFlags.isEnabled(SCREENSHOT_WORK_PROFILE_POLICY)) {
-            mWorkProfileMessageController.onScreenshotTaken(imageData.owner,
-                    mMessageContainerController);
-        }
     }
 
     /**

packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotData.kt

@@ -8,6 +8,7 @@ import android.net.Uri
 import android.os.UserHandle
 import android.view.WindowManager.ScreenshotSource
 import android.view.WindowManager.ScreenshotType
+import androidx.annotation.VisibleForTesting
 import com.android.internal.util.ScreenshotRequest
 
 /** ScreenshotData represents the current state of a single screenshot being acquired. */
@@ -42,5 +43,10 @@ data class ScreenshotData(
                 request.bitmap,
             )
         }
+
+        @VisibleForTesting
+        fun forTesting(): ScreenshotData {
+            return ScreenshotData(0, 0, null, null, null, 0, Insets.NONE, null)
+        }
     }
 }

packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotDetectionController.kt

@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2023 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.screenshot
+
+import android.content.pm.PackageManager
+import android.view.IWindowManager
+import android.view.ViewGroup
+import android.widget.TextView
+import com.android.systemui.R
+import javax.inject.Inject
+
+class ScreenshotDetectionController
+@Inject
+constructor(
+    private val windowManager: IWindowManager,
+    private val packageManager: PackageManager,
+) {
+    /**
+     * Notify potentially listening apps of the screenshot. Return a list of the names of the apps
+     * notified.
+     */
+    fun maybeNotifyOfScreenshot(data: ScreenshotData): List<CharSequence> {
+        // TODO: actually ask the window manager once API is available.
+        return listOf()
+    }
+
+    fun populateView(view: ViewGroup, appNames: List<CharSequence>) {
+        assert(appNames.isNotEmpty())
+
+        val textView: TextView = view.requireViewById(R.id.screenshot_detection_notice_text)
+        if (appNames.size == 1) {
+            textView.text =
+                view.resources.getString(R.string.screenshot_detected_template, appNames[0])
+        } else {
+            textView.text =
+                view.resources.getString(
+                    R.string.screenshot_detected_multiple_template,
+                    appNames[0]
+                )
+        }
+    }
+}

packages/SystemUI/src/com/android/systemui/screenshot/WorkProfileMessageController.kt

@@ -23,13 +23,16 @@ import android.graphics.drawable.Drawable
 import android.os.UserHandle
 import android.os.UserManager
 import android.util.Log
+import android.view.View
+import android.view.ViewGroup
+import android.widget.ImageView
+import android.widget.TextView
 import com.android.systemui.R
 import javax.inject.Inject
 
 /**
- * Handles all the non-UI portions of the work profile first run:
+ * Handles work profile first run, determining whether a first run UI should be shown and populating
- * - Track whether the user has already dismissed it.
+ * that UI if needed.
- * - Load the proper icon and app name.
  */
 class WorkProfileMessageController
 @Inject
@@ -40,10 +43,12 @@ constructor(
 ) {
 
     /**
-     * Determine if a message should be shown to the user, send message details to
+     * @return a populated WorkProfileFirstRunData object if a work profile first run message should
-     * MessageContainerController if appropriate.
+     * be shown
      */
-    fun onScreenshotTaken(userHandle: UserHandle, messageContainer: MessageContainerController) {
+    fun onScreenshotTaken(userHandle: UserHandle?): WorkProfileFirstRunData? {
+        if (userHandle == null) return null
+
         if (userManager.isManagedProfile(userHandle.identifier) && !messageAlreadyDismissed()) {
             var badgedIcon: Drawable? = null
             var label: CharSequence? = null
@@ -65,9 +70,27 @@ constructor(
             val badgedLabel =
                 packageManager.getUserBadgedLabel(label ?: defaultFileAppName(), userHandle)
 
-            messageContainer.showWorkProfileMessage(badgedLabel, badgedIcon) {
+            return WorkProfileFirstRunData(badgedLabel, badgedIcon)
-                onMessageDismissed()
+        }
-            }
+        return null
+    }
+
+    /**
+     * Use the provided WorkProfileFirstRunData to populate the work profile first run UI in the
+     * given view.
+     */
+    fun populateView(view: ViewGroup, data: WorkProfileFirstRunData, animateOut: () -> Unit) {
+        if (data.icon != null) {
+            // Replace the default icon if one is provided.
+            val imageView: ImageView = view.requireViewById<ImageView>(R.id.screenshot_message_icon)
+            imageView.setImageDrawable(data.icon)
+        }
+        val messageContent = view.requireViewById<TextView>(R.id.screenshot_message_content)
+        messageContent.text =
+            view.context.getString(R.string.screenshot_work_profile_notification, data.appName)
+        view.requireViewById<View>(R.id.message_dismiss_button).setOnClickListener {
+            animateOut()
+            onMessageDismissed()
         }
     }
 
@@ -91,6 +114,8 @@ constructor(
 
     private fun defaultFileAppName() = context.getString(R.string.screenshot_default_files_app_name)
 
+    data class WorkProfileFirstRunData constructor(val appName: CharSequence, val icon: Drawable?)
+
     companion object {
         const val TAG = "WorkProfileMessageCtrl"
         const val SHARED_PREFERENCES_NAME = "com.android.systemui.screenshot"

packages/SystemUI/tests/src/com/android/systemui/screenshot/MessageContainerControllerTest.kt

@@ -0,0 +1,143 @@
+package com.android.systemui.screenshot
+
+import android.graphics.drawable.Drawable
+import android.os.UserHandle
+import android.testing.AndroidTestingRunner
+import android.view.View
+import android.view.ViewGroup
+import android.widget.FrameLayout
+import androidx.constraintlayout.widget.ConstraintLayout
+import androidx.constraintlayout.widget.Guideline
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.flags.FakeFeatureFlags
+import com.android.systemui.flags.Flags
+import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.eq
+import com.android.systemui.util.mockito.whenever
+import junit.framework.Assert.assertEquals
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.Mockito.never
+import org.mockito.Mockito.verify
+import org.mockito.MockitoAnnotations
+
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+class MessageContainerControllerTest : SysuiTestCase() {
+    lateinit var messageContainer: MessageContainerController
+
+    @Mock lateinit var workProfileMessageController: WorkProfileMessageController
+
+    @Mock lateinit var screenshotDetectionController: ScreenshotDetectionController
+
+    @Mock lateinit var icon: Drawable
+
+    lateinit var workProfileFirstRunView: ViewGroup
+    lateinit var detectionNoticeView: ViewGroup
+    lateinit var container: FrameLayout
+
+    var featureFlags = FakeFeatureFlags()
+    lateinit var screenshotView: ViewGroup
+
+    val userHandle = UserHandle.of(5)
+    val screenshotData = ScreenshotData.forTesting()
+
+    val appName = "app name"
+    lateinit var workProfileData: WorkProfileMessageController.WorkProfileFirstRunData
+
+    @Before
+    fun setup() {
+        MockitoAnnotations.initMocks(this)
+        messageContainer =
+            MessageContainerController(
+                workProfileMessageController,
+                screenshotDetectionController,
+                featureFlags
+            )
+        screenshotView = ConstraintLayout(mContext)
+        workProfileData = WorkProfileMessageController.WorkProfileFirstRunData(appName, icon)
+
+        val guideline = Guideline(mContext)
+        guideline.id = com.android.systemui.R.id.guideline
+        screenshotView.addView(guideline)
+
+        container = FrameLayout(mContext)
+        container.id = com.android.systemui.R.id.screenshot_message_container
+        screenshotView.addView(container)
+
+        workProfileFirstRunView = FrameLayout(mContext)
+        workProfileFirstRunView.id = com.android.systemui.R.id.work_profile_first_run
+        container.addView(workProfileFirstRunView)
+
+        detectionNoticeView = FrameLayout(mContext)
+        detectionNoticeView.id = com.android.systemui.R.id.screenshot_detection_notice
+        container.addView(detectionNoticeView)
+
+        messageContainer.setView(screenshotView)
+
+        screenshotData.userHandle = userHandle
+    }
+
+    @Test
+    fun testOnScreenshotTakenUserHandle_noWorkProfileFirstRun() {
+        featureFlags.set(Flags.SCREENSHOT_WORK_PROFILE_POLICY, true)
+        // (just being explicit here)
+        whenever(workProfileMessageController.onScreenshotTaken(eq(userHandle))).thenReturn(null)
+
+        messageContainer.onScreenshotTaken(userHandle)
+
+        verify(workProfileMessageController, never()).populateView(any(), any(), any())
+    }
+
+    @Test
+    fun testOnScreenshotTakenUserHandle_noWorkProfileFlag() {
+        featureFlags.set(Flags.SCREENSHOT_WORK_PROFILE_POLICY, false)
+
+        messageContainer.onScreenshotTaken(userHandle)
+
+        verify(workProfileMessageController, never()).onScreenshotTaken(any())
+        verify(workProfileMessageController, never()).populateView(any(), any(), any())
+    }
+
+    @Test
+    fun testOnScreenshotTakenUserHandle_withWorkProfileFirstRun() {
+        featureFlags.set(Flags.SCREENSHOT_WORK_PROFILE_POLICY, true)
+        whenever(workProfileMessageController.onScreenshotTaken(eq(userHandle)))
+            .thenReturn(workProfileData)
+        messageContainer.onScreenshotTaken(userHandle)
+
+        verify(workProfileMessageController)
+            .populateView(eq(workProfileFirstRunView), eq(workProfileData), any())
+        assertEquals(View.VISIBLE, workProfileFirstRunView.visibility)
+        assertEquals(View.GONE, detectionNoticeView.visibility)
+    }
+
+    @Test
+    fun testOnScreenshotTakenScreenshotData_flagsOff() {
+        featureFlags.set(Flags.SCREENSHOT_WORK_PROFILE_POLICY, false)
+        featureFlags.set(Flags.SCREENSHOT_DETECTION, false)
+
+        messageContainer.onScreenshotTaken(screenshotData)
+
+        verify(workProfileMessageController, never()).onScreenshotTaken(any())
+        verify(screenshotDetectionController, never()).maybeNotifyOfScreenshot(any())
+
+        assertEquals(View.GONE, container.visibility)
+    }
+
+    @Test
+    fun testOnScreenshotTakenScreenshotData_nothingToShow() {
+        featureFlags.set(Flags.SCREENSHOT_WORK_PROFILE_POLICY, true)
+        featureFlags.set(Flags.SCREENSHOT_DETECTION, true)
+
+        messageContainer.onScreenshotTaken(screenshotData)
+
+        verify(workProfileMessageController, never()).populateView(any(), any(), any())
+        verify(screenshotDetectionController, never()).populateView(any(), any())
+
+        assertEquals(View.GONE, container.visibility)
+    }
+}

packages/SystemUI/tests/src/com/android/systemui/screenshot/WorkProfileMessageControllerTest.java

@@ -16,13 +16,11 @@
 
 package com.android.systemui.screenshot;
 
-import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.ArgumentMatchers.nullable;
-import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.content.ComponentName;
@@ -33,24 +31,33 @@ import android.graphics.drawable.Drawable;
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.testing.AndroidTestingRunner;
+import android.view.LayoutInflater;
+import android.view.ViewGroup;
+import android.widget.ImageView;
+import android.widget.TextView;
 
 import androidx.test.filters.SmallTest;
 
+import com.android.systemui.R;
+import com.android.systemui.SysuiTestCase;
 import com.android.systemui.util.FakeSharedPreferences;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.ArgumentCaptor;
 import org.mockito.ArgumentMatchers;
-import org.mockito.Captor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+import kotlin.Unit;
+
 
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
-public class WorkProfileMessageControllerTest {
+public class WorkProfileMessageControllerTest extends SysuiTestCase {
     private static final String DEFAULT_LABEL = "default label";
     private static final String BADGED_DEFAULT_LABEL = "badged default label";
     private static final String APP_LABEL = "app label";
@@ -63,17 +70,13 @@ public class WorkProfileMessageControllerTest {
     @Mock
     private PackageManager mPackageManager;
     @Mock
-    private Context mContext;
+    private Context mMockContext;
-    @Mock
-    private MessageContainerController mMessageDisplay;
     @Mock
     private Drawable mActivityIcon;
     @Mock
     private Drawable mBadgedActivityIcon;
     @Mock
     private ActivityInfo mActivityInfo;
-    @Captor
-    private ArgumentCaptor<Runnable> mRunnableArgumentCaptor;
 
     private FakeSharedPreferences mSharedPreferences = new FakeSharedPreferences();
 
@@ -84,10 +87,10 @@ public class WorkProfileMessageControllerTest {
         MockitoAnnotations.initMocks(this);
 
         when(mUserManager.isManagedProfile(eq(WORK_USER.getIdentifier()))).thenReturn(true);
-        when(mContext.getSharedPreferences(
+        when(mMockContext.getSharedPreferences(
                 eq(WorkProfileMessageController.SHARED_PREFERENCES_NAME),
                 eq(Context.MODE_PRIVATE))).thenReturn(mSharedPreferences);
-        when(mContext.getString(ArgumentMatchers.anyInt())).thenReturn(DEFAULT_LABEL);
+        when(mMockContext.getString(ArgumentMatchers.anyInt())).thenReturn(DEFAULT_LABEL);
         when(mPackageManager.getUserBadgedLabel(eq(DEFAULT_LABEL), any()))
                 .thenReturn(BADGED_DEFAULT_LABEL);
         when(mPackageManager.getUserBadgedLabel(eq(APP_LABEL), any()))
@@ -103,16 +106,13 @@ public class WorkProfileMessageControllerTest {
         mSharedPreferences.edit().putBoolean(
                 WorkProfileMessageController.PREFERENCE_KEY, false).apply();
 
-        mMessageController = new WorkProfileMessageController(mContext, mUserManager,
+        mMessageController = new WorkProfileMessageController(mMockContext, mUserManager,
                 mPackageManager);
     }
 
     @Test
     public void testOnScreenshotTaken_notManaged() {
-        mMessageController.onScreenshotTaken(NON_WORK_USER, mMessageDisplay);
+        assertNull(mMessageController.onScreenshotTaken(NON_WORK_USER));
-
-        verify(mMessageDisplay, never())
-                .showWorkProfileMessage(any(), nullable(Drawable.class), any());
     }
 
     @Test
@@ -120,10 +120,7 @@ public class WorkProfileMessageControllerTest {
         mSharedPreferences.edit().putBoolean(
                 WorkProfileMessageController.PREFERENCE_KEY, true).apply();
 
-        mMessageController.onScreenshotTaken(WORK_USER, mMessageDisplay);
+        assertNull(mMessageController.onScreenshotTaken(WORK_USER));
-
-        verify(mMessageDisplay, never())
-                .showWorkProfileMessage(any(), nullable(Drawable.class), any());
     }
 
     @Test
@@ -133,28 +130,45 @@ public class WorkProfileMessageControllerTest {
                 any(PackageManager.ComponentInfoFlags.class))).thenThrow(
                 new PackageManager.NameNotFoundException());
 
-        mMessageController.onScreenshotTaken(WORK_USER, mMessageDisplay);
+        WorkProfileMessageController.WorkProfileFirstRunData data =
+                mMessageController.onScreenshotTaken(WORK_USER);
 
-        verify(mMessageDisplay).showWorkProfileMessage(
+        assertEquals(BADGED_DEFAULT_LABEL, data.getAppName());
-                eq(BADGED_DEFAULT_LABEL), eq(null), any());
+        assertNull(data.getIcon());
     }
 
     @Test
     public void testOnScreenshotTaken() {
-        mMessageController.onScreenshotTaken(WORK_USER, mMessageDisplay);
+        WorkProfileMessageController.WorkProfileFirstRunData data =
+                mMessageController.onScreenshotTaken(WORK_USER);
 
-        verify(mMessageDisplay).showWorkProfileMessage(
+        assertEquals(BADGED_APP_LABEL, data.getAppName());
-                eq(BADGED_APP_LABEL), eq(mBadgedActivityIcon), mRunnableArgumentCaptor.capture());
+        assertEquals(mBadgedActivityIcon, data.getIcon());
+    }
 
-        // Dismiss hasn't been tapped, preference untouched.
+    @Test
-        assertFalse(
+    public void testPopulateView() throws InterruptedException {
-                mSharedPreferences.getBoolean(WorkProfileMessageController.PREFERENCE_KEY, false));
+        ViewGroup layout = (ViewGroup) LayoutInflater.from(mContext).inflate(
+                R.layout.screenshot_work_profile_first_run, null);
+        WorkProfileMessageController.WorkProfileFirstRunData data =
+                new WorkProfileMessageController.WorkProfileFirstRunData(BADGED_APP_LABEL,
+                        mBadgedActivityIcon);
+        final CountDownLatch countdown = new CountDownLatch(1);
+        mMessageController.populateView(layout, data, () -> {
+            countdown.countDown();
+            return Unit.INSTANCE;
+        });
 
-        mRunnableArgumentCaptor.getValue().run();
+        ImageView image = layout.findViewById(R.id.screenshot_message_icon);
+        assertEquals(mBadgedActivityIcon, image.getDrawable());
+        TextView text = layout.findViewById(R.id.screenshot_message_content);
+        // The app name is used in a template, but at least validate that it was inserted.
+        assertTrue(text.getText().toString().contains(BADGED_APP_LABEL));
 
-        // After dismiss has been tapped, the setting should be updated.
+        // Validate that clicking the dismiss button calls back properly.
-        assertTrue(
+        assertEquals(1, countdown.getCount());
-                mSharedPreferences.getBoolean(WorkProfileMessageController.PREFERENCE_KEY, false));
+        layout.findViewById(R.id.message_dismiss_button).callOnClick();
+        countdown.await(1000, TimeUnit.MILLISECONDS);
     }
 }