Merge "Only honor passed-in userid value if work profile flag enabled" into tm-qpr-dev

4 files changed, 82 insertions, 11 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ImageExporter.java b/packages/SystemUI/src/com/android/systemui/screenshot/ImageExporter.java
index e3658defc52a..c8c133774766 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ImageExporter.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ImageExporter.java
@@ -38,6 +38,8 @@ import androidx.concurrent.futures.CallbackToFutureAdapter;
 import androidx.exifinterface.media.ExifInterface;
 
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.systemui.flags.FeatureFlags;
+import com.android.systemui.flags.Flags;
 
 import com.google.common.util.concurrent.ListenableFuture;
 
@@ -85,10 +87,12 @@ class ImageExporter {
     private final ContentResolver mResolver;
     private CompressFormat mCompressFormat = CompressFormat.PNG;
     private int mQuality = 100;
+    private final FeatureFlags mFlags;
 
     @Inject
-    ImageExporter(ContentResolver resolver) {
+    ImageExporter(ContentResolver resolver, FeatureFlags flags) {
         mResolver = resolver;
+        mFlags = flags;
     }
 
     /**
@@ -161,7 +165,7 @@ class ImageExporter {
             ZonedDateTime captureTime, UserHandle owner) {
 
         final Task task = new Task(mResolver, requestId, bitmap, captureTime, mCompressFormat,
-                mQuality, /* publish */ true, owner);
+                mQuality, /* publish */ true, owner, mFlags);
 
         return CallbackToFutureAdapter.getFuture(
                 (completer) -> {
@@ -209,9 +213,11 @@ class ImageExporter {
         private final UserHandle mOwner;
         private final String mFileName;
         private final boolean mPublish;
+        private final FeatureFlags mFlags;
 
         Task(ContentResolver resolver, UUID requestId, Bitmap bitmap, ZonedDateTime captureTime,
-                CompressFormat format, int quality, boolean publish, UserHandle owner) {
+                CompressFormat format, int quality, boolean publish, UserHandle owner,
+                FeatureFlags flags) {
             mResolver = resolver;
             mRequestId = requestId;
             mBitmap = bitmap;
@@ -221,6 +227,7 @@ class ImageExporter {
             mOwner = owner;
             mFileName = createFilename(mCaptureTime, mFormat);
             mPublish = publish;
+            mFlags = flags;
         }
 
         public Result execute() throws ImageExportException, InterruptedException {
@@ -234,7 +241,7 @@ class ImageExporter {
                     start = Instant.now();
                 }
 
-                uri = createEntry(mResolver, mFormat, mCaptureTime, mFileName, mOwner);
+                uri = createEntry(mResolver, mFormat, mCaptureTime, mFileName, mOwner, mFlags);
                 throwIfInterrupted();
 
                 writeImage(mResolver, mBitmap, mFormat, mQuality, uri);
@@ -278,13 +285,15 @@ class ImageExporter {
     }
 
     private static Uri createEntry(ContentResolver resolver, CompressFormat format,
-            ZonedDateTime time, String fileName, UserHandle owner) throws ImageExportException {
+            ZonedDateTime time, String fileName, UserHandle owner, FeatureFlags flags)
+            throws ImageExportException {
         Trace.beginSection("ImageExporter_createEntry");
         try {
             final ContentValues values = createMetadata(time, format, fileName);
 
             Uri baseUri = MediaStore.Images.Media.EXTERNAL_CONTENT_URI;
-            if (UserHandle.myUserId() != owner.getIdentifier()) {
+            if (flags.isEnabled(Flags.SCREENSHOT_WORK_PROFILE_POLICY)
+                    && UserHandle.myUserId() != owner.getIdentifier()) {
                 baseUri = ContentProvider.maybeAddUserId(baseUri, owner.getIdentifier());
             }
             Uri uri = resolver.insert(baseUri, values);
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
index d524a356a323..9b5295d1bb3d 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
@@ -1043,8 +1043,13 @@ public class ScreenshotController {
     }
 
     private boolean isUserSetupComplete(UserHandle owner) {
-        return Settings.Secure.getInt(mContext.createContextAsUser(owner, 0)
-                .getContentResolver(), SETTINGS_SECURE_USER_SETUP_COMPLETE, 0) == 1;
+        if (mFlags.isEnabled(SCREENSHOT_WORK_PROFILE_POLICY)) {
+            return Settings.Secure.getInt(mContext.createContextAsUser(owner, 0)
+                    .getContentResolver(), SETTINGS_SECURE_USER_SETUP_COMPLETE, 0) == 1;
+        } else {
+            return Settings.Secure.getInt(mContext.getContentResolver(),
+                    SETTINGS_SECURE_USER_SETUP_COMPLETE, 0) == 1;
+        }
     }
 
     /**
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ImageExporterTest.java b/packages/SystemUI/tests/src/com/android/systemui/screenshot/ImageExporterTest.java
index 4c44dacab1a2..f4bc232702e8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ImageExporterTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/ImageExporterTest.java
@@ -23,6 +23,7 @@ import static org.junit.Assert.assertTrue;
 
 import static java.nio.charset.StandardCharsets.US_ASCII;
 
+import android.content.ContentProvider;
 import android.content.ContentResolver;
 import android.content.ContentValues;
 import android.graphics.Bitmap;
@@ -31,9 +32,11 @@ import android.graphics.BitmapFactory;
 import android.graphics.Canvas;
 import android.graphics.Color;
 import android.graphics.Paint;
+import android.net.Uri;
 import android.os.Build;
 import android.os.ParcelFileDescriptor;
 import android.os.Process;
+import android.os.UserHandle;
 import android.provider.MediaStore;
 import android.testing.AndroidTestingRunner;
 
@@ -41,11 +44,18 @@ import androidx.exifinterface.media.ExifInterface;
 import androidx.test.filters.MediumTest;
 
 import com.android.systemui.SysuiTestCase;
+import com.android.systemui.flags.FakeFeatureFlags;
+import com.android.systemui.flags.Flags;
 
 import com.google.common.util.concurrent.ListenableFuture;
 
+import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.MockitoAnnotations;
 
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
@@ -60,7 +70,6 @@ import java.util.concurrent.Executor;
 @RunWith(AndroidTestingRunner.class)
 @MediumTest // file I/O
 public class ImageExporterTest extends SysuiTestCase {
-
     /** Executes directly in the caller's thread */
     private static final Executor DIRECT_EXECUTOR = Runnable::run;
     private static final byte[] EXIF_FILE_TAG = "Exif\u0000\u0000".getBytes(US_ASCII);
@@ -68,6 +77,15 @@ public class ImageExporterTest extends SysuiTestCase {
     private static final ZonedDateTime CAPTURE_TIME =
             ZonedDateTime.of(LocalDateTime.of(2020, 12, 15, 13, 15), ZoneId.of("EST"));
 
+    private FakeFeatureFlags mFeatureFlags = new FakeFeatureFlags();
+    @Mock
+    private ContentResolver mMockContentResolver;
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+    }
+
     @Test
     public void testImageFilename() {
         assertEquals("image file name", "Screenshot_20201215-131500.png",
@@ -92,7 +110,8 @@ public class ImageExporterTest extends SysuiTestCase {
     @Test
     public void testImageExport() throws ExecutionException, InterruptedException, IOException {
         ContentResolver contentResolver = mContext.getContentResolver();
-        ImageExporter exporter = new ImageExporter(contentResolver);
+        mFeatureFlags.set(Flags.SCREENSHOT_WORK_PROFILE_POLICY, true);
+        ImageExporter exporter = new ImageExporter(contentResolver, mFeatureFlags);
 
         UUID requestId = UUID.fromString("3c11da99-9284-4863-b1d5-6f3684976814");
         Bitmap original = createCheckerBitmap(10, 10, 10);
@@ -168,6 +187,44 @@ public class ImageExporterTest extends SysuiTestCase {
                 values.getAsLong(MediaStore.MediaColumns.DATE_EXPIRES));
     }
 
+    @Test
+    public void testSetUser() {
+        mFeatureFlags.set(Flags.SCREENSHOT_WORK_PROFILE_POLICY, true);
+        ImageExporter exporter = new ImageExporter(mMockContentResolver, mFeatureFlags);
+
+        UserHandle imageUserHande = UserHandle.of(10);
+
+        ArgumentCaptor<Uri> uriCaptor = ArgumentCaptor.forClass(Uri.class);
+        // Capture the URI and then return null to bail out of export.
+        Mockito.when(mMockContentResolver.insert(uriCaptor.capture(), Mockito.any())).thenReturn(
+                null);
+        exporter.export(DIRECT_EXECUTOR, UUID.fromString("3c11da99-9284-4863-b1d5-6f3684976814"),
+                null, CAPTURE_TIME, imageUserHande);
+
+        Uri expected = MediaStore.Images.Media.EXTERNAL_CONTENT_URI;
+        expected = ContentProvider.maybeAddUserId(expected, imageUserHande.getIdentifier());
+
+        assertEquals(expected, uriCaptor.getValue());
+    }
+
+    @Test
+    public void testSetUser_noWorkProfile() {
+        mFeatureFlags.set(Flags.SCREENSHOT_WORK_PROFILE_POLICY, false);
+        ImageExporter exporter = new ImageExporter(mMockContentResolver, mFeatureFlags);
+
+        UserHandle imageUserHandle = UserHandle.of(10);
+
+        ArgumentCaptor<Uri> uriCaptor = ArgumentCaptor.forClass(Uri.class);
+        // Capture the URI and then return null to bail out of export.
+        Mockito.when(mMockContentResolver.insert(uriCaptor.capture(), Mockito.any())).thenReturn(
+                null);
+        exporter.export(DIRECT_EXECUTOR, UUID.fromString("3c11da99-9284-4863-b1d5-6f3684976814"),
+                null, CAPTURE_TIME, imageUserHandle);
+
+        // The user handle should be ignored here since the flag is off.
+        assertEquals(MediaStore.Images.Media.EXTERNAL_CONTENT_URI, uriCaptor.getValue());
+    }
+
     @SuppressWarnings("SameParameterValue")
     private Bitmap createCheckerBitmap(int tileSize, int w, int h) {
         Bitmap bitmap = Bitmap.createBitmap(w * tileSize, h * tileSize, Bitmap.Config.ARGB_8888);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotServiceTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotServiceTest.kt
index 3a4da86b8045..fa1fedb7c119 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotServiceTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotServiceTest.kt
@@ -62,7 +62,7 @@ import org.mockito.Mockito.verifyZeroInteractions
 import org.mockito.Mockito.`when` as whenever
 
 private const val USER_ID = 1
-private const val TASK_ID = 1
+private const val TASK_ID = 11
 
 @RunWith(AndroidTestingRunner::class)
 @SmallTest