Allow returned screenshot URI to be null

Occurs in failure cases (if the screenshot does not successfully save).

Flag: NONE
Test: manual, ensure that ending screenshot process before screenshot
saves does not cause TakeScreenshotExecutor null error

Change-Id: I891eac48ccf77101dc5ff2fcd0dd492c53343e59

3 files changed, 46 insertions, 26 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt
index ec7707c83980..e56a4f45d7aa 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt
@@ -22,7 +22,7 @@ import kotlinx.coroutines.launch
 interface TakeScreenshotExecutor {
     suspend fun executeScreenshots(
         screenshotRequest: ScreenshotRequest,
-        onSaved: (Uri) -> Unit,
+        onSaved: (Uri?) -> Unit,
         requestCallback: RequestCallback
     )
     fun onCloseSystemDialogsReceived()
@@ -30,7 +30,7 @@ interface TakeScreenshotExecutor {
     fun onDestroy()
     fun executeScreenshotsAsync(
         screenshotRequest: ScreenshotRequest,
-        onSaved: Consumer<Uri>,
+        onSaved: Consumer<Uri?>,
         requestCallback: RequestCallback
     )
 }
@@ -65,7 +65,7 @@ constructor(
      */
     override suspend fun executeScreenshots(
         screenshotRequest: ScreenshotRequest,
-        onSaved: (Uri) -> Unit,
+        onSaved: (Uri?) -> Unit,
         requestCallback: RequestCallback
     ) {
         val displayIds = getDisplaysToScreenshot(screenshotRequest.type)
@@ -86,7 +86,7 @@ constructor(
     /** All logging should be triggered only by this method. */
     private suspend fun dispatchToController(
         rawScreenshotData: ScreenshotData,
-        onSaved: (Uri) -> Unit,
+        onSaved: (Uri?) -> Unit,
         callback: RequestCallback
     ) {
         // Let's wait before logging "screenshot requested", as we should log the processed
@@ -185,7 +185,7 @@ constructor(
     /** For java compatibility only. see [executeScreenshots] */
     override fun executeScreenshotsAsync(
         screenshotRequest: ScreenshotRequest,
-        onSaved: Consumer<Uri>,
+        onSaved: Consumer<Uri?>,
         requestCallback: RequestCallback
     ) {
         mainScope.launch {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt
index c900463c7159..0f3714385725 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt
@@ -26,6 +26,7 @@ import com.android.systemui.util.mockito.nullable
 import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
 import java.lang.IllegalStateException
+import java.util.function.Consumer
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import kotlinx.coroutines.test.runCurrent
@@ -78,7 +79,7 @@ class TakeScreenshotExecutorTest : SysuiTestCase() {
     fun executeScreenshots_severalDisplays_callsControllerForEachOne() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
-            val onSaved = { _: Uri -> }
+            val onSaved = { _: Uri? -> }
             screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
 
             verify(controllerFactory).create(eq(0), any())
@@ -107,7 +108,7 @@ class TakeScreenshotExecutorTest : SysuiTestCase() {
     fun executeScreenshots_providedImageType_callsOnlyDefaultDisplayController() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
-            val onSaved = { _: Uri -> }
+            val onSaved = { _: Uri? -> }
             screenshotExecutor.executeScreenshots(
                 createScreenshotRequest(TAKE_SCREENSHOT_PROVIDED_IMAGE),
                 onSaved,
@@ -136,7 +137,7 @@ class TakeScreenshotExecutorTest : SysuiTestCase() {
     fun executeScreenshots_onlyVirtualDisplays_noInteractionsWithControllers() =
         testScope.runTest {
             setDisplays(display(TYPE_VIRTUAL, id = 0), display(TYPE_VIRTUAL, id = 1))
-            val onSaved = { _: Uri -> }
+            val onSaved = { _: Uri? -> }
             screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
 
             verifyNoMoreInteractions(controllerFactory)
@@ -154,7 +155,7 @@ class TakeScreenshotExecutorTest : SysuiTestCase() {
                 display(TYPE_OVERLAY, id = 2),
                 display(TYPE_WIFI, id = 3)
             )
-            val onSaved = { _: Uri -> }
+            val onSaved = { _: Uri? -> }
             screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
 
             verify(controller0, times(4)).handleScreenshot(any(), any(), any())
@@ -165,7 +166,7 @@ class TakeScreenshotExecutorTest : SysuiTestCase() {
     fun executeScreenshots_reportsOnFinishedOnlyWhenBothFinished() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
-            val onSaved = { _: Uri -> }
+            val onSaved = { _: Uri? -> }
             screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
 
             val capturer0 = ArgumentCaptor<TakeScreenshotService.RequestCallback>()
@@ -190,7 +191,7 @@ class TakeScreenshotExecutorTest : SysuiTestCase() {
     fun executeScreenshots_oneFinishesOtherFails_reportFailsOnlyAtTheEnd() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
-            val onSaved = { _: Uri -> }
+            val onSaved = { _: Uri? -> }
             screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
 
             val capturer0 = ArgumentCaptor<TakeScreenshotService.RequestCallback>()
@@ -217,7 +218,7 @@ class TakeScreenshotExecutorTest : SysuiTestCase() {
     fun executeScreenshots_allDisplaysFail_reportsFail() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
-            val onSaved = { _: Uri -> }
+            val onSaved = { _: Uri? -> }
             screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
 
             val capturer0 = ArgumentCaptor<TakeScreenshotService.RequestCallback>()
@@ -244,7 +245,7 @@ class TakeScreenshotExecutorTest : SysuiTestCase() {
     fun onDestroy_propagatedToControllers() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
-            val onSaved = { _: Uri -> }
+            val onSaved = { _: Uri? -> }
             screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
 
             screenshotExecutor.onDestroy()
@@ -256,7 +257,7 @@ class TakeScreenshotExecutorTest : SysuiTestCase() {
     fun removeWindows_propagatedToControllers() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
-            val onSaved = { _: Uri -> }
+            val onSaved = { _: Uri? -> }
             screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
 
             screenshotExecutor.removeWindows()
@@ -270,7 +271,7 @@ class TakeScreenshotExecutorTest : SysuiTestCase() {
     fun onCloseSystemDialogsReceived_propagatedToControllers() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
-            val onSaved = { _: Uri -> }
+            val onSaved = { _: Uri? -> }
             screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
 
             screenshotExecutor.onCloseSystemDialogsReceived()
@@ -286,7 +287,7 @@ class TakeScreenshotExecutorTest : SysuiTestCase() {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
             whenever(controller0.isPendingSharedTransition).thenReturn(true)
             whenever(controller1.isPendingSharedTransition).thenReturn(false)
-            val onSaved = { _: Uri -> }
+            val onSaved = { _: Uri? -> }
             screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
 
             screenshotExecutor.onCloseSystemDialogsReceived()
@@ -304,7 +305,7 @@ class TakeScreenshotExecutorTest : SysuiTestCase() {
             val toBeReturnedByProcessor = ScreenshotData.forTesting()
             requestProcessor.toReturn = toBeReturnedByProcessor
 
-            val onSaved = { _: Uri -> }
+            val onSaved = { _: Uri? -> }
             screenshotExecutor.executeScreenshots(screenshotRequest, onSaved, callback)
 
             assertThat(requestProcessor.processed)
@@ -321,7 +322,7 @@ class TakeScreenshotExecutorTest : SysuiTestCase() {
     fun executeScreenshots_errorFromProcessor_logsScreenshotRequested() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
-            val onSaved = { _: Uri -> }
+            val onSaved = { _: Uri? -> }
             requestProcessor.shouldThrowException = true
 
             screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
@@ -338,7 +339,7 @@ class TakeScreenshotExecutorTest : SysuiTestCase() {
     fun executeScreenshots_errorFromProcessor_logsUiError() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
-            val onSaved = { _: Uri -> }
+            val onSaved = { _: Uri? -> }
             requestProcessor.shouldThrowException = true
 
             screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
@@ -355,7 +356,7 @@ class TakeScreenshotExecutorTest : SysuiTestCase() {
     fun executeScreenshots_errorFromProcessorOnDefaultDisplay_showsErrorNotification() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
-            val onSaved = { _: Uri -> }
+            val onSaved = { _: Uri? -> }
             requestProcessor.shouldThrowException = true
 
             screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
@@ -368,7 +369,7 @@ class TakeScreenshotExecutorTest : SysuiTestCase() {
     fun executeScreenshots_errorFromProcessorOnSecondaryDisplay_showsErrorNotification() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0))
-            val onSaved = { _: Uri -> }
+            val onSaved = { _: Uri? -> }
             requestProcessor.shouldThrowException = true
 
             screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
@@ -381,7 +382,7 @@ class TakeScreenshotExecutorTest : SysuiTestCase() {
     fun executeScreenshots_errorFromScreenshotController_reportsRequested() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
-            val onSaved = { _: Uri -> }
+            val onSaved = { _: Uri? -> }
             whenever(controller0.handleScreenshot(any(), any(), any()))
                 .thenThrow(IllegalStateException::class.java)
             whenever(controller1.handleScreenshot(any(), any(), any()))
@@ -401,7 +402,7 @@ class TakeScreenshotExecutorTest : SysuiTestCase() {
     fun executeScreenshots_errorFromScreenshotController_reportsError() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
-            val onSaved = { _: Uri -> }
+            val onSaved = { _: Uri? -> }
             whenever(controller0.handleScreenshot(any(), any(), any()))
                 .thenThrow(IllegalStateException::class.java)
             whenever(controller1.handleScreenshot(any(), any(), any()))
@@ -421,7 +422,7 @@ class TakeScreenshotExecutorTest : SysuiTestCase() {
     fun executeScreenshots_errorFromScreenshotController_showsErrorNotification() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
-            val onSaved = { _: Uri -> }
+            val onSaved = { _: Uri? -> }
             whenever(controller0.handleScreenshot(any(), any(), any()))
                 .thenThrow(IllegalStateException::class.java)
             whenever(controller1.handleScreenshot(any(), any(), any()))
@@ -434,6 +435,25 @@ class TakeScreenshotExecutorTest : SysuiTestCase() {
             screenshotExecutor.onDestroy()
         }
 
+    @Test
+    fun executeScreenshots_finisherCalledWithNullUri_succeeds() =
+        testScope.runTest {
+            setDisplays(display(TYPE_INTERNAL, id = 0))
+            var onSavedCallCount = 0
+            val onSaved: (Uri?) -> Unit = {
+                assertThat(it).isNull()
+                onSavedCallCount += 1
+            }
+            whenever(controller0.handleScreenshot(any(), any(), any())).thenAnswer {
+                (it.getArgument(1) as Consumer<Uri?>).accept(null)
+            }
+
+            screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
+            assertThat(onSavedCallCount).isEqualTo(1)
+
+            screenshotExecutor.onDestroy()
+        }
+
     private suspend fun TestScope.setDisplays(vararg displays: Display) {
         fakeDisplayRepository.emit(displays.toSet())
         runCurrent()
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 0776aa7d1845..77b5c9115295 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotServiceTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotServiceTest.kt
@@ -232,7 +232,7 @@ internal class FakeScreenshotExecutor : TakeScreenshotExecutor {
     override fun onCloseSystemDialogsReceived() {}
     override suspend fun executeScreenshots(
         screenshotRequest: ScreenshotRequest,
-        onSaved: (Uri) -> Unit,
+        onSaved: (Uri?) -> Unit,
         requestCallback: RequestCallback,
     ) {
         requestReceived = screenshotRequest
@@ -248,7 +248,7 @@ internal class FakeScreenshotExecutor : TakeScreenshotExecutor {
 
     override fun executeScreenshotsAsync(
         screenshotRequest: ScreenshotRequest,
-        onSaved: Consumer<Uri>,
+        onSaved: Consumer<Uri?>,
         requestCallback: RequestCallback,
     ) {
         runBlocking {