PM-34231: feat: Support renaming attachments during creation (#6742)

app/src/main/kotlin/com/x8bit/bitwarden/ui/vault/feature/attachments/AttachmentsContent.kt

@@ -1,5 +1,6 @@
 package com.x8bit.bitwarden.ui.vault.feature.attachments
 
+import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.Row
 import androidx.compose.foundation.layout.Spacer
 import androidx.compose.foundation.layout.defaultMinSize
@@ -29,8 +30,10 @@ import com.bitwarden.ui.platform.base.util.toListItemCardStyle
 import com.bitwarden.ui.platform.components.button.BitwardenOutlinedButton
 import com.bitwarden.ui.platform.components.button.BitwardenStandardIconButton
 import com.bitwarden.ui.platform.components.dialog.BitwardenTwoButtonDialog
+import com.bitwarden.ui.platform.components.field.BitwardenTextField
 import com.bitwarden.ui.platform.components.header.BitwardenListHeaderText
 import com.bitwarden.ui.platform.components.model.CardStyle
+import com.bitwarden.ui.platform.components.util.nonEditableExtensionVisualTransformation
 import com.bitwarden.ui.platform.resource.BitwardenDrawable
 import com.bitwarden.ui.platform.resource.BitwardenString
 import com.bitwarden.ui.platform.theme.BitwardenTheme
@@ -39,9 +42,34 @@ import com.x8bit.bitwarden.ui.vault.feature.attachments.handlers.AttachmentsHand
 /**
  * The top level content UI state for the [AttachmentsScreen] when viewing a content.
  */
-@Suppress("LongMethod")
 @Composable
 fun AttachmentsContent(
+    viewState: AttachmentsState.ViewState.Content,
+    attachmentsHandlers: AttachmentsHandlers,
+    isAttachmentUpdatesEnabled: Boolean,
+    modifier: Modifier = Modifier,
+) {
+    if (isAttachmentUpdatesEnabled) {
+        AttachmentsContentV2(
+            viewState = viewState,
+            attachmentsHandlers = attachmentsHandlers,
+            modifier = modifier,
+        )
+    } else {
+        AttachmentsContentV1(
+            viewState = viewState,
+            attachmentsHandlers = attachmentsHandlers,
+            modifier = modifier,
+        )
+    }
+}
+
+/**
+ * The top level content UI state for the [AttachmentsScreen] when viewing a content.
+ */
+@Suppress("LongMethod")
+@Composable
+private fun AttachmentsContentV1(
     viewState: AttachmentsState.ViewState.Content,
     attachmentsHandlers: AttachmentsHandlers,
     modifier: Modifier = Modifier,
@@ -96,7 +124,7 @@ fun AttachmentsContent(
             Text(
                 text = viewState
                     .newAttachment
-                    ?.displayName
+                    ?.completeFileName
                     ?: stringResource(id = BitwardenString.no_file_chosen),
                 color = BitwardenTheme.colorScheme.text.secondary,
                 style = BitwardenTheme.typography.bodySmall,
@@ -138,6 +166,116 @@ fun AttachmentsContent(
     }
 }
 
+/**
+ * The top level content UI state for the [AttachmentsScreen] when viewing a content.
+ */
+@Suppress("LongMethod")
+@Composable
+private fun AttachmentsContentV2(
+    viewState: AttachmentsState.ViewState.Content,
+    attachmentsHandlers: AttachmentsHandlers,
+    modifier: Modifier = Modifier,
+) {
+    LazyColumn(modifier = modifier) {
+        item {
+            Spacer(modifier = Modifier.height(height = 12.dp))
+            BitwardenListHeaderText(
+                label = stringResource(id = BitwardenString.attachments),
+                modifier = Modifier
+                    .fillMaxWidth()
+                    .standardHorizontalMargin()
+                    .padding(horizontal = 16.dp),
+            )
+            Spacer(modifier = Modifier.height(height = 8.dp))
+        }
+
+        if (viewState.attachments.isEmpty()) {
+            item {
+                Box(
+                    contentAlignment = Alignment.CenterStart,
+                    modifier = Modifier
+                        .fillMaxWidth()
+                        .standardHorizontalMargin()
+                        .defaultMinSize(minHeight = 60.dp)
+                        .cardStyle(cardStyle = CardStyle.Full),
+                ) {
+                    Text(
+                        text = stringResource(id = BitwardenString.no_attachments),
+                        style = BitwardenTheme.typography.bodyLarge,
+                        color = BitwardenTheme.colorScheme.text.secondary,
+                        modifier = Modifier
+                            .fillMaxWidth()
+                            .padding(horizontal = 16.dp)
+                            .testTag(tag = "NoAttachmentsLabel"),
+                    )
+                }
+            }
+        } else {
+            itemsIndexed(items = viewState.attachments) { index, attachment ->
+                AttachmentListEntry(
+                    attachmentItem = attachment,
+                    onDeleteClick = attachmentsHandlers.onDeleteClick,
+                    onItemClick = attachmentsHandlers.onItemClick,
+                    cardStyle = viewState.attachments.toListItemCardStyle(index = index),
+                    modifier = Modifier
+                        .fillMaxWidth()
+                        .standardHorizontalMargin()
+                        .testTag(tag = "AttachmentList"),
+                )
+            }
+        }
+
+        item {
+            Spacer(modifier = Modifier.height(height = 16.dp))
+            BitwardenListHeaderText(
+                label = stringResource(id = BitwardenString.add_new_attachment),
+                modifier = Modifier
+                    .fillMaxWidth()
+                    .standardHorizontalMargin()
+                    .padding(horizontal = 16.dp),
+            )
+            Spacer(modifier = Modifier.height(height = 8.dp))
+        }
+
+        item {
+            BitwardenTextField(
+                label = stringResource(id = BitwardenString.file_name),
+                value = viewState.newAttachment?.displayName
+                    ?: stringResource(id = BitwardenString.no_file_chosen),
+                textFieldTestTag = "SelectedFileNameLabel",
+                onValueChange = attachmentsHandlers.onFileNameChange,
+                enabled = viewState.newAttachment != null,
+                supportingText = stringResource(id = BitwardenString.max_file_size),
+                visualTransformation = nonEditableExtensionVisualTransformation(
+                    fileExtension = viewState.newAttachment?.extension,
+                ),
+                cardStyle = CardStyle.Full,
+                modifier = Modifier
+                    .fillMaxWidth()
+                    .standardHorizontalMargin(),
+            )
+        }
+
+        item {
+            Spacer(modifier = Modifier.height(height = 8.dp))
+            BitwardenOutlinedButton(
+                label = stringResource(id = BitwardenString.choose_file),
+                onClick = attachmentsHandlers.onChooseFileClick,
+                isExternalLink = true,
+                modifier = Modifier
+                    .fillMaxWidth()
+                    .standardHorizontalMargin()
+                    .testTag(tag = "AttachmentSelectFileButton"),
+            )
+        }
+
+        item {
+            Spacer(modifier = Modifier.height(height = 16.dp))
+            Spacer(modifier = Modifier.navigationBarsPadding())
+        }
+    }
+}
+
 @Composable
 private fun AttachmentListEntry(
     attachmentItem: AttachmentsState.AttachmentItem,

app/src/main/kotlin/com/x8bit/bitwarden/ui/vault/feature/attachments/AttachmentsScreen.kt

@@ -104,6 +104,7 @@ fun AttachmentsScreen(
             is AttachmentsState.ViewState.Content -> AttachmentsContent(
                 viewState = viewState,
                 attachmentsHandlers = attachmentsHandlers,
+                isAttachmentUpdatesEnabled = state.isAttachmentUpdatesEnabled,
                 modifier = Modifier.fillMaxSize(),
             )
 

app/src/main/kotlin/com/x8bit/bitwarden/ui/vault/feature/attachments/AttachmentsViewModel.kt

@@ -4,6 +4,7 @@ import android.net.Uri
 import android.os.Parcelable
 import androidx.lifecycle.SavedStateHandle
 import androidx.lifecycle.viewModelScope
+import com.bitwarden.core.data.manager.model.FlagKey
 import com.bitwarden.core.data.repository.model.DataState
 import com.bitwarden.ui.platform.base.BaseViewModel
 import com.bitwarden.ui.platform.components.snackbar.model.BitwardenSnackbarData
@@ -15,6 +16,7 @@ import com.bitwarden.ui.util.concat
 import com.bitwarden.vault.CipherView
 import com.x8bit.bitwarden.data.auth.repository.AuthRepository
 import com.x8bit.bitwarden.data.auth.repository.model.UserState
+import com.x8bit.bitwarden.data.platform.manager.FeatureFlagManager
 import com.x8bit.bitwarden.data.vault.repository.VaultRepository
 import com.x8bit.bitwarden.data.vault.repository.model.CreateAttachmentResult
 import com.x8bit.bitwarden.data.vault.repository.model.DeleteAttachmentResult
@@ -45,6 +47,7 @@ private const val MAX_FILE_SIZE_BYTES: Long = 100 * 1024 * 1024
 class AttachmentsViewModel @Inject constructor(
     private val authRepo: AuthRepository,
     private val vaultRepo: VaultRepository,
+    featureFlagManager: FeatureFlagManager,
     savedStateHandle: SavedStateHandle,
 ) : BaseViewModel<AttachmentsState, AttachmentsEvent, AttachmentsAction>(
     // We load the state from the savedStateHandle for testing purposes.
@@ -60,6 +63,9 @@ class AttachmentsViewModel @Inject constructor(
                 )
                     .takeUnless { isPremiumUser },
                 isPremiumUser = isPremiumUser,
+                isAttachmentUpdatesEnabled = featureFlagManager.getFeatureFlag(
+                    key = FlagKey.AttachmentUpdates,
+                ),
             )
         },
 ) {
@@ -75,6 +81,12 @@ class AttachmentsViewModel @Inject constructor(
             .map { AttachmentsAction.Internal.UserStateReceive(it) }
             .onEach(::sendAction)
             .launchIn(viewModelScope)
+
+        featureFlagManager
+            .getFeatureFlagFlow(key = FlagKey.AttachmentUpdates)
+            .map { AttachmentsAction.Internal.AttachmentUpdatesFlagReceive(it) }
+            .onEach(::sendAction)
+            .launchIn(viewModelScope)
     }
 
     override fun handleAction(action: AttachmentsAction) {
@@ -83,6 +95,7 @@ class AttachmentsViewModel @Inject constructor(
             AttachmentsAction.SaveClick -> handleSaveClick()
             AttachmentsAction.DismissDialogClick -> handleDismissDialogClick()
             AttachmentsAction.ChooseFileClick -> handleChooseFileClick()
+            is AttachmentsAction.FileNameChange -> handleFileNameChange(action)
             is AttachmentsAction.FileChoose -> handleFileChoose(action)
             is AttachmentsAction.DeleteClick -> handleDeleteClick(action)
             is AttachmentsAction.ItemClick -> handleItemClick(action)
@@ -145,7 +158,7 @@ class AttachmentsViewModel @Inject constructor(
                     cipherId = state.cipherId,
                     cipherView = requireNotNull(content.originalCipher),
                     fileSizeBytes = content.newAttachment.sizeBytes.toString(),
-                    fileName = content.newAttachment.displayName,
+                    fileName = content.newAttachment.completeFileName,
                     fileUri = content.newAttachment.uri,
                 )
                 sendAction(AttachmentsAction.Internal.CreateAttachmentResultReceive(result))
@@ -161,12 +174,31 @@ class AttachmentsViewModel @Inject constructor(
         sendEvent(AttachmentsEvent.ShowChooserSheet)
     }
 
+    private fun handleFileNameChange(action: AttachmentsAction.FileNameChange) {
+        onContent { content ->
+            mutableStateFlow.update {
+                it.copy(
+                    viewState = content.copy(
+                        newAttachment = content.newAttachment?.copy(
+                            displayName = action.fileName,
+                        ),
+                    ),
+                )
+            }
+        }
+    }
+
     private fun handleFileChoose(action: AttachmentsAction.FileChoose) {
         updateContent {
             it.copy(
                 newAttachment = AttachmentsState.NewAttachment(
                     uri = action.fileData.uri,
-                    displayName = action.fileData.fileName,
+                    extension = action
+                        .fileData
+                        .fileName
+                        .substringAfterLast(delimiter = '.', missingDelimiterValue = "")
+                        .takeUnless { extension -> extension.isBlank() },
+                    displayName = action.fileData.fileName.substringBeforeLast(delimiter = '.'),
                     sizeBytes = action.fileData.sizeBytes,
                 ),
             )
@@ -213,6 +245,9 @@ class AttachmentsViewModel @Inject constructor(
 
             is AttachmentsAction.Internal.DeleteResultReceive -> handleDeleteResultReceive(action)
             is AttachmentsAction.Internal.UserStateReceive -> handleUserStateReceive(action)
+            is AttachmentsAction.Internal.AttachmentUpdatesFlagReceive -> {
+                handleAttachmentUpdatesFlagReceive(action)
+            }
         }
     }
 
@@ -337,6 +372,12 @@ class AttachmentsViewModel @Inject constructor(
         }
     }
 
+    private fun handleAttachmentUpdatesFlagReceive(
+        action: AttachmentsAction.Internal.AttachmentUpdatesFlagReceive,
+    ) {
+        mutableStateFlow.update { it.copy(isAttachmentUpdatesEnabled = action.isEnabled) }
+    }
+
     private inline fun onContent(
         crossinline block: (AttachmentsState.ViewState.Content) -> Unit,
     ) {
@@ -365,6 +406,7 @@ data class AttachmentsState(
     val viewState: ViewState,
     val dialogState: DialogState?,
     val isPremiumUser: Boolean,
+    val isAttachmentUpdatesEnabled: Boolean,
 ) : Parcelable {
     /**
      * Represents the specific view states for the [AttachmentsScreen].
@@ -401,9 +443,13 @@ data class AttachmentsState(
     @Parcelize
     data class NewAttachment(
         val uri: Uri,
+        val extension: String?,
         val displayName: String,
         val sizeBytes: Long,
-    ) : Parcelable
+    ) : Parcelable {
+        val completeFileName: String
+            get() = extension?.let { "$displayName.$it" } ?: displayName
+    }
 
     /**
      * Represents an individual attachment that is already saved to the cipher.
@@ -508,6 +554,11 @@ sealed class AttachmentsAction {
      */
     data object ChooseFileClick : AttachmentsAction()
 
+    /**
+     * User edited the new attachment file name.
+     */
+    data class FileNameChange(val fileName: String) : AttachmentsAction()
+
     /**
      * User has chosen the file attachment.
      */
@@ -533,6 +584,13 @@ sealed class AttachmentsAction {
      * Internal ViewModel actions.
      */
     sealed class Internal : AttachmentsAction() {
+        /**
+         * Updates about the state of the attachment updates flag have been received.
+         */
+        data class AttachmentUpdatesFlagReceive(
+            val isEnabled: Boolean,
+        ) : Internal()
+
         /**
          * The cipher data has been received.
          */

app/src/main/kotlin/com/x8bit/bitwarden/ui/vault/feature/attachments/handlers/AttachmentsHandlers.kt

@@ -16,6 +16,7 @@ data class AttachmentsHandlers(
     val onDeleteClick: (attachmentId: String) -> Unit,
     val onItemClick: (attachment: AttachmentsState.AttachmentItem) -> Unit,
     val onDismissRequest: () -> Unit,
+    val onFileNameChange: (String) -> Unit,
 ) {
     @Suppress("UndocumentedPublicClass")
     companion object {
@@ -34,6 +35,9 @@ data class AttachmentsHandlers(
                 onDismissRequest = {
                     viewModel.trySendAction(AttachmentsAction.DismissDialogClick)
                 },
+                onFileNameChange = {
+                    viewModel.trySendAction(AttachmentsAction.FileNameChange(it))
+                },
             )
     }
 }

app/src/test/kotlin/com/x8bit/bitwarden/ui/vault/feature/attachments/AttachmentsScreenTest.kt

@@ -10,6 +10,7 @@ import androidx.compose.ui.test.onAllNodesWithText
 import androidx.compose.ui.test.onNodeWithContentDescription
 import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.performClick
+import androidx.compose.ui.test.performTextInput
 import com.bitwarden.core.data.repository.util.bufferedMutableSharedFlow
 import com.bitwarden.ui.platform.manager.IntentManager
 import com.bitwarden.ui.util.asText
@@ -103,6 +104,29 @@ class AttachmentsScreenTest : BitwardenComposeTest() {
         }
     }
 
+    @Test
+    fun `on edit file name should send FileNameChange`() {
+        mutableStateFlow.update {
+            it.copy(
+                viewState = DEFAULT_CONTENT_WITHOUT_ATTACHMENTS.copy(
+                    newAttachment = AttachmentsState.NewAttachment(
+                        extension = "png",
+                        displayName = "cool_file",
+                        uri = mockk(),
+                        sizeBytes = 100L,
+                    ),
+                ),
+            )
+        }
+        composeTestRule
+            .onNodeWithText("cool_file")
+            .performTextInput("5")
+
+        verify(exactly = 1) {
+            viewModel.trySendAction(AttachmentsAction.FileNameChange(fileName = "cool_file5"))
+        }
+    }
+
     @Test
     fun `progressbar should be displayed according to state`() {
         mutableStateFlow.update { it.copy(viewState = AttachmentsState.ViewState.Loading) }
@@ -260,6 +284,7 @@ private val DEFAULT_STATE: AttachmentsState = AttachmentsState(
     viewState = AttachmentsState.ViewState.Loading,
     dialogState = null,
     isPremiumUser = false,
+    isAttachmentUpdatesEnabled = true,
 )
 
 private val DEFAULT_CONTENT_WITHOUT_ATTACHMENTS: AttachmentsState.ViewState.Content =

app/src/test/kotlin/com/x8bit/bitwarden/ui/vault/feature/attachments/AttachmentsViewModelTest.kt

@@ -3,6 +3,7 @@ package com.x8bit.bitwarden.ui.vault.feature.attachments
 import android.net.Uri
 import androidx.lifecycle.SavedStateHandle
 import app.cash.turbine.test
+import com.bitwarden.core.data.manager.model.FlagKey
 import com.bitwarden.core.data.repository.model.DataState
 import com.bitwarden.data.repository.model.Environment
 import com.bitwarden.ui.platform.base.BaseViewModelTest
@@ -15,6 +16,7 @@ import com.x8bit.bitwarden.data.auth.datasource.disk.model.OnboardingStatus
 import com.x8bit.bitwarden.data.auth.repository.AuthRepository
 import com.x8bit.bitwarden.data.auth.repository.model.UserState
 import com.x8bit.bitwarden.data.platform.error.NoActiveUserException
+import com.x8bit.bitwarden.data.platform.manager.FeatureFlagManager
 import com.x8bit.bitwarden.data.platform.manager.model.FirstTimeState
 import com.x8bit.bitwarden.data.vault.datasource.sdk.model.createMockCipherView
 import com.x8bit.bitwarden.data.vault.repository.VaultRepository
@@ -29,6 +31,7 @@ import io.mockk.mockkStatic
 import io.mockk.unmockkStatic
 import kotlinx.collections.immutable.persistentListOf
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.test.runTest
 import org.junit.jupiter.api.AfterEach
 import org.junit.jupiter.api.Assertions.assertEquals
@@ -46,6 +49,13 @@ class AttachmentsViewModelTest : BaseViewModelTest() {
     private val vaultRepository: VaultRepository = mockk {
         every { getVaultItemStateFlow(any()) } returns mutableVaultItemStateFlow
     }
+    private val mutableAttachmentUpdatesFlow = MutableStateFlow(true)
+    private val featureFlagManager: FeatureFlagManager = mockk {
+        every {
+            getFeatureFlag(FlagKey.AttachmentUpdates)
+        } answers { mutableAttachmentUpdatesFlow.value }
+        every { getFeatureFlagFlow(FlagKey.AttachmentUpdates) } returns mutableAttachmentUpdatesFlow
+    }
 
     @BeforeEach
     fun setup() {
@@ -105,6 +115,45 @@ class AttachmentsViewModelTest : BaseViewModelTest() {
         }
     }
 
+    @Test
+    fun `FileNameChange should update the newAttachment state`() = runTest {
+        val cipherView = createMockCipherView(number = 1)
+        mutableVaultItemStateFlow.value = DataState.Loaded(cipherView)
+        mutableUserStateFlow.value = DEFAULT_USER_STATE
+        val uri = mockk<Uri>()
+        val newAttachment = AttachmentsState.NewAttachment(
+            uri = uri,
+            extension = "png",
+            displayName = "cool_file",
+            sizeBytes = 100L,
+        )
+        val initialState = DEFAULT_STATE.copy(viewState = DEFAULT_CONTENT_WITH_ATTACHMENTS)
+
+        val fileData = FileData(
+            fileName = "cool_file.png",
+            uri = uri,
+            sizeBytes = 100L,
+        )
+        mutableVaultItemStateFlow.value = DataState.Loaded(cipherView)
+        mutableUserStateFlow.value = DEFAULT_USER_STATE
+
+        val viewModel = createViewModel()
+        // Need to populate the VM with a file
+        viewModel.trySendAction(AttachmentsAction.FileChoose(fileData))
+
+        // Then alter the name
+        viewModel.trySendAction(AttachmentsAction.FileNameChange(fileName = "cool_file5"))
+
+        assertEquals(
+            initialState.copy(
+                viewState = DEFAULT_CONTENT_WITH_ATTACHMENTS.copy(
+                    newAttachment = newAttachment.copy(displayName = "cool_file5"),
+                ),
+            ),
+            viewModel.stateFlow.value,
+        )
+    }
+
     @Test
     fun `SaveClick should display error dialog when user is not Premium`() = runTest {
         val cipherView = createMockCipherView(number = 1)
@@ -172,7 +221,8 @@ class AttachmentsViewModelTest : BaseViewModelTest() {
         val state = DEFAULT_STATE.copy(
             viewState = DEFAULT_CONTENT_WITH_ATTACHMENTS.copy(
                 newAttachment = AttachmentsState.NewAttachment(
-                    displayName = fileName,
+                    extension = "png",
+                    displayName = "test",
                     uri = uri,
                     sizeBytes = sizeToBig,
                 ),
@@ -217,7 +267,8 @@ class AttachmentsViewModelTest : BaseViewModelTest() {
             val state = DEFAULT_STATE.copy(
                 viewState = DEFAULT_CONTENT_WITH_ATTACHMENTS.copy(
                     newAttachment = AttachmentsState.NewAttachment(
-                        displayName = fileName,
+                        extension = "png",
+                        displayName = "test",
                         uri = uri,
                         sizeBytes = sizeJustRight,
                     ),
@@ -292,7 +343,8 @@ class AttachmentsViewModelTest : BaseViewModelTest() {
             val state = DEFAULT_STATE.copy(
                 viewState = DEFAULT_CONTENT_WITH_ATTACHMENTS.copy(
                     newAttachment = AttachmentsState.NewAttachment(
-                        displayName = fileName,
+                        extension = "png",
+                        displayName = "test",
                         uri = uri,
                         sizeBytes = sizeJustRight,
                     ),
@@ -363,7 +415,8 @@ class AttachmentsViewModelTest : BaseViewModelTest() {
         val state = DEFAULT_STATE.copy(
             viewState = DEFAULT_CONTENT_WITH_ATTACHMENTS.copy(
                 newAttachment = AttachmentsState.NewAttachment(
-                    displayName = fileName,
+                    extension = "png",
+                    displayName = "test",
                     uri = uri,
                     sizeBytes = sizeJustRight,
                 ),
@@ -435,7 +488,7 @@ class AttachmentsViewModelTest : BaseViewModelTest() {
     fun `ChooseFile should update state with new file data`() = runTest {
         val uri = createMockUri()
         val fileData = FileData(
-            fileName = "filename-1",
+            fileName = "filename-1.png",
             uri = uri,
             sizeBytes = 100L,
         )
@@ -449,6 +502,7 @@ class AttachmentsViewModelTest : BaseViewModelTest() {
             initialState.copy(
                 viewState = DEFAULT_CONTENT_WITH_ATTACHMENTS.copy(
                     newAttachment = AttachmentsState.NewAttachment(
+                        extension = "png",
                         displayName = "filename-1",
                         uri = uri,
                         sizeBytes = 100L,
@@ -693,11 +747,27 @@ class AttachmentsViewModelTest : BaseViewModelTest() {
         )
     }
 
+    @Test
+    fun `AttachmentUpdatesFlow should update isAttachmentUpdatesEnabled state`() = runTest {
+        val viewModel = createViewModel()
+
+        viewModel.stateFlow.test {
+            assertEquals(DEFAULT_STATE, awaitItem())
+
+            mutableAttachmentUpdatesFlow.update { false }
+            assertEquals(DEFAULT_STATE.copy(isAttachmentUpdatesEnabled = false), awaitItem())
+
+            mutableAttachmentUpdatesFlow.update { true }
+            assertEquals(DEFAULT_STATE, awaitItem())
+        }
+    }
+
     private fun createViewModel(
         initialState: AttachmentsState? = null,
     ): AttachmentsViewModel = AttachmentsViewModel(
         authRepo = authRepository,
         vaultRepo = vaultRepository,
+        featureFlagManager = featureFlagManager,
         savedStateHandle = SavedStateHandle().apply {
             set("state", initialState)
             every {
@@ -739,6 +809,7 @@ private val DEFAULT_STATE: AttachmentsState = AttachmentsState(
     viewState = AttachmentsState.ViewState.Loading,
     dialogState = null,
     isPremiumUser = true,
+    isAttachmentUpdatesEnabled = true,
 )
 
 private val DEFAULT_ATTACHMENT_ITEM: AttachmentsState.AttachmentItem =