fix(storage): Prevent missed transfer status updates

aws-storage-s3/src/main/java/com/amplifyframework/storage/s3/AWSS3StoragePlugin.java

@@ -94,6 +94,7 @@
 import com.amplifyframework.storage.s3.service.AWSS3StorageServiceContainer;
 import com.amplifyframework.storage.s3.transfer.S3StorageTransferClientProvider;
 import com.amplifyframework.storage.s3.transfer.StorageTransferClientProvider;
+import com.amplifyframework.storage.s3.transfer.TransferDB;
 import com.amplifyframework.storage.s3.transfer.TransferObserver;
 import com.amplifyframework.storage.s3.transfer.TransferRecord;
 import com.amplifyframework.storage.s3.transfer.TransferStatusUpdater;
@@ -154,6 +155,8 @@ public final class AWSS3StoragePlugin extends StoragePlugin<S3Client> {
                 return defaultStorageService.getClient();
             });
 
+    private TransferStatusUpdater transferStatusUpdater;
+
     /**
      * Constructs the AWS S3 Storage Plugin initializing the executor service.
      */
@@ -176,14 +179,15 @@ public AWSS3StoragePlugin(AWSS3StoragePluginConfiguration awsS3StoragePluginConf
 
     @VisibleForTesting
     AWSS3StoragePlugin(AuthCredentialsProvider authCredentialsProvider) {
-        this((context, region, bucket, clientProvider) ->
+        this((context, region, bucket, clientProvider, transferStatusUpdater) ->
                 new AWSS3StorageService(
                     context,
                     region,
                     bucket,
                     authCredentialsProvider,
                     AWS_S3_STORAGE_PLUGIN_KEY,
-                    clientProvider
+                    clientProvider,
+                    transferStatusUpdater
                 ),
             authCredentialsProvider,
             new AWSS3StoragePluginConfiguration.Builder().build());
@@ -193,14 +197,15 @@ public AWSS3StoragePlugin(AWSS3StoragePluginConfiguration awsS3StoragePluginConf
     AWSS3StoragePlugin(AuthCredentialsProvider authCredentialsProvider,
                        AWSS3StoragePluginConfiguration awss3StoragePluginConfiguration) {
 
-        this((context, region, bucket, clientProvider) ->
+        this((context, region, bucket, clientProvider, transferStatusUpdater) ->
                 new AWSS3StorageService(
                     context,
                     region,
                     bucket,
                     authCredentialsProvider,
                     AWS_S3_STORAGE_PLUGIN_KEY,
-                    clientProvider
+                    clientProvider,
+                    transferStatusUpdater
                 ),
             authCredentialsProvider,
             awss3StoragePluginConfiguration);
@@ -298,14 +303,20 @@ private void configure(
             @NonNull ResolvedStorageBucket bucket
     ) throws StorageException {
         try {
+            this.transferStatusUpdater = new TransferStatusUpdater(TransferDB.Companion.getInstance(context));
+
+            this.awss3StorageServiceContainer = new AWSS3StorageServiceContainer(
+                    context, storageServiceFactory,
+                    (S3StorageTransferClientProvider) clientProvider,
+                    transferStatusUpdater
+            );
             this.defaultStorageService = storageServiceFactory.create(
                     context,
                     region,
                     bucket.getBucketInfo().getBucketName(),
-                    clientProvider);
-            this.awss3StorageServiceContainer = new AWSS3StorageServiceContainer(
-                    context, storageServiceFactory,
-                    (S3StorageTransferClientProvider) clientProvider);
+                    clientProvider,
+                    transferStatusUpdater
+            );
             this.awss3StorageServiceContainer.put(bucket.getBucketInfo().getBucketName(), this.defaultStorageService);
         } catch (RuntimeException exception) {
             throw new StorageException(
@@ -927,7 +938,7 @@ public void getTransfer(
                     TransferObserver transferObserver =
                         new TransferObserver(
                             transferRecord.getId(),
-                            defaultStorageService.getTransferManager().getTransferStatusUpdater(),
+                            transferStatusUpdater,
                             transferRecord.getBucketName(),
                             transferRecord.getRegion(),
                             transferRecord.getKey(),

aws-storage-s3/src/main/java/com/amplifyframework/storage/s3/service/AWSS3StorageService.kt

@@ -37,6 +37,7 @@ import com.amplifyframework.storage.s3.transfer.StorageTransferClientProvider
 import com.amplifyframework.storage.s3.transfer.TransferManager
 import com.amplifyframework.storage.s3.transfer.TransferObserver
 import com.amplifyframework.storage.s3.transfer.TransferRecord
+import com.amplifyframework.storage.s3.transfer.TransferStatusUpdater
 import com.amplifyframework.storage.s3.transfer.UploadOptions
 import com.amplifyframework.storage.s3.utils.S3Keys
 import java.io.File
@@ -57,13 +58,14 @@ internal class AWSS3StorageService(
     private val s3BucketName: String,
     private val authCredentialsProvider: AuthCredentialsProvider,
     private val awsS3StoragePluginKey: String,
-    private val clientProvider: StorageTransferClientProvider
+    private val clientProvider: StorageTransferClientProvider,
+    private val transferStatusUpdater: TransferStatusUpdater
 ) : StorageService {
 
     private var s3Client: S3Client = S3StorageTransferClientProvider.getS3Client(awsRegion, authCredentialsProvider)
 
     val transferManager: TransferManager =
-        TransferManager(context, clientProvider, awsS3StoragePluginKey)
+        TransferManager(context, clientProvider, awsS3StoragePluginKey, transferStatusUpdater)
 
     /**
      * Generate pre-signed URL for an object.
@@ -421,7 +423,8 @@ internal class AWSS3StorageService(
             context: Context,
             region: String,
             bucketName: String,
-            clientProvider: StorageTransferClientProvider
+            clientProvider: StorageTransferClientProvider,
+            transferStatusUpdater: TransferStatusUpdater
         ): AWSS3StorageService
     }
 }

aws-storage-s3/src/main/java/com/amplifyframework/storage/s3/service/AWSS3StorageServiceContainer.kt

@@ -18,6 +18,7 @@ import android.content.Context
 import com.amplifyframework.storage.ResolvedStorageBucket
 import com.amplifyframework.storage.s3.transfer.S3StorageTransferClientProvider
 import com.amplifyframework.storage.s3.transfer.StorageTransferClientProvider
+import com.amplifyframework.storage.s3.transfer.TransferStatusUpdater
 import java.util.concurrent.ConcurrentHashMap
 
 /**
@@ -28,13 +29,15 @@ internal class AWSS3StorageServiceContainer(
     private val context: Context,
     private val storageServiceFactory: AWSS3StorageService.Factory,
     private val clientProvider: StorageTransferClientProvider,
-    private val awsS3StorageServicesByBucketName: ConcurrentHashMap<String, AWSS3StorageService>
+    private val awsS3StorageServicesByBucketName: ConcurrentHashMap<String, AWSS3StorageService>,
+    private val transferStatusUpdater: TransferStatusUpdater
 ) {
     constructor(
         context: Context,
         storageServiceFactory: AWSS3StorageService.Factory,
-        clientProvider: S3StorageTransferClientProvider
-    ) : this(context, storageServiceFactory, clientProvider, ConcurrentHashMap())
+        clientProvider: S3StorageTransferClientProvider,
+        transferStatusUpdater: TransferStatusUpdater
+    ) : this(context, storageServiceFactory, clientProvider, ConcurrentHashMap(), transferStatusUpdater)
 
     private val lock = Any()
 
@@ -61,7 +64,8 @@ internal class AWSS3StorageServiceContainer(
             var service = awsS3StorageServicesByBucketName.get(bucketName)
             if (service == null) {
                 val region: String = resolvedStorageBucket.bucketInfo.region
-                service = storageServiceFactory.create(context, region, bucketName, clientProvider)
+                service =
+                    storageServiceFactory.create(context, region, bucketName, clientProvider, transferStatusUpdater)
                 awsS3StorageServicesByBucketName[bucketName] = service
             }
             return service
@@ -78,7 +82,8 @@ internal class AWSS3StorageServiceContainer(
         synchronized(lock) {
             var service = awsS3StorageServicesByBucketName[bucketName]
             if (service == null) {
-                service = storageServiceFactory.create(context, region, bucketName, clientProvider)
+                service =
+                    storageServiceFactory.create(context, region, bucketName, clientProvider, transferStatusUpdater)
                 awsS3StorageServicesByBucketName[bucketName] = service
             }
 

aws-storage-s3/src/main/java/com/amplifyframework/storage/s3/transfer/TransferManager.kt

@@ -48,11 +48,11 @@ internal class TransferManager(
     context: Context,
     clientProvider: StorageTransferClientProvider,
     private val pluginKey: String,
+    private val transferStatusUpdater: TransferStatusUpdater,
     private val workManager: WorkManager = WorkManager.getInstance(context)
 ) {
 
     private val transferDB: TransferDB = TransferDB.getInstance(context)
-    val transferStatusUpdater: TransferStatusUpdater = TransferStatusUpdater(transferDB)
 
     private val logger =
         Amplify.Logging.logger(

aws-storage-s3/src/test/java/com/amplifyframework/storage/AWSS3StorageServiceContainerTest.kt

@@ -33,7 +33,7 @@ import org.junit.Test
 class AWSS3StorageServiceContainerTest {
 
     private val storageServiceFactory = mockk<AWSS3StorageService.Factory> {
-        every { create(any(), any(), any(), any()) } returns mockk<AWSS3StorageService>()
+        every { create(any(), any(), any(), any(), any()) } returns mockk<AWSS3StorageService>()
     }
     private val context = mockk<Context>()
     private val clientProvider = mockk<StorageTransferClientProvider>()
@@ -50,13 +50,14 @@ class AWSS3StorageServiceContainerTest {
             context,
             storageServiceFactory,
             clientProvider,
-            serviceContainerHashMap
+            serviceContainerHashMap,
+            mockk()
         )
     }
 
     @Test
     fun `put default AWSS3Service in container`() {
-        val service = storageServiceFactory.create(context, region, bucketName, clientProvider)
+        val service = storageServiceFactory.create(context, region, bucketName, clientProvider, mockk())
         serviceContainer.put(bucketName, service)
 
         serviceContainerHashMap.size shouldBe 1

aws-storage-s3/src/test/java/com/amplifyframework/storage/s3/AWSS3StoragePluginTest.kt

@@ -15,6 +15,8 @@
 
 package com.amplifyframework.storage.s3
 
+import android.content.Context
+import androidx.test.core.app.ApplicationProvider
 import com.amplifyframework.storage.BucketInfo
 import com.amplifyframework.storage.InvalidStorageBucketException
 import com.amplifyframework.storage.StorageBucket
@@ -29,11 +31,16 @@ import io.mockk.every
 import io.mockk.mockk
 import io.mockk.verify
 import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.RobolectricTestRunner
 
+@RunWith(RobolectricTestRunner::class)
 class AWSS3StoragePluginTest {
 
+    var context: Context = ApplicationProvider.getApplicationContext()
+
     private val storageServiceFactory = mockk<AWSS3StorageService.Factory> {
-        every { create(any(), any(), any(), any()) } returns mockk<AWSS3StorageService>()
+        every { create(any(), any(), any(), any(), any()) } returns mockk<AWSS3StorageService>()
     }
 
     private val plugin = AWSS3StoragePlugin(
@@ -51,10 +58,10 @@ class AWSS3StoragePluginTest {
             }
         }
 
-        plugin.configure(data, mockk())
+        plugin.configure(data, context)
 
         verify {
-            storageServiceFactory.create(any(), "test-region", "test-bucket", any())
+            storageServiceFactory.create(any(), "test-region", "test-bucket", any(), any())
         }
     }
 
@@ -65,7 +72,7 @@ class AWSS3StoragePluginTest {
         }
 
         shouldThrow<StorageException> {
-            plugin.configure(data, mockk())
+            plugin.configure(data, context)
         }
     }
 
@@ -83,7 +90,7 @@ class AWSS3StoragePluginTest {
             }
         }
 
-        plugin.configure(data, mockk())
+        plugin.configure(data, context)
         val service = plugin.getStorageService(null)
         service shouldNotBe null
     }
@@ -102,7 +109,7 @@ class AWSS3StoragePluginTest {
             }
         }
 
-        plugin.configure(data, mockk())
+        plugin.configure(data, context)
         val bucketInfo = BucketInfo("test-bucket", "test-region")
         val bucket = StorageBucket.fromBucketInfo(bucketInfo)
         val service = plugin.getStorageService(bucket)
@@ -123,7 +130,7 @@ class AWSS3StoragePluginTest {
             }
         }
 
-        plugin.configure(data, mockk())
+        plugin.configure(data, context)
         val bucket = StorageBucket.fromOutputs("test=name")
         val service = plugin.getStorageService(bucket)
         service shouldNotBe null
@@ -143,7 +150,7 @@ class AWSS3StoragePluginTest {
             }
         }
 
-        plugin.configure(data, mockk())
+        plugin.configure(data, context)
         val bucket = StorageBucket.fromOutputs("myBucket")
         val exception = shouldThrow<StorageException> {
             plugin.getStorageService(bucket)
@@ -165,7 +172,7 @@ class AWSS3StoragePluginTest {
             }
         }
 
-        plugin.configure(data, mockk())
+        plugin.configure(data, context)
         val bucket = StorageBucket.fromOutputs("test-name")
         val result = plugin.getStorageServiceResult(bucket)
         val service = result.storageService
@@ -188,7 +195,7 @@ class AWSS3StoragePluginTest {
             }
         }
 
-        plugin.configure(data, mockk())
+        plugin.configure(data, context)
         val bucket = StorageBucket.fromOutputs("myBucket")
         val exception = plugin.getStorageServiceResult(bucket).storageException
         exception shouldNotBe null

aws-storage-s3/src/test/java/com/amplifyframework/storage/s3/StorageComponentTest.java

@@ -90,8 +90,13 @@ public final class StorageComponentTest {
     public void setup() throws AmplifyException {
         this.storage = new StorageCategory();
         this.storageService = mock(AWSS3StorageService.class);
-        AWSS3StorageService.Factory storageServiceFactory
-                = (context, region, bucket, clientProvider) -> (AWSS3StorageService) storageService;
+        AWSS3StorageService.Factory storageServiceFactory = (
+                context,
+                region,
+                bucket,
+                clientProvider,
+                transferStatusUpdater
+        ) -> (AWSS3StorageService) storageService;
         AuthCredentialsProvider cognitoAuthProvider = mock(AuthCredentialsProvider.class);
         doReturn(RandomString.string()).when(cognitoAuthProvider).getIdentityId(null);
         this.storage.addPlugin(new AWSS3StoragePlugin(storageServiceFactory,