From 8c4dfdfa40ded129a60951db365e736a5a0bf88a Mon Sep 17 00:00:00 2001
From: Guo Xiang Tan <tgx_world@hotmail.com>
Date: Mon, 16 Jan 2017 22:51:57 +0800
Subject: [PATCH] FIX: Votes lost when editing a poll option with votes.

---
 plugins/poll/lib/polls_updater.rb           | 19 +++++++++--
 plugins/poll/spec/lib/polls_updater_spec.rb | 37 +++++++++++++++++----
 2 files changed, 47 insertions(+), 9 deletions(-)

diff --git a/plugins/poll/lib/polls_updater.rb b/plugins/poll/lib/polls_updater.rb
index 42d67f219d7..696a62eed11 100644
--- a/plugins/poll/lib/polls_updater.rb
+++ b/plugins/poll/lib/polls_updater.rb
@@ -64,7 +64,10 @@ module DiscoursePoll
           end
 
           polls[poll_name]["voters"] = previous_polls[poll_name]["voters"]
-          polls[poll_name]["anonymous_voters"] = previous_polls[poll_name]["anonymous_voters"] if previous_polls[poll_name].has_key?("anonymous_voters")
+
+          if previous_polls[poll_name].has_key?("anonymous_voters")
+            polls[poll_name]["anonymous_voters"] = previous_polls[poll_name]["anonymous_voters"]
+          end
 
           previous_options = previous_polls[poll_name]["options"]
           public_poll = polls[poll_name]["public"] == "true"
@@ -72,7 +75,19 @@ module DiscoursePoll
           polls[poll_name]["options"].each_with_index do |option, index|
             previous_option = previous_options[index]
             option["votes"] = previous_option["votes"]
-            option["anonymous_votes"] = previous_option["anonymous_votes"] if previous_option.has_key?("anonymous_votes")
+
+            if previous_option["id"] != option["id"]
+              if votes_fields = post.custom_fields[DiscoursePoll::VOTES_CUSTOM_FIELD]
+                votes_fields.each do |key, value|
+                  index = value[poll_name].index(previous_option["id"])
+                  votes_fields[key][poll_name][index] = option["id"] if index
+                end
+              end
+            end
+
+            if previous_option.has_key?("anonymous_votes")
+              option["anonymous_votes"] = previous_option["anonymous_votes"]
+            end
 
             if public_poll && previous_option.has_key?("voter_ids")
               option["voter_ids"] = previous_option["voter_ids"]
diff --git a/plugins/poll/spec/lib/polls_updater_spec.rb b/plugins/poll/spec/lib/polls_updater_spec.rb
index 4210b12d5e7..3d679d071a4 100644
--- a/plugins/poll/spec/lib/polls_updater_spec.rb
+++ b/plugins/poll/spec/lib/polls_updater_spec.rb
@@ -1,6 +1,8 @@
 require 'rails_helper'
 
 describe DiscoursePoll::PollsUpdater do
+  let(:user) { Fabricate(:user) }
+
   let(:post_with_two_polls) do
     raw = <<-RAW.strip_heredoc
     [poll]
@@ -127,8 +129,6 @@ describe DiscoursePoll::PollsUpdater do
         DiscoursePoll::PollsValidator.new(Fabricate(:post, raw: raw)).validate_polls
       end
 
-      let(:user) { Fabricate(:user) }
-
       before do
         DiscoursePoll::Poll.vote(post.id, "poll", ["5c24fc1df56d764b550ceae1b9319125"], user)
         post.reload
@@ -222,9 +222,16 @@ describe DiscoursePoll::PollsUpdater do
       let(:another_post) { Fabricate(:post, created_at: Time.zone.now - poll_edit_window_mins.minutes) }
 
       before do
-        polls.each { |key, value| value["voters"] = 2 }
         described_class.update(another_post, polls)
+        another_post.reload
         SiteSetting.poll_edit_window_mins = poll_edit_window_mins
+
+        DiscoursePoll::Poll.vote(
+          another_post.id,
+          "poll",
+          [polls["poll"]["options"].first["id"]],
+          user
+        )
       end
 
       it "should not allow new polls to be added" do
@@ -254,6 +261,8 @@ describe DiscoursePoll::PollsUpdater do
       end
 
       context "staff" do
+        let(:another_user) { Fabricate(:user) }
+
         it "should not allow staff to add options if votes have been casted" do
           another_post.update_attributes!(last_editor_id: User.staff.first.id)
 
@@ -284,8 +293,15 @@ describe DiscoursePoll::PollsUpdater do
           expect(message.data[:polls]).to eq(polls_with_3_options)
         end
 
-        it "should allow staff to edit options if votes have been casted" do
-          another_post.update_attributes!(last_editor_id: User.staff.first.id)
+        it "should allow staff to edit options even if votes have been casted" do
+          another_post.update!(last_editor_id: User.staff.first.id)
+
+          DiscoursePoll::Poll.vote(
+            another_post.id,
+            "poll",
+            [polls["poll"]["options"].first["id"]],
+            another_user
+          )
 
           raw = <<-RAW.strip_heredoc
           [poll]
@@ -301,9 +317,16 @@ describe DiscoursePoll::PollsUpdater do
             described_class.update(another_post, different_polls)
           end.first
 
-          different_polls.each { |key, value| value["voters"] = 2 }
+          custom_fields = another_post.reload.custom_fields
+
+          expect(custom_fields[DiscoursePoll::POLLS_CUSTOM_FIELD])
+            .to eq(different_polls)
+
+          [user, another_user].each do |u|
+            expect(custom_fields[DiscoursePoll::VOTES_CUSTOM_FIELD][u.id.to_s]["poll"])
+              .to eq(["68b434ff88aeae7054e42cd05a4d9056"])
+          end
 
-          expect(another_post.reload.custom_fields[DiscoursePoll::POLLS_CUSTOM_FIELD]).to eq(different_polls)
           expect(message.data[:post_id]).to eq(another_post.id)
           expect(message.data[:polls]).to eq(different_polls)
         end