Highlight my vote even if it was made on another device (#7202)

src/components/views/messages/MPollBody.tsx

@@ -29,7 +29,7 @@ import {
 import StyledRadioButton from '../elements/StyledRadioButton';
 import { MatrixEvent } from "matrix-js-sdk/src/models/event";
 import { Relations } from 'matrix-js-sdk/src/models/relations';
-import { MatrixClientPeg } from '../../../MatrixClientPeg';
+import MatrixClientContext from "../../../contexts/MatrixClientContext";
 import ErrorDialog from '../dialogs/ErrorDialog';
 
 // TODO: [andyb] Use extensible events library when ready
@@ -42,20 +42,16 @@ interface IState {
 
 @replaceableComponent("views.messages.MPollBody")
 export default class MPollBody extends React.Component<IBodyProps, IState> {
+    static contextType = MatrixClientContext;
+    public context!: React.ContextType<typeof MatrixClientContext>;
+
     constructor(props: IBodyProps) {
         super(props);
 
-        const pollRelations = this.fetchPollRelations();
+        this.state = {
-        let selected = null;
+            selected: null,
-
+            pollRelations: this.fetchPollRelations(),
-        const userVotes = collectUserVotes(allVotes(pollRelations), null);
+        };
-        const userId = MatrixClientPeg.get().getUserId();
-        const currentVote = userVotes.get(userId);
-        if (currentVote) {
-            selected = currentVote.answers[0];
-        }
-
-        this.state = { selected, pollRelations };
 
         this.addListeners(this.state.pollRelations);
         this.props.mxEvent.on("Event.relationsCreated", this.onPollRelationsCreated);
@@ -119,7 +115,8 @@ export default class MPollBody extends React.Component<IBodyProps, IState> {
                 "rel_type": "m.reference",
             },
         };
-        MatrixClientPeg.get().sendEvent(
+
+        this.context.sendEvent(
             this.props.mxEvent.getRoomId(),
             POLL_RESPONSE_EVENT_TYPE.name,
             responseContent,
@@ -158,12 +155,13 @@ export default class MPollBody extends React.Component<IBodyProps, IState> {
     }
 
     /**
-     * @returns answer-id -> number-of-votes
+     * @returns userId -> UserVote
      */
-    private collectVotes(): Map<string, number> {
+    private collectUserVotes(): Map<string, UserVote> {
-        return countVotes(
+        return collectUserVotes(
-            collectUserVotes(allVotes(this.state.pollRelations), this.state.selected),
+            allVotes(this.state.pollRelations),
-            this.props.mxEvent.getContent(),
+            this.context.getUserId(),
+            this.state.selected,
         );
     }
 
@@ -184,15 +182,18 @@ export default class MPollBody extends React.Component<IBodyProps, IState> {
         }
 
         const pollId = this.props.mxEvent.getId();
-        const votes = this.collectVotes();
+        const userVotes = this.collectUserVotes();
+        const votes = countVotes(userVotes, this.props.mxEvent.getContent());
         const totalVotes = this.totalVotes(votes);
+        const userId = this.context.getUserId();
+        const myVote = userVotes.get(userId)?.answers[0];
 
         return <div className="mx_MPollBody">
             <h2>{ pollInfo.question[TEXT_NODE_TYPE] }</h2>
             <div className="mx_MPollBody_allOptions">
                 {
                     pollInfo.answers.map((answer: IPollAnswer) => {
-                        const checked = this.state.selected === answer.id;
+                        const checked = myVote === answer.id;
                         const classNames = `mx_MPollBody_option${
                             checked ? " mx_MPollBody_option_checked": ""
                         }`;
@@ -207,7 +208,7 @@ export default class MPollBody extends React.Component<IBodyProps, IState> {
                             <StyledRadioButton
                                 name={`poll_answer_select-${pollId}`}
                                 value={answer.id}
-                                checked={this.state.selected === answer.id}
+                                checked={checked}
                                 onChange={this.onOptionSelected}
                             >
                                 <div className="mx_MPollBody_optionDescription">
@@ -275,6 +276,7 @@ export function allVotes(pollRelations: Relations): Array<UserVote> {
  */
 function collectUserVotes(
     userResponses: Array<UserVote>,
+    userId: string,
     selected?: string,
 ): Map<string, UserVote> {
     const userVotes: Map<string, UserVote> = new Map();
@@ -287,8 +289,6 @@ function collectUserVotes(
     }
 
     if (selected) {
-        const client = MatrixClientPeg.get();
-        const userId = client.getUserId();
         userVotes.set(userId, new UserVote(0, userId, [selected]));
     }
 

test/components/views/messages/MPollBody-test.tsx

@@ -27,6 +27,8 @@ import { IPollAnswer, IPollContent } from "../../../../src/polls/consts";
 import { UserVote, allVotes } from "../../../../src/components/views/messages/MPollBody";
 import { MatrixClientPeg } from "../../../../src/MatrixClientPeg";
 
+const CHECKED = "mx_MPollBody_option_checked";
+
 const _MPollBody = sdk.getComponent("views.messages.MPollBody");
 const MPollBody = TestUtils.wrapInMatrixClientContext(_MPollBody);
 
@@ -142,6 +144,25 @@ describe("MPollBody", () => {
         expect(votesCount(body, "wings")).toBe("1 vote");
 
         expect(body.find(".mx_MPollBody_totalVotes").text()).toBe("Based on 2 votes");
+
+        // And my vote is highlighted
+        expect(voteButton(body, "wings").hasClass(CHECKED)).toBe(true);
+        expect(voteButton(body, "italian").hasClass(CHECKED)).toBe(false);
+    });
+
+    it("highlights my vote even if I did it on another device", () => {
+        // Given I voted italian
+        const votes = [
+            responseEvent("@me:example.com", "italian"),
+            responseEvent("@nf:example.com", "wings"),
+        ];
+        const body = newMPollBody(votes);
+
+        // But I didn't click anything locally
+
+        // Then my vote is highlighted, and others are not
+        expect(voteButton(body, "italian").hasClass(CHECKED)).toBe(true);
+        expect(voteButton(body, "wings").hasClass(CHECKED)).toBe(false);
     });
 
     it("ignores extra answers", () => {
@@ -380,6 +401,12 @@ function clickRadio(wrapper: ReactWrapper, value: string) {
     wrapper.find(`StyledRadioButton[value="${value}"]`).simulate("click");
 }
 
+function voteButton(wrapper: ReactWrapper, value: string): ReactWrapper {
+    return wrapper.find(
+        `div.mx_MPollBody_option`,
+    ).findWhere(w => w.key() === value);
+}
+
 function votesCount(wrapper: ReactWrapper, value: string): string {
     return wrapper.find(
         `StyledRadioButton[value="${value}"] .mx_MPollBody_optionVoteCount`,