Fix poll data overridden by message.updated event

CHANGELOG.md

@@ -34,6 +34,7 @@
 
 ## stream-chat-android-state
 ### 🐞 Fixed
+- Fix poll state getting overridden by `message.new` events. [#5963](https://github.com/GetStream/stream-chat-android/pull/5963)
 
 ### ⬆️ Improved
 

stream-chat-android-client/src/main/java/io/getstream/chat/android/client/extensions/internal/Thread.kt

@@ -44,8 +44,11 @@ public fun Thread.updateParentOrReply(message: Message): Thread {
 public fun Thread.updateParent(parent: Message): Thread {
     // Skip update if [parent] is not related to this Thread
     if (this.parentMessageId != parent.id) return this
+    // Enrich the poll as it might not be present in the event
+    val poll = parent.poll ?: this.parentMessage.poll
+    val updatedParent = parent.copy(poll = poll)
     return this.copy(
-        parentMessage = parent,
+        parentMessage = updatedParent,
         deletedAt = parent.deletedAt,
         updatedAt = parent.updatedAt ?: this.updatedAt,
     )

stream-chat-android-client/src/test/java/io/getstream/chat/android/client/extensions/internal/ThreadExtensionsTests.kt

@@ -18,11 +18,15 @@ package io.getstream.chat.android.client.extensions.internal
 
 import io.getstream.chat.android.models.ChannelUserRead
 import io.getstream.chat.android.models.Message
+import io.getstream.chat.android.models.Option
+import io.getstream.chat.android.models.Poll
 import io.getstream.chat.android.models.Thread
 import io.getstream.chat.android.models.ThreadInfo
 import io.getstream.chat.android.models.ThreadParticipant
 import io.getstream.chat.android.models.User
+import io.getstream.chat.android.models.VotingVisibility
 import org.amshove.kluent.shouldBeEqualTo
+import org.amshove.kluent.shouldNotBeNull
 import org.junit.Test
 import java.util.Date
 
@@ -157,6 +161,139 @@ internal class ThreadExtensionsTests {
         result shouldBeEqualTo baseThread
     }
 
+    @Test
+    fun `updateParent should preserve existing poll when updated parent has no poll`() {
+        // given
+        val originalPoll = Poll(
+            id = "poll1",
+            name = "Test Poll",
+            description = "Test Description",
+            options = listOf(
+                Option(id = "option1", text = "Option 1"),
+                Option(id = "option2", text = "Option 2"),
+            ),
+            votingVisibility = VotingVisibility.PUBLIC,
+            enforceUniqueVote = true,
+            maxVotesAllowed = 1,
+            allowUserSuggestedOptions = false,
+            allowAnswers = false,
+            voteCountsByOption = emptyMap(),
+            votes = emptyList(),
+            ownVotes = emptyList(),
+            createdAt = now,
+            updatedAt = now,
+            closed = false,
+        )
+        val threadWithPoll = baseThread.copy(
+            parentMessage = parentMessage.copy(poll = originalPoll),
+        )
+        val updatedParent = parentMessage.copy(
+            text = "Updated parent message",
+            poll = null,
+        )
+
+        // when
+        val result = threadWithPoll.updateParent(updatedParent)
+
+        // then
+        result.parentMessage.poll.shouldNotBeNull()
+        result.parentMessage.poll shouldBeEqualTo originalPoll
+    }
+
+    @Test
+    fun `updateParent should use new poll when updated parent has poll`() {
+        // given
+        val originalPoll = Poll(
+            id = "poll1",
+            name = "Original Poll",
+            description = "Original Description",
+            options = listOf(
+                Option(id = "option1", text = "Option 1"),
+            ),
+            votingVisibility = VotingVisibility.PUBLIC,
+            enforceUniqueVote = true,
+            maxVotesAllowed = 1,
+            allowUserSuggestedOptions = false,
+            allowAnswers = false,
+            voteCountsByOption = emptyMap(),
+            votes = emptyList(),
+            ownVotes = emptyList(),
+            createdAt = now,
+            updatedAt = now,
+            closed = false,
+        )
+        val newPoll = Poll(
+            id = "poll2",
+            name = "Updated Poll",
+            description = "Updated Description",
+            options = listOf(
+                Option(id = "option1", text = "Option 1"),
+                Option(id = "option2", text = "Option 2"),
+            ),
+            votingVisibility = VotingVisibility.ANONYMOUS,
+            enforceUniqueVote = false,
+            maxVotesAllowed = 2,
+            allowUserSuggestedOptions = true,
+            allowAnswers = true,
+            voteCountsByOption = mapOf("option1" to 5),
+            votes = emptyList(),
+            ownVotes = emptyList(),
+            createdAt = now,
+            updatedAt = now,
+            closed = true,
+        )
+        val threadWithPoll = baseThread.copy(
+            parentMessage = parentMessage.copy(poll = originalPoll),
+        )
+        val updatedParent = parentMessage.copy(
+            text = "Updated parent message",
+            poll = newPoll,
+        )
+
+        // when
+        val result = threadWithPoll.updateParent(updatedParent)
+
+        // then
+        result.parentMessage.poll.shouldNotBeNull()
+        result.parentMessage.poll shouldBeEqualTo newPoll
+    }
+
+    @Test
+    fun `updateParent should add poll when original parent has no poll but updated parent has poll`() {
+        // given
+        val newPoll = Poll(
+            id = "poll1",
+            name = "New Poll",
+            description = "New Description",
+            options = listOf(
+                Option(id = "option1", text = "Option 1"),
+                Option(id = "option2", text = "Option 2"),
+            ),
+            votingVisibility = VotingVisibility.PUBLIC,
+            enforceUniqueVote = true,
+            maxVotesAllowed = 1,
+            allowUserSuggestedOptions = false,
+            allowAnswers = false,
+            voteCountsByOption = emptyMap(),
+            votes = emptyList(),
+            ownVotes = emptyList(),
+            createdAt = now,
+            updatedAt = now,
+            closed = false,
+        )
+        val updatedParent = parentMessage.copy(
+            text = "Updated parent message",
+            poll = newPoll,
+        )
+
+        // when
+        val result = baseThread.updateParent(updatedParent)
+
+        // then
+        result.parentMessage.poll.shouldNotBeNull()
+        result.parentMessage.poll shouldBeEqualTo newPoll
+    }
+
     @Test
     fun `upsertReply should add new reply and update related fields`() {
         // given

stream-chat-android-state/src/main/java/io/getstream/chat/android/state/event/handler/internal/EventHandlerSequential.kt

@@ -572,7 +572,9 @@ internal class EventHandlerSequential(
                         }
                 }
                 is MessageUpdatedEvent -> {
-                    val enrichedMessage = event.message.enrichWithOwnReactions(batch, currentUserId, event.user)
+                    val enrichedMessage = event.message
+                        .enrichWithOwnReactions(batch, currentUserId, event.user)
+                        .enrichWithOwnPoll(batch)
                     batch.addMessageData(
                         event.createdAt,
                         event.cid,
@@ -907,6 +909,11 @@ internal class EventHandlerSequential(
         },
     )
 
+    private fun Message.enrichWithOwnPoll(batch: EventBatchUpdate): Message {
+        val localMessage = batch.getCurrentMessage(id)
+        return copy(poll = poll ?: localMessage?.poll)
+    }
+
     private infix fun UserId.mustBe(currentUserId: UserId?) {
         if (this != currentUserId) {
             throw InputMismatchException(

stream-chat-android-state/src/main/java/io/getstream/chat/android/state/plugin/logic/channel/internal/ChannelLogic.kt

@@ -549,10 +549,14 @@ internal class ChannelLogic(
                     }
 
                     is MessageUpdatedEvent -> {
+                        val originalMessage = mutableState.getMessageById(event.message.id)
+                        // Enrich the poll as it might not be present in the event
+                        val poll = event.message.poll ?: originalMessage?.poll
                         event.message.copy(
                             replyTo = event.message.replyMessageId
                                 ?.let { mutableState.getMessageById(it) }
                                 ?: event.message.replyTo,
+                            poll = poll,
                         ).let(::upsertEventMessage)
                         channelStateLogic.toggleHidden(false)
                     }

stream-chat-android-state/src/main/java/io/getstream/chat/android/state/plugin/logic/channel/thread/internal/ThreadLogic.kt

@@ -94,11 +94,15 @@ internal class ThreadLogic(
     internal fun handleMessageEvents(events: List<HasMessage>) {
         val messages = events
             .map { event ->
-                val ownReactions = getMessage(event.message.id)?.ownReactions ?: event.message.ownReactions
+                val originalMessage = getMessage(event.message.id)
+                val ownReactions = originalMessage?.ownReactions ?: event.message.ownReactions
+                // Enrich the poll as might not be present in the event
+                val poll = event.message.poll ?: originalMessage?.poll
                 if (event is MessageUpdatedEvent) {
                     event.message.copy(
                         replyTo = mutableState.messages.value.firstOrNull { it.id == event.message.replyMessageId },
                         ownReactions = ownReactions,
+                        poll = poll,
                     )
                 } else {
                     event.message.copy(

stream-chat-android-state/src/test/java/io/getstream/chat/android/state/channel/controller/WhenHandleEvent.kt

@@ -148,6 +148,78 @@ internal class WhenHandleEvent : SynchronizedCoroutineTest {
         verify(channelStateLogic, times(2)).upsertMessage(messageUpdateEvent.message)
     }
 
+    @Test
+    fun `when message update event arrives, channel should be toggled to not hidden`() = runTest {
+        val message = randomMessage(
+            id = randomString(),
+            user = User(id = "otherUserId"),
+            silent = false,
+            showInChannel = true,
+        )
+        channelMutableState.setMessages(listOf(message))
+
+        val messageUpdateEvent = randomMessageUpdateEvent(message = message)
+
+        channelLogic.handleEvent(messageUpdateEvent)
+
+        verify(channelStateLogic).toggleHidden(false)
+    }
+
+    @Test
+    fun `when message update event arrives without poll but original message has poll, poll should be preserved`() = runTest {
+        val poll = randomPoll()
+        val originalMessage = randomMessage(
+            id = randomString(),
+            user = User(id = "otherUserId"),
+            silent = false,
+            showInChannel = true,
+            poll = poll,
+        )
+        channelMutableState.setMessages(listOf(originalMessage))
+
+        val updatedMessageWithoutPoll = originalMessage.copy(
+            text = "Updated text",
+            poll = null,
+        )
+        val messageUpdateEvent = randomMessageUpdateEvent(message = updatedMessageWithoutPoll)
+
+        channelLogic.handleEvent(messageUpdateEvent)
+
+        verify(channelStateLogic).upsertMessage(
+            org.mockito.kotlin.argThat { message ->
+                message.poll == poll && message.text == "Updated text"
+            },
+        )
+    }
+
+    @Test
+    fun `when message update event arrives with poll, event poll should be used`() = runTest {
+        val originalPoll = randomPoll(id = "poll1")
+        val originalMessage = randomMessage(
+            id = randomString(),
+            user = User(id = "otherUserId"),
+            silent = false,
+            showInChannel = true,
+            poll = originalPoll,
+        )
+        channelMutableState.setMessages(listOf(originalMessage))
+
+        val eventPoll = randomPoll(id = "poll2")
+        val updatedMessage = originalMessage.copy(
+            text = "Updated text",
+            poll = eventPoll,
+        )
+        val messageUpdateEvent = randomMessageUpdateEvent(message = updatedMessage)
+
+        channelLogic.handleEvent(messageUpdateEvent)
+
+        verify(channelStateLogic).upsertMessage(
+            org.mockito.kotlin.argThat { message ->
+                message.poll == eventPoll && message.text == "Updated text"
+            },
+        )
+    }
+
     // Member added event
     @Test
     fun `when member is added, it should be propagated`(): Unit = runTest {