Fix poll votes list not updating on poll events

Sources/StreamChat/Controllers/PollController/PollVoteListController+Combine.swift

@@ -11,6 +11,11 @@ extension PollVoteListController {
         basePublishers.state.keepAlive(self)
     }
 
+    /// A publisher emitting a new value every time the votes change.
+    public var pollPublisher: AnyPublisher<Poll, Never> {
+        basePublishers.poll.keepAlive(self)
+    }
+
     /// A publisher emitting a new value every time the votes change.
     public var voteChangesPublisher: AnyPublisher<[ListChange<PollVote>], Never> {
         basePublishers.voteChanges.keepAlive(self)
@@ -26,6 +31,9 @@ extension PollVoteListController {
         /// A backing subject for `statePublisher`.
         let state: CurrentValueSubject<DataController.State, Never>
 
+        /// A backing subject for `pollPublisher`.
+        let poll: PassthroughSubject<Poll, Never> = .init()
+
         /// A backing subject for `voteChangesPublisher`.
         let voteChanges: PassthroughSubject<[ListChange<PollVote>], Never> = .init()
 
@@ -49,4 +57,8 @@ extension PollVoteListController.BasePublishers: PollVoteListControllerDelegate
     ) {
         voteChanges.send(changes)
     }
+
+    func controller(_ controller: PollVoteListController, didUpdatePoll poll: Poll) {
+        self.poll.send(poll)
+    }
 }

Sources/StreamChat/Controllers/PollController/PollVoteListController+SwiftUI.swift

@@ -18,6 +18,9 @@ extension PollVoteListController {
         /// The poll votes.
         @Published public private(set) var votes: LazyCachedMapCollection<PollVote> = []
 
+        /// The poll which the votes belong to.
+        @Published public private(set) var poll: Poll?
+
         /// The current state of the controller.
         @Published public private(set) var state: DataController.State
 
@@ -29,6 +32,7 @@ extension PollVoteListController {
             controller.multicastDelegate.add(additionalDelegate: self)
 
             votes = controller.votes
+            poll = controller.poll
         }
     }
 }
@@ -41,6 +45,10 @@ extension PollVoteListController.ObservableObject: PollVoteListControllerDelegat
         votes = controller.votes
     }
 
+    public func controller(_ controller: PollVoteListController, didUpdatePoll poll: Poll) {
+        self.poll = poll
+    }
+
     public func controller(_ controller: DataController, didChangeState state: DataController.State) {
         self.state = state
     }

Sources/StreamChat/Controllers/PollController/PollVoteListController.swift

@@ -26,6 +26,24 @@ public protocol PollVoteListControllerDelegate: DataControllerStateDelegate {
         _ controller: PollVoteListController,
         didChangeVotes changes: [ListChange<PollVote>]
     )
+
+    /// The controller updated the poll.
+    ///
+    /// - Parameters:
+    ///   - controller: The controller emitting the change callback.
+    ///   - poll: The poll with the new data.
+    func controller(
+        _ controller: PollVoteListController,
+        didUpdatePoll poll: Poll
+    )
+}
+
+/// Optional delegate methods.
+public extension PollVoteListControllerDelegate {
+    func controller(
+        _ controller: PollVoteListController,
+        didUpdatePoll poll: Poll
+    ) {}
 }
 
 /// A controller which allows querying and filtering the votes of a poll.
@@ -37,14 +55,17 @@ public class PollVoteListController: DataController, DelegateCallable, DataStore
     public let client: ChatClient
 
     /// The votes of the poll the controller represents.
-    ///
-    /// To observe changes of the votes, set your class as a delegate of this controller or use the provided
-    /// `Combine` publishers.
     public var votes: LazyCachedMapCollection<PollVote> {
-        startPollVotesListObserverIfNeeded()
+        startObserversIfNeeded()
         return pollVotesObserver.items
     }
 
+    /// Returns the poll that this controller represents.
+    public var poll: Poll? {
+        startObserversIfNeeded()
+        return pollObserver?.item
+    }
+
     /// A Boolean value that returns whether pagination is finished.
     public private(set) var hasLoadedAllVotes: Bool = false
 
@@ -59,9 +80,7 @@ public class PollVoteListController: DataController, DelegateCallable, DataStore
         didSet {
             stateMulticastDelegate.set(mainDelegate: multicastDelegate.mainDelegate)
             stateMulticastDelegate.set(additionalDelegates: multicastDelegate.additionalDelegates)
-
-            // After setting delegate local changes will be fetched and observed.
-            startPollVotesListObserverIfNeeded()
+            startObserversIfNeeded()
         }
     }
 
@@ -78,7 +97,6 @@ public class PollVoteListController: DataController, DelegateCallable, DataStore
         observer.onDidChange = { [weak self] changes in
             self?.delegateCallback { [weak self] in
                 guard let self = self else {
-                    log.warning("Callback called while self is nil")
                     return
                 }
 
@@ -89,6 +107,31 @@ public class PollVoteListController: DataController, DelegateCallable, DataStore
         return observer
     }()
 
+    /// Used for observing the poll for changes.
+    private lazy var pollObserver: BackgroundEntityDatabaseObserver<Poll, PollDTO>? = { [weak self] in
+        guard let self = self else {
+            return nil
+        }
+
+        let observer = environment.pollObserverBuilder(
+            self.client.databaseContainer,
+            PollDTO.fetchRequest(for: query.pollId),
+            { try $0.asModel() as Poll },
+            NSFetchedResultsController<PollDTO>.self
+        )
+        .onChange { [weak self] change in
+            self?.delegateCallback { [weak self] delegate in
+                guard let self = self else {
+                    log.warning("Callback called while self is nil")
+                    return
+                }
+                delegate.controller(self, didUpdatePoll: change.item)
+            }
+        }
+
+        return observer
+    }()
+
     var _basePublishers: Any?
     /// An internal backing object for all publicly available Combine publishers. We use it to simplify the way we expose
     /// publishers. Instead of creating custom `Publisher` types, we use `CurrentValueSubject` and `PassthroughSubject` internally,
@@ -122,7 +165,7 @@ public class PollVoteListController: DataController, DelegateCallable, DataStore
     }
 
     override public func synchronize(_ completion: ((_ error: Error?) -> Void)? = nil) {
-        startPollVotesListObserverIfNeeded()
+        startObserversIfNeeded()
 
         pollsRepository.queryPollVotes(query: query) { [weak self] result in
             guard let self else { return }
@@ -140,12 +183,13 @@ public class PollVoteListController: DataController, DelegateCallable, DataStore
     }
 
     /// If the `state` of the controller is `initialized`, this method calls `startObserving` on the
-    /// `pollVotesObserver` to fetch the local data and start observing the changes. It also changes
+    /// `pollVotesObserver` and `pollObserver` to fetch the local data and start observing the changes. It also changes
     /// `state` based on the result.
-    private func startPollVotesListObserverIfNeeded() {
+    private func startObserversIfNeeded() {
         guard state == .initialized else { return }
         do {
             try pollVotesObserver.startObserving()
+            try pollObserver?.startObserving()
             state = .localDataFetched
         } catch {
             state = .localDataFetchFailed(ClientError(with: error))
@@ -198,6 +242,20 @@ extension PollVoteListController {
                     itemReuseKeyPaths: (\PollVote.id, \PollVoteDTO.id)
                 )
             }
+
+        var pollObserverBuilder: (
+            _ database: DatabaseContainer,
+            _ fetchRequest: NSFetchRequest<PollDTO>,
+            _ itemCreator: @escaping (PollDTO) throws -> Poll,
+            _ fetchedResultsControllerType: NSFetchedResultsController<PollDTO>.Type
+        ) -> BackgroundEntityDatabaseObserver<Poll, PollDTO> = {
+            BackgroundEntityDatabaseObserver(
+                database: $0,
+                fetchRequest: $1,
+                itemCreator: $2,
+                fetchedResultsControllerType: $3
+            )
+        }
     }
 }
 
@@ -210,7 +268,13 @@ extension PollVoteListController: EventsControllerDelegate {
             vote = event.vote
         }
         guard let vote else { return }
-        if vote.isAnswer == true && query.pollId == vote.pollId && query.optionId == nil {
+        if vote.isAnswer == true
+            && query.pollId == vote.pollId
+            && query.optionId == nil {
+            pollsRepository.link(pollVote: vote, to: query)
+        } else if vote.isAnswer == false
+            && query.pollId == vote.pollId
+            && query.optionId == vote.optionId {
             pollsRepository.link(pollVote: vote, to: query)
         }
     }

Sources/StreamChatUI/ChatMessageList/Attachments/Poll/PollResultsVC/PollResultsVoteListVC/PollResultsVoteListVC.swift

@@ -139,15 +139,6 @@ open class PollResultsVoteListVC:
         return view
     }
 
-    // MARK: - PollVoteListControllerDelegate
-
-    public func controller(_ controller: PollVoteListController, didChangeVotes changes: [ListChange<PollVote>]) {
-        var snapshot = NSDiffableDataSourceSnapshot<PollOption, PollVote>()
-        snapshot.appendSections([option])
-        snapshot.appendItems(Array(controller.votes))
-        dataSource.apply(snapshot, animatingDifferences: true)
-    }
-
     // MARK: - Actions
 
     /// Loads the next page of votes.
@@ -166,4 +157,20 @@ open class PollResultsVoteListVC:
     open func didFinishLoadingMoreVotes(with error: Error?) {
         isPaginatingVotes = false
     }
+
+    // MARK: - PollVoteListControllerDelegate
+
+    public func controller(_ controller: PollVoteListController, didChangeVotes changes: [ListChange<PollVote>]) {
+        var snapshot = NSDiffableDataSourceSnapshot<PollOption, PollVote>()
+        snapshot.appendSections([option])
+        snapshot.appendItems(Array(controller.votes))
+        dataSource.apply(snapshot, animatingDifferences: true)
+    }
+
+    public func controller(_ controller: PollVoteListController, didUpdatePoll poll: Poll) {
+        self.poll = poll
+        var newSnapshot = dataSource.snapshot()
+        newSnapshot.reloadSections([option])
+        dataSource.apply(newSnapshot, animatingDifferences: true)
+    }
 }

TestTools/StreamChatTestTools/Mocks/StreamChat/Controllers/PollVoteListController_Mock.swift

@@ -15,6 +15,11 @@ final class PollVoteListController_Mock: PollVoteListController {
     override var votes: LazyCachedMapCollection<PollVote> {
         votes_simulated
     }
+
+    var poll_simulated: Poll?
+    override var poll: Poll? {
+        poll_simulated
+    }
 
     var state_simulated: DataController.State?
     override var state: DataController.State {

TestTools/StreamChatTestTools/Mocks/StreamChat/Repositories/PollsRepository_Mock.swift

@@ -15,6 +15,9 @@ final class PollsRepository_Mock: PollsRepository, Spy {
 
     var recordedFunctions: [String] = []
     var spyState: SpyState = .init()
+
+    // Mock for link method
+    var link: ((PollVote, PollVoteListQuery) -> Void)?
 
     override func queryPollVotes(
         query: PollVoteListQuery,
@@ -65,4 +68,8 @@ final class PollsRepository_Mock: PollsRepository, Spy {
     ) {
         deletePoll_completion = completion
     }
+
+    override func link(pollVote: PollVote, to query: PollVoteListQuery) {
+        link?(pollVote, query)
+    }
 }

Tests/StreamChatTests/Controllers/PollsControllers/PollVoteListController+Combine_Tests.swift

@@ -79,4 +79,26 @@ final class PollVoteListController_Combine_Tests: iOS13TestCase {
 
         XCTAssertEqual(recording.output, .init(arrayLiteral: [.insert(vote, index: .init())]))
     }
+
+    func test_pollPublisher() {
+        // Setup Recording publishers
+        var recording = Record<Poll, Never>.Recording()
+
+        // Setup the chain
+        voteListController
+            .pollPublisher
+            .sink(receiveValue: { recording.receive($0) })
+            .store(in: &cancellables)
+
+        // Keep only the weak reference to the controller. The existing publisher should keep it alive.
+        weak var controller: PollVoteListController? = voteListController
+        voteListController = nil
+
+        let poll: Poll = .unique
+        controller?.delegateCallback {
+            $0.controller(controller!, didUpdatePoll: poll)
+        }
+
+        XCTAssertEqual(recording.output, [poll])
+    }
 }

Tests/StreamChatTests/Controllers/PollsControllers/PollVoteListController+SwiftUI_Tests.swift

@@ -66,4 +66,25 @@ final class PollVoteListController_SwiftUI_Tests: iOS13TestCase {
 
         AssertAsync.willBeEqual(observableObject.state, newState)
     }
+
+    func test_observableObject_initialPollValue() {
+        let observableObject = voteListController.observableObject
+
+        // Initially poll should be nil
+        XCTAssertNil(observableObject.poll)
+    }
+
+    func test_observableObject_reactsToDelegateUpdatePollCallback() {
+        let observableObject = voteListController.observableObject
+
+        // Simulate poll update
+        let poll: Poll = .unique
+        voteListController.poll_simulated = poll
+
+        voteListController.delegateCallback {
+            $0.controller(self.voteListController, didUpdatePoll: poll)
+        }
+
+        AssertAsync.willBeEqual(observableObject.poll, poll)
+    }
 }