From a3199009a5070f19de85e7dd6f4d8e7df93d2ef9 Mon Sep 17 00:00:00 2001
From: Ivan Sekovanikj <ivan.sekovanikj@getstream.io>
Date: Mon, 14 Oct 2024 20:48:44 +0200
Subject: [PATCH] feat: impl improved poll state management system

---
 package/src/components/Chat/Chat.tsx          |  2 ++
 .../Message/MessageSimple/MessageContent.tsx  | 11 +++++--
 package/src/components/Poll/Poll.tsx          | 29 +++----------------
 3 files changed, 15 insertions(+), 27 deletions(-)

diff --git a/package/src/components/Chat/Chat.tsx b/package/src/components/Chat/Chat.tsx
index 54ad2ac36..2e2999408 100644
--- a/package/src/components/Chat/Chat.tsx
+++ b/package/src/components/Chat/Chat.tsx
@@ -219,9 +219,11 @@ const ChatWithContext = <
     if (!client) return;
 
     client.threads.registerSubscriptions();
+    client.polls.registerSubscriptions();
 
     return () => {
       client.threads.unregisterSubscriptions();
+      client.polls.unregisterSubscriptions();
     };
   }, [client]);
 
diff --git a/package/src/components/Message/MessageSimple/MessageContent.tsx b/package/src/components/Message/MessageSimple/MessageContent.tsx
index c4085bbd1..ca859bb3e 100644
--- a/package/src/components/Message/MessageSimple/MessageContent.tsx
+++ b/package/src/components/Message/MessageSimple/MessageContent.tsx
@@ -9,6 +9,7 @@ import {
 
 import { MessageTextContainer } from './MessageTextContainer';
 
+import { useChatContext } from '../../../contexts';
 import {
   MessageContextValue,
   useMessageContext,
@@ -143,6 +144,7 @@ const MessageContentWithContext = <
     showMessageStatus,
     threadList,
   } = props;
+  const { client } = useChatContext();
 
   const {
     theme: {
@@ -379,8 +381,13 @@ const MessageContentWithContext = <
               case 'gallery':
                 return <Gallery key={`gallery_${messageContentOrderIndex}`} />;
               case 'poll':
-                return message.poll_id && message.poll ? (
-                  <Poll key={`poll_${message.poll_id}`} poll={message.poll} />
+                return message.poll_id &&
+                  message.poll &&
+                  client.polls.fromState(message.poll_id) ? (
+                  <Poll
+                    key={`poll_${message.poll_id}`}
+                    poll={client.polls.fromState(message.poll_id)}
+                  />
                 ) : null;
               case 'text':
               default:
diff --git a/package/src/components/Poll/Poll.tsx b/package/src/components/Poll/Poll.tsx
index 9b50c39e5..e8c2e7e87 100644
--- a/package/src/components/Poll/Poll.tsx
+++ b/package/src/components/Poll/Poll.tsx
@@ -1,12 +1,7 @@
-import React, { useCallback, useEffect, useMemo, useState } from 'react';
+import React, { useCallback, useMemo, useState } from 'react';
 import { Modal, SafeAreaView, Text, View } from 'react-native';
 
-import {
-  Poll as PollClass,
-  PollOption as PollOptionClass,
-  PollResponse,
-  PollState,
-} from 'stream-chat';
+import { Poll as PollClass, PollOption as PollOptionClass, PollState } from 'stream-chat';
 
 import {
   AddCommentButton,
@@ -21,12 +16,7 @@ import { PollInputDialog } from './components/PollInputDialog';
 import { PollOption, ShowAllOptionsContent } from './components/PollOption';
 import { PollResults } from './components/PollResults';
 
-import {
-  PollContextProvider,
-  useChatContext,
-  useMessageContext,
-  usePollContext,
-} from '../../contexts';
+import { PollContextProvider, useMessageContext, usePollContext } from '../../contexts';
 import { useStateStore } from '../../hooks';
 // import * as dbApi from '../../store/apis';
 
@@ -135,15 +125,9 @@ const PollWithContext = () => {
   );
 };
 
-export const Poll = ({ poll: pollData }: { poll: PollResponse }) => {
-  const { client } = useChatContext();
+export const Poll = ({ poll }: { poll: PollClass }) => {
   const { message } = useMessageContext();
 
-  const poll = useMemo<PollClass>(
-    () => new PollClass({ client, poll: pollData }),
-    [client, pollData],
-  );
-
   const [
     vote_counts_by_option,
     ownVotesByOptionId,
@@ -169,11 +153,6 @@ export const Poll = ({ poll: pollData }: { poll: PollResponse }) => {
   );
   const endVote = useCallback(() => poll.close(), [poll]);
 
-  useEffect(() => {
-    poll.registerSubscriptions();
-    return poll.unregisterSubscriptions;
-  }, [poll]);
-
   return (
     <PollContextProvider
       value={{