Fix race between local updates and setting up API request waiting in send message tests in Chat (#3887)

Author: laevandus

Sources/StreamChat/StateLayer/Chat.swift

@@ -406,8 +406,9 @@ public class Chat {
     /// - Returns: An instance of `ChatMessage` which was resent.
     @discardableResult public func resendMessage(_ messageId: MessageId) async throws -> ChatMessage {
         let messageSender = try client.backgroundWorker(of: MessageSender.self)
+        async let resentMessage = try await messageSender.waitForAPIRequest(messageId: messageId)
         try await messageUpdater.resendMessage(with: messageId)
-        return try await messageSender.waitForAPIRequest(messageId: messageId)
+        return try await resentMessage
     }
 
     /// Downloads the specified attachment and stores it locally on the device.
@@ -509,9 +510,11 @@ public class Chat {
     ) async throws -> ChatMessage {
         Task { try await stopTyping() } // errors explicitly ignored
         let cid = try await self.cid
+        let newMessageId = messageId ?? .newUniqueId
+        async let sentMessage = waitForMessageSender(messageId: newMessageId)
         let localMessage = try await channelUpdater.createNewMessage(
             in: cid,
-            messageId: messageId,
+            messageId: newMessageId,
             text: text,
             pinning: pinning,
             isSilent: silent,
@@ -526,8 +529,6 @@ public class Chat {
             restrictedVisibility: restrictedVisibility,
             extraData: extraData
         )
-        // Important to set up the waiter immediately
-        async let sentMessage = try await waitForAPIRequest(localMessage: localMessage)
         eventNotificationCenter.process(NewMessagePendingEvent(message: localMessage, cid: cid))
         return try await sentMessage
     }
@@ -547,9 +548,11 @@ public class Chat {
         extraData: [String: RawJSON] = [:]
     ) async throws -> ChatMessage {
         let cid = try await self.cid
+        let newMessageId = messageId ?? .newUniqueId
+        async let sentMessage = waitForMessageSender(messageId: newMessageId)
         let localMessage = try await channelUpdater.createNewMessage(
             in: cid,
-            messageId: messageId,
+            messageId: newMessageId,
             text: text,
             pinning: nil,
             isSilent: false,
@@ -564,8 +567,6 @@ public class Chat {
             restrictedVisibility: restrictedVisibility,
             extraData: extraData
         )
-        // Important to set up the waiter immediately
-        async let sentMessage = try await waitForAPIRequest(localMessage: localMessage)
         eventNotificationCenter.process(NewMessagePendingEvent(message: localMessage, cid: cid))
         return try await sentMessage
     }
@@ -593,7 +594,8 @@ public class Chat {
         skipPush: Bool = false
     ) async throws -> ChatMessage {
         Task { try await stopTyping() } // errors explicitly ignored
-        let localMessage = try await messageUpdater.editMessage(
+        async let updatedMessage = waitForMessageEditor(messageId: messageId)
+        _ = try await messageUpdater.editMessage(
             messageId: messageId,
             text: text,
             skipEnrichUrl: skipEnrichURL,
@@ -602,7 +604,7 @@ public class Chat {
             restrictedVisibility: restrictedVisibility,
             extraData: extraData
         )
-        return try await waitForAPIRequest(localMessage: localMessage)
+        return try await updatedMessage
     }
 
     // MARK: - Message Pagination
@@ -992,9 +994,11 @@ public class Chat {
     ) async throws -> ChatMessage {
         Task { try await stopTyping() } // errors explicitly ignored
         let cid = try await self.cid
+        let newMessageId = messageId ?? .newUniqueId
+        async let sentMessage = waitForMessageSender(messageId: newMessageId)
         let localMessage = try await messageUpdater.createNewReply(
             in: cid,
-            messageId: messageId,
+            messageId: newMessageId,
             text: text,
             pinning: pinning,
             command: nil,
@@ -1009,7 +1013,6 @@ public class Chat {
             skipEnrichUrl: skipEnrichURL,
             extraData: extraData
         )
-        async let sentMessage = try await waitForAPIRequest(localMessage: localMessage)
         eventNotificationCenter.process(NewMessagePendingEvent(message: localMessage, cid: cid))
         return try await sentMessage
     }
@@ -1536,18 +1539,14 @@ extension Chat {
         }
     }
 
-    /// Depending on the local state we use different workers.
-    func waitForAPIRequest(localMessage: ChatMessage) async throws -> ChatMessage {
-        switch localMessage.localState {
-        case .pendingSend:
-            let messageSender = try client.backgroundWorker(of: MessageSender.self)
-            return try await messageSender.waitForAPIRequest(messageId: localMessage.id)
-        case .pendingSync:
-            let messageEditor = try client.backgroundWorker(of: MessageEditor.self)
-            return try await messageEditor.waitForAPIRequest(messageId: localMessage.id)
-        default:
-            return localMessage
-        }
+    func waitForMessageEditor(messageId: MessageId) async throws -> ChatMessage {
+        let messageEditor = try client.backgroundWorker(of: MessageEditor.self)
+        return try await messageEditor.waitForAPIRequest(messageId: messageId)
+    }
+    
+    func waitForMessageSender(messageId: MessageId) async throws -> ChatMessage {
+        let messageSender = try client.backgroundWorker(of: MessageSender.self)
+        return try await messageSender.waitForAPIRequest(messageId: messageId)
     }
 }
 

Tests/StreamChatTests/StateLayer/Chat_Tests.swift

@@ -673,12 +673,6 @@ final class Chat_Tests: XCTestCase {
     func test_sendMessage_whenAPIRequestSucceeds_thenSendMessageSucceeds() async throws {
         try await setUpChat(usesMockedUpdaters: false)
         await XCTAssertEqual(0, chat.state.messages.count)
-
-        let notificationExpectation = expectation(
-            forNotification: .NewEventReceived,
-            object: nil,
-            notificationCenter: env.client.eventNotificationCenter
-        )
 
         let typingIndicatorResponse = EmptyResponse()
         env.client.mockAPIClient.test_mockResponseResult(.success(typingIndicatorResponse))
@@ -696,8 +690,6 @@ final class Chat_Tests: XCTestCase {
             messageId: apiResponse.message.id
         )
 
-        await fulfillment(of: [notificationExpectation], timeout: defaultTimeout)
-        
         XCTAssertEqual(text, message.text)
         await XCTAssertEqual(1, chat.state.messages.count)
         let messages = await chat.state.messages
@@ -735,21 +727,29 @@ final class Chat_Tests: XCTestCase {
         XCTAssertEqual(LocalMessageState.sendingFailed, stateMessage.localState)
     }
 
-    func test_sendSystemMessage_thenCallsChannelUpdaterWithIsSystemTrue() async throws {
-        try await setUpChat(usesMockedUpdaters: true)
+    func test_sendSystemMessage_whenAPIRequestSucceeds_thenSendMessageSucceeds() async throws {
+        try await setUpChat(usesMockedUpdaters: false)
         await XCTAssertEqual(0, chat.state.messages.count)
 
         let text = "Text"
-        env.channelUpdaterMock.createNewMessage_completion_result = .success((
-            .mock(id: "0", text: text, type: .system)
-        ))
+        let apiResponse = MessagePayload.Boxed(
+            message: .dummy(
+                type: .system,
+                messageId: "0",
+                text: text
+            )
+        )
+        env.client.mockAPIClient.test_mockResponseResult(.success(apiResponse))
+        
         let message = try await chat.sendSystemMessage(
             with: text,
             messageId: "0"
         )
 
-        XCTAssertEqual(env.typingEventsSenderMock.stopTyping_cid, nil)
-        XCTAssertEqual(env.channelUpdaterMock.createNewMessage_isSystem, true)
+        let body = env.client.mockAPIClient.request_endpoint?.body?.encodable as? [String: AnyEncodable]
+        let messageRequestBody = body?["message"]?.encodable as? MessageRequestBody
+        XCTAssertEqual(messageRequestBody?.type, MessageType.system.rawValue)
+        
         XCTAssertEqual(text, message.text)
         XCTAssertEqual(.system, message.type)
     }