Ensure attachment save to external completes successfully

ET-6116

mail-attachments/presentation/build.gradle.kts

@@ -56,6 +56,7 @@ dependencies {
     implementation(libs.dagger.hilt.android)
     implementation(libs.bundles.module.presentation)
     implementation(project(":mail-attachments:domain"))
+    implementation(project(":mail-common:domain"))
 
     testImplementation(libs.bundles.test)
     testImplementation(libs.androidx.test.core)

mail-attachments/presentation/src/main/kotlin/ch/protonmail/android/mailattachments/presentation/ExternalAttachmentsHandlerImpl.kt

@@ -27,13 +27,16 @@ import arrow.core.Either
 import arrow.core.raise.either
 import ch.protonmail.android.mailattachments.presentation.model.FileContent
 import ch.protonmail.android.mailattachments.presentation.usecase.GenerateUniqueFileName
+import ch.protonmail.android.mailcommon.domain.coroutines.IODispatcher
 import dagger.hilt.android.qualifiers.ApplicationContext
-import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.NonCancellable
 import kotlinx.coroutines.withContext
 import javax.inject.Inject
 
 class ExternalAttachmentsHandlerImpl @Inject constructor(
     @ApplicationContext private val context: Context,
+    @IODispatcher private val ioDispatcher: CoroutineDispatcher,
     private val generateUniqueFileName: GenerateUniqueFileName
 ) : ExternalAttachmentsHandler {
 
@@ -41,7 +44,7 @@ class ExternalAttachmentsHandlerImpl @Inject constructor(
         sourceUri: Uri,
         destinationUri: Uri
     ): Either<ExternalAttachmentErrorResult, Unit> = either {
-        withContext(Dispatchers.IO) {
+        withContext(ioDispatcher + NonCancellable) {
             context.contentResolver.openInputStream(sourceUri)?.use { inputStream ->
                 context.contentResolver.openOutputStream(destinationUri)?.use { outputStream ->
                     inputStream.copyTo(outputStream, bufferSize = 8 * 1024)
@@ -53,7 +56,7 @@ class ExternalAttachmentsHandlerImpl @Inject constructor(
     override suspend fun saveFileToDownloadsFolder(
         fileContent: FileContent
     ): Either<ExternalAttachmentErrorResult, Unit> = either {
-        withContext(Dispatchers.IO) {
+        withContext(ioDispatcher + NonCancellable) {
             val values = ContentValues().apply {
                 put(MediaStore.Downloads.DISPLAY_NAME, generateUniqueFileName(fileContent.name))
                 put(MediaStore.Downloads.MIME_TYPE, fileContent.mimeType)
@@ -75,7 +78,7 @@ class ExternalAttachmentsHandlerImpl @Inject constructor(
         mimeType: String,
         data: ByteArray
     ): Either<ExternalAttachmentErrorResult, Unit> = either {
-        withContext(Dispatchers.IO) {
+        withContext(ioDispatcher + NonCancellable) {
             try {
                 context.contentResolver.openOutputStream(destinationUri)?.use { outputStream ->
                     outputStream.write(data)
@@ -91,12 +94,11 @@ class ExternalAttachmentsHandlerImpl @Inject constructor(
         mimeType: String,
         data: ByteArray
     ): Either<ExternalAttachmentErrorResult, Unit> = either {
-        withContext(Dispatchers.IO) {
+        withContext(ioDispatcher + NonCancellable) {
             val values = ContentValues().apply {
                 put(MediaStore.Downloads.DISPLAY_NAME, fileName)
                 put(MediaStore.Downloads.MIME_TYPE, mimeType)
             }
-
             val uri = context.contentResolver.insert(MediaStore.Downloads.EXTERNAL_CONTENT_URI, values)
                 ?: raise(ExternalAttachmentErrorResult.UnableToCreateUri)
             try {

mail-attachments/presentation/src/test/kotlin/ch/protonmail/android/mailattachments/presentation/ExternalAttachmentsHandlerImplSaveTest.kt

@@ -34,6 +34,7 @@ import io.mockk.every
 import io.mockk.mockk
 import io.mockk.spyk
 import io.mockk.verify
+import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.test.runTest
 import org.junit.runner.RunWith
 import org.robolectric.RobolectricTestRunner
@@ -47,7 +48,6 @@ import kotlin.test.assertEquals
 internal class ExternalAttachmentsHandlerImplSaveTest {
 
     private lateinit var context: Context
-
     private lateinit var attachmentsHandler: ExternalAttachmentsHandlerImpl
 
     private val generateUniqueFileName = mockk<GenerateUniqueFileName> {
@@ -57,7 +57,7 @@ internal class ExternalAttachmentsHandlerImplSaveTest {
     @BeforeTest
     fun setup() {
         context = spyk(RuntimeEnvironment.getApplication().applicationContext)
-        attachmentsHandler = ExternalAttachmentsHandlerImpl(context, generateUniqueFileName)
+        attachmentsHandler = ExternalAttachmentsHandlerImpl(context, Dispatchers.Unconfined, generateUniqueFileName)
     }
 
     @AfterTest

mail-attachments/presentation/src/test/kotlin/ch/protonmail/android/mailattachments/presentation/ExternalAttachmentsHandlerImplTest.kt

@@ -31,6 +31,11 @@ import io.mockk.confirmVerified
 import io.mockk.every
 import io.mockk.mockk
 import io.mockk.verify
+import kotlinx.coroutines.CoroutineStart
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.advanceUntilIdle
 import kotlinx.coroutines.test.runTest
 import kotlin.test.AfterTest
 import kotlin.test.BeforeTest
@@ -41,7 +46,6 @@ internal class ExternalAttachmentsHandlerImplTest {
 
     private val context = mockk<Context>()
     private lateinit var contentResolver: ContentResolver
-
     private lateinit var attachmentsHandler: ExternalAttachmentsHandlerImpl
 
     private val generateUniqueFileName = mockk<GenerateUniqueFileName> {
@@ -50,9 +54,9 @@ internal class ExternalAttachmentsHandlerImplTest {
 
     @BeforeTest
     fun setup() {
-        attachmentsHandler = ExternalAttachmentsHandlerImpl(context, generateUniqueFileName)
         contentResolver = mockk(relaxed = true)
         every { context.contentResolver } returns contentResolver
+        attachmentsHandler = ExternalAttachmentsHandlerImpl(context, Dispatchers.Unconfined, generateUniqueFileName)
     }
 
     @AfterTest
@@ -153,4 +157,33 @@ internal class ExternalAttachmentsHandlerImplTest {
         verify(exactly = 1) { contentResolver.openOutputStream(destinationUri) }
         confirmVerified(contentResolver)
     }
+
+    @Test
+    fun `copyUriToDestination should complete even when coroutine is cancelled`() = runTest {
+        // Given
+        val handler = ExternalAttachmentsHandlerImpl(
+            context, StandardTestDispatcher(testScheduler), generateUniqueFileName
+        )
+        val sourceUri = mockk<Uri>()
+        val destinationUri = mockk<Uri>()
+        val content = "test content".toByteArray()
+        val outputStream = ByteArrayOutputStream()
+
+        val inputStream = ByteArrayInputStream(content)
+
+        every { contentResolver.openInputStream(sourceUri) } returns inputStream
+        every { contentResolver.openOutputStream(destinationUri) } returns outputStream
+
+        // When
+        // CoroutineStart.UNDISPATCHED lets the coroutine run until it suspends at withContext(ioDispatcher),
+        // then we cancel the parent job before the IO block is dispatched
+        val job = launch(start = CoroutineStart.UNDISPATCHED) {
+            handler.copyUriToDestination(sourceUri, destinationUri)
+        }
+        job.cancel()
+        advanceUntilIdle()
+
+        // Then
+        assertEquals(content.toList(), outputStream.toByteArray().toList())
+    }
 }