Fix draft not deleted when attachments are removed from the composer

Author: nuno-vieira

Sources/StreamChatSwiftUI/ChatChannel/Composer/MessageComposerViewModel.swift

@@ -30,6 +30,10 @@ open class MessageComposerViewModel: ObservableObject {
     @Published public private(set) var addedAssets = [AddedAsset]() {
         didSet {
             checkPickerSelectionState()
+
+            if shouldDeleteDraftMessage(oldValue: oldValue) {
+                deleteDraftMessage()
+            }
         }
     }
 
@@ -55,7 +59,7 @@ open class MessageComposerViewModel: ObservableObject {
                 suggestions = [String: Any]()
                 mentionedUsers = Set<ChatUser>()
 
-                if oldValue != "" && !sendButtonEnabled {
+                if shouldDeleteDraftMessage(oldValue: oldValue) {
                     deleteDraftMessage()
                 }
             }
@@ -71,18 +75,30 @@ open class MessageComposerViewModel: ObservableObject {
                 addedFileURLs.removeLast()
             }
             checkPickerSelectionState()
+
+            if shouldDeleteDraftMessage(oldValue: oldValue) {
+                deleteDraftMessage()
+            }
         }
     }
 
     @Published public var addedVoiceRecordings = [AddedVoiceRecording]() {
         didSet {
             checkPickerSelectionState()
+
+            if shouldDeleteDraftMessage(oldValue: oldValue) {
+                deleteDraftMessage()
+            }
         }
     }
 
     @Published public var addedCustomAttachments = [CustomAttachment]() {
         didSet {
             checkPickerSelectionState()
+
+            if shouldDeleteDraftMessage(oldValue: oldValue) {
+                deleteDraftMessage()
+            }
         }
     }
 
@@ -251,10 +267,10 @@ open class MessageComposerViewModel: ObservableObject {
         quotedMessage?.wrappedValue = message.quotedMessage
         showReplyInChannel = message.showReplyInChannel
 
-        addedAssets.removeAll()
-        addedFileURLs.removeAll()
-        addedVoiceRecordings.removeAll()
-        addedCustomAttachments.removeAll()
+        var addedAssets: [AddedAsset] = []
+        var addedFileURLs: [URL] = []
+        var addedVoiceRecordings: [AddedVoiceRecording] = []
+        var addedCustomAttachments: [CustomAttachment] = []
 
         message.attachments.forEach { attachment in
             switch attachment.type {
@@ -277,6 +293,11 @@ open class MessageComposerViewModel: ObservableObject {
                 addedCustomAttachments.append(customAttachment)
             }
         }
+
+        self.addedAssets = addedAssets
+        self.addedFileURLs = addedFileURLs
+        self.addedVoiceRecordings = addedVoiceRecordings
+        self.addedCustomAttachments = addedCustomAttachments
     }
 
     /// Updates the draft message locally and on the server.
@@ -318,7 +339,8 @@ open class MessageComposerViewModel: ObservableObject {
         )
     }
 
-    private func deleteDraftMessage() {
+    /// Deletes the draft message locally and on the server if it exists.
+    public func deleteDraftMessage() {
         guard draftMessage != nil else {
             return
         }
@@ -330,6 +352,11 @@ open class MessageComposerViewModel: ObservableObject {
         }
     }
 
+    /// Checks if the previous value of the content in the composer was not empty and the current value is empty.
+    private func shouldDeleteDraftMessage(oldValue: any Collection) -> Bool {
+        !oldValue.isEmpty && !sendButtonEnabled
+    }
+
     open func sendMessage(
         quotedMessage: ChatMessage?,
         editedMessage: ChatMessage?,

StreamChatSwiftUITests/Tests/ChatChannel/MessageComposerViewModel_Tests.swift

@@ -984,6 +984,101 @@ class MessageComposerViewModel_Tests: StreamChatTestCase {
         XCTAssertEqual(viewModel.text, "Draft")
     }
 
+    func test_messageComposerVM_whenLastAssetRemoved_shouldDeleteDraft() {
+        // Given
+        let channelController = makeChannelController()
+        let draftMessage = DraftMessage.mock(text: "")
+        channelController.channel_mock = .mock(cid: .unique, draftMessage: draftMessage)
+        let viewModel = makeComposerDraftsViewModel(
+            channelController: channelController,
+            messageController: nil
+        )
+        let asset = defaultAsset
+        viewModel.imageTapped(asset)
+
+        // When
+        viewModel.imageTapped(asset) // Remove the asset by tapping again
+
+        // Then
+        XCTAssertEqual(channelController.deleteDraftMessage_callCount, 1)
+    }
+
+    func test_messageComposerVM_whenLastFileRemoved_shouldDeleteDraft() {
+        // Given
+        let channelController = makeChannelController()
+        let draftMessage = DraftMessage.mock(text: "")
+        channelController.channel_mock = .mock(cid: .unique, draftMessage: draftMessage)
+        let viewModel = makeComposerDraftsViewModel(
+            channelController: channelController,
+            messageController: nil
+        )
+        viewModel.addedFileURLs = [mockURL]
+
+        // When
+        viewModel.removeAttachment(with: mockURL.absoluteString)
+
+        // Then
+        XCTAssertEqual(channelController.deleteDraftMessage_callCount, 1)
+    }
+
+    func test_messageComposerVM_whenLastVoiceRecordingRemoved_shouldDeleteDraft() {
+        // Given
+        let channelController = makeChannelController()
+        let draftMessage = DraftMessage.mock(text: "")
+        channelController.channel_mock = .mock(cid: .unique, draftMessage: draftMessage)
+        let viewModel = makeComposerDraftsViewModel(
+            channelController: channelController,
+            messageController: nil
+        )
+        let recording = AddedVoiceRecording(url: mockURL, duration: 1.0, waveform: [0.5])
+        viewModel.addedVoiceRecordings = [recording]
+
+        // When
+        viewModel.removeAttachment(with: mockURL.absoluteString)
+
+        // Then
+        XCTAssertEqual(channelController.deleteDraftMessage_callCount, 1)
+    }
+
+    func test_messageComposerVM_whenLastCustomAttachmentRemoved_shouldDeleteDraft() {
+        // Given
+        let channelController = makeChannelController()
+        let draftMessage = DraftMessage.mock(text: "")
+        channelController.channel_mock = .mock(cid: .unique, draftMessage: draftMessage)
+        let viewModel = makeComposerDraftsViewModel(
+            channelController: channelController,
+            messageController: nil
+        )
+        let attachment = CustomAttachment(id: .unique, content: .mockFile)
+        viewModel.customAttachmentTapped(attachment)
+
+        // When
+        viewModel.customAttachmentTapped(attachment) // Remove by tapping again
+
+        // Then
+        XCTAssertEqual(channelController.deleteDraftMessage_callCount, 1)
+    }
+
+    func test_messageComposerVM_whenRemovingAttachment_withTextPresent_shouldNotDeleteDraft() {
+        // Given
+        let channelController = makeChannelController()
+        let draftMessage = DraftMessage.mock(text: "Hello")
+        channelController.channel_mock = .mock(cid: .unique, draftMessage: draftMessage)
+        let viewModel = makeComposerDraftsViewModel(
+            channelController: channelController,
+            messageController: nil
+        )
+        viewModel.text = "Hello"
+        let asset = defaultAsset
+        viewModel.imageTapped(asset)
+
+        // When
+        viewModel.imageTapped(asset) // Remove the asset by tapping again
+
+        // Then
+        XCTAssertEqual(channelController.deleteDraftMessage_callCount, 0)
+    }
+
     // MARK: - private
 
     private func makeComposerDraftsViewModel(