From fa5c5f9ee1cd8b272716a7a7e0f052fe3a79ae74 Mon Sep 17 00:00:00 2001
From: "neiljp (Neil Pilgrim)" <github@kepier.clara.net>
Date: Mon, 24 Jul 2017 11:43:54 -0700
Subject: [PATCH 1/2] Initial pollbot from previous PR.

---
 zulip_bots/zulip_bots/bots/poll/__init__.py  |   0
 zulip_bots/zulip_bots/bots/poll/poll.py      | 257 +++++++++++++++++++
 zulip_bots/zulip_bots/bots/poll/readme.md    | 154 +++++++++++
 zulip_bots/zulip_bots/bots/poll/test_poll.py | 142 ++++++++++
 4 files changed, 553 insertions(+)
 create mode 100644 zulip_bots/zulip_bots/bots/poll/__init__.py
 create mode 100644 zulip_bots/zulip_bots/bots/poll/poll.py
 create mode 100644 zulip_bots/zulip_bots/bots/poll/readme.md
 create mode 100644 zulip_bots/zulip_bots/bots/poll/test_poll.py

diff --git a/zulip_bots/zulip_bots/bots/poll/__init__.py b/zulip_bots/zulip_bots/bots/poll/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/zulip_bots/zulip_bots/bots/poll/poll.py b/zulip_bots/zulip_bots/bots/poll/poll.py
new file mode 100644
index 000000000..1d59bb673
--- /dev/null
+++ b/zulip_bots/zulip_bots/bots/poll/poll.py
@@ -0,0 +1,257 @@
+# See readme.md for instructions on running this code.
+
+from __future__ import absolute_import
+from __future__ import print_function
+from six.moves import range
+
+from collections import OrderedDict, namedtuple
+
+def input_from_message_content(message_content):
+    lines = message_content.split('\n')
+    main_line = lines[0].split(' ')
+    command = main_line[0]
+    options = main_line[1:]
+    title   = ""
+    if len(lines) > 1:
+        title = lines[1]
+    vote_options = []
+    if len(lines) > 2:
+        vote_options = lines[2:]
+    vote_options = [v for v in vote_options if len(v) > 0]
+    Input = namedtuple('Input', ['command', 'options', 'title', 'vote_options'])
+    return Input(command, options, title, vote_options)
+
+class PollHandler(object):
+    def usage(self):
+        return '''
+            This bot maintains up to one poll per user, per topic, in streams only.
+            It currently keeps a running count of the votes, as they are made, with one
+            mesage in the stream being updated to show the current status.
+            Message the bot privately, appending the stream and topic, or mention it
+            within a topic (for new, vote and end commands); if the stream or
+            topic contain spaces use a '+' where the space would be.
+            '''
+
+    def handle_message(self, message, bot_handler, state_handler):
+
+        help_msg = OrderedDict([
+            ('about', "gives a simple summary of this bot."),
+            ('help', "produces this help."),
+            ('commands', "a concise form of help, listing the supported commands."),
+            ('new', ("start a new poll: specify a title on the following line "
+                     "and at least two options on subsequent lines.")),
+            ('vote', ("vote in an ongoing poll: specify a poll id given in the poll message "
+                      "followed by the number of the option to vote for.")),
+            ('end', ("end your own ongoing poll.")),
+        ])
+        stream_topic_notgiven = "\nPlease specify a stream & topic if messaging the bot privately."
+        space_equivalent = "+"
+
+        sender = message["sender_email"]
+        sender_id = message["sender_id"]
+
+        def private_reply(text):
+            bot_handler.send_message(dict(type='private', to=sender, content=text))
+
+        # Break down the text supplied into potential input sets
+        inp = input_from_message_content(message['content'])
+
+        # Updates the poll text, based on the current data in the poll (may modify poll)
+        # Return True/False, depending if message was successfully posted
+        def update_poll_text(poll, stream_topic, force_end=False):
+            # Construct initial/updated poll text
+            msg = ("Poll by {} (id: {})\n{}\n"
+                   .format(poll['creator'], stream_topic[2], poll['title']))
+            for i in range(poll['n']):
+                msg += ("{}. [{}] {}\n"
+                        .format(i+1, len(poll['tallies'][i]), poll[i]))
+            if force_end:
+                msg += "**This poll has ended**\n"
+            # Set/update poll text
+            if poll['msg_id'] is None:
+                result = bot_handler.send_message(dict(type='stream',
+                                                       to=stream_topic[0],
+                                                       subject=stream_topic[1],
+                                                       content=msg))
+                if result['result'] == 'success':
+                    poll['msg_id'] = result['id']
+                    return (msg, True)
+                else:
+                    return (msg, False)
+            else:
+                result = bot_handler.update_message(dict(message_id = poll['msg_id'],
+                                                         content = msg))
+                if result['result'] == 'success':
+                    return (msg, True)
+                else:
+                    return (msg, False)  # FIXME Cannot always update a message; should poll end?
+                    # FIXME Should add handling of return value
+
+        # Quickly return for simple help response cases
+        if inp.command == "" or inp.command == "help":
+            txt = ("{}\n\nIt supports the following commands:\n"
+                   .format(" ".join(self.usage().split())))
+            for k, v in help_msg.items():
+                txt += "\n**{}** : {}".format(k, v)
+            private_reply(txt)
+            return
+        elif inp.command == "about":
+            private_reply(" ".join(self.usage().split()))
+            return
+        elif inp.command == "commands":
+            private_reply("Commands: " + ", ".join((k for k in help_msg)))
+            return
+
+        src_is_private = (message['type'] == 'private')
+
+        # Ensure we have some state
+        active_polls = state_handler.get_state()
+        if active_polls is None:
+            active_polls = {}
+
+        if inp.command == "new":
+            # Check where the poll will be -> obtain stream_topic
+            stream_topic = None
+            if src_is_private:
+                if len(inp.options) != 2:
+                    private_reply(stream_topic_notgiven)
+                    return
+                else:
+                    stream = inp.options[0].replace(space_equivalent, " ")
+                    topic = inp.options[1].replace(space_equivalent, " ")
+                    stream_topic = (stream, topic, sender_id)
+            else:
+                stream_topic = (message['display_recipient'], message['subject'], sender_id)
+            # Check if a poll is already active with this id
+            if stream_topic in active_polls:
+                private_reply("Already have a poll with this id; end it explicitly first")
+                return
+            # Check we have at least a poll title and 2(+) options
+            if inp.title == "" or len(inp.vote_options) < 2:
+                private_reply("To " + help_msg['new'])
+                return
+            # Create new poll data
+            new_poll = {
+                'title': inp.title,  # Poll title
+                'tallies': [],   # List of list of sender_id's who voted
+                'msg_id': None,  # Message id containing poll text
+                'n': len(inp.vote_options),  # How many voting options
+                'creator': message['sender_full_name']  # Name of poll creator
+            }
+            for i, v in enumerate(inp.vote_options):      # Set text & tallies for each vote_option
+                new_poll[i] = v
+                new_poll['tallies'].append([])
+            # Try to send initial poll message to stream/topic, and alert user of result
+            (update_msg, success) = update_poll_text(new_poll, stream_topic)
+            if success:
+                # Insert the new poll and update the state
+                active_polls[stream_topic] = new_poll
+                state_handler.set_state(active_polls)
+                # Generate private message to creator, to indicate successful creation
+                msg = ("Poll created in stream '#{}' with topic '{}':\n{}"
+                       .format(stream_topic[0], stream_topic[1], update_msg))
+            else:
+                msg = ("Could not create poll in stream '#{}' with topic '{}'"
+                       .format(stream_topic[0], stream_topic[1]))
+            private_reply(msg)
+        elif inp.command == "vote":
+            # Translate options into requested poll and vote
+            requested_poll_id = 0
+            requested_vote_option = ""
+            stream_topic = None
+            if src_is_private:
+                if len(inp.options) != 4:
+                    private_reply("To " + help_msg['vote'] + stream_topic_notgiven)
+                    return
+                else:
+                    requested_poll_id = inp.options[2]
+                    requested_vote_option = inp.options[3]
+                    stream = inp.options[0].replace(space_equivalent, " ")
+                    topic = inp.options[1].replace(space_equivalent, " ")
+                    stream_topic = (stream, topic, requested_poll_id)
+            else:
+                if len(inp.options) != 2:
+                    private_reply("To " + help_msg['vote'])
+                    return
+                else:
+                    requested_poll_id = inp.options[0]
+                    requested_vote_option = inp.options[1]
+                    stream_topic = (message['display_recipient'],
+                                    message['subject'], requested_poll_id)
+            # Validate requested_poll_id before setting stream_topic & poll
+            try:
+                requested_poll_id = int(requested_poll_id)
+            except ValueError:
+                private_reply("To " + help_msg['vote'])
+                return
+            stream_topic = (stream_topic[0], stream_topic[1], requested_poll_id)
+            if stream_topic not in active_polls:
+                private_reply("To " + help_msg['vote'])
+                return
+            poll = active_polls[stream_topic]
+            # Check if this user has voted for any option already
+            context_txt = (" in the poll on stream '#{}' (topic '{}') titled: '{}'"
+                           .format(stream_topic[0], stream_topic[1], poll['title']))
+            for i, tally in enumerate(poll['tallies']):
+                if sender_id in tally:
+                    msg = ("You have already voted{}\n(You voted for {}: {})"
+                           .format(context_txt, i+1, poll[i]))
+                    private_reply(msg)
+                    return
+            # Check that vote index is within expected bounds
+            vote_index = 0
+            try:
+                vote_index = int(requested_vote_option)
+            except ValueError:
+                private_reply("Please select one number to vote for{}".format(context_txt))
+                return
+            max_vote_index = poll['n']
+            if 0 < vote_index <= max_vote_index:  # Indexed from 1
+                # Use the vote
+                poll['tallies'][vote_index-1].append(sender_id)
+                (update_msg, success) = update_poll_text(poll, stream_topic)
+                if success:
+                    state_handler.set_state(active_polls)
+                    msg = ("You just voted{}\n(You voted for {}: {})"
+                           .format(context_txt, vote_index, poll[vote_index-1]))
+                    private_reply(msg)
+                else:
+                    private_reply("Could not update the poll with your vote.")
+                    # FIXME Should we end the poll automatically here?
+            else:
+                private_reply(("Please select a number to vote for, between 1-{},{}"
+                               .format(max_vote_index, context_txt)))
+            return
+        elif inp.command == "end":
+            # Translate options into stream & topic
+            stream_topic = None
+            if src_is_private:
+                if len(inp.options) != 2:
+                    private_reply(stream_topic_notgiven)
+                    return
+                else:
+                    stream = inp.options[0].replace(space_equivalent, " ")
+                    topic = inp.options[1].replace(space_equivalent, " ")
+                    stream_topic = (stream, topic, sender_id)
+            else:
+                stream_topic = (message['display_recipient'], message['subject'], sender_id)
+            # End the poll if present
+            if stream_topic in active_polls:
+                (update_msg, success) = update_poll_text(active_polls[stream_topic],
+                                                         stream_topic, force_end=True)
+                private_reply((("Ending your poll in '#{}' and topic '{}', "
+                                "final results were:\n\n{}")
+                               .format(stream, topic, update_msg)))
+                if not success:
+                    private_reply("NOTE: Your poll ended, but the poll message could not be updated.")
+                del active_polls[stream_topic]
+                state_handler.set_state(active_polls)
+            else:
+                private_reply(("You do not have a poll in '#{}' and topic '{}'"
+                               .format(stream_topic[0], stream_topic[1])))
+        else:
+            msg = "Unsupported command."
+            private_reply(msg)
+
+
+handler_class = PollHandler
diff --git a/zulip_bots/zulip_bots/bots/poll/readme.md b/zulip_bots/zulip_bots/bots/poll/readme.md
new file mode 100644
index 000000000..f9122c9ad
--- /dev/null
+++ b/zulip_bots/zulip_bots/bots/poll/readme.md
@@ -0,0 +1,154 @@
+# Poll bot
+
+The poll bot maintains up to one poll per user, per topic, in streams only.
+It currently keeps a running count of the votes, as they are made, with one
+message in the stream being updated to show the current status.
+
+## Usage
+
+The bot can be messaged:
+
+* privately; the stream and topic must be mentioned explicitly in some commands
+(as in `[<stream> <topic>]` below), and a '+' must be used in place of any
+spaces in the stream or topic names.
+
+* in a stream/topic; the poll location is then implied to be in that stream and
+topic. The bot must be mentioned in this case, eg. each of the commands must
+be preceded with `@botname`, where botname is the bot user running this code.  
+
+The bot has the following commands, the success or failure of each being
+reported privately to the user attempting them:
+
+* **commands**:
+This command simply provides a concise listing of the available bot commands,
+which should give a list of the following.
+
+* **about**:
+This command explains what the bot is for.
+
+* **help**:
+This command provides a more detailed help listing than just the commands.
+
+* **new**:
+This command adds a new poll in the specified stream/topic, for the user.
+The format of this command is:
+```
+new [<stream> <topic>]
+<poll title>
+<option 1>
+<option 2>
+  ...
+```
+This produces a message in the stream/topic with the new poll details.
+It will fail if the user already has a poll in the stream/topic, in which
+case a different stream/topic can be used, or the current poll ended with
+the **end** command. A poll must have a title/question, and at least 2 options.
+
+* **vote**:
+This command allows any user to vote on a currently running poll, started
+with **new**.
+The format of this command is:
+```
+vote [<stream> <topic>] <poll id> <option number>
+```
+where `<poll id>` is the id of an active poll, as given in the main poll message
+output after the **new** command has been made.
+This currently attempts to update (edit) the poll message to indicate the vote.
+
+* **end**:
+This command allows a user to end a currently running poll.
+The format of this command is:
+```
+end [<stream> <topic>]  
+```
+The bot will attempt to update (edit) the poll message to indicate the poll has
+closed.
+This will fail if there is no active poll running in the specified topic for
+the user.
+
+### Example usage
+
+#### In a particular stream & topic
+
+```
+@pollbot commands
+```
+
+```
+@pollbot about
+```
+
+```
+@pollbot help
+```
+
+```
+@pollbot new
+How do you like the new poll bot service?
+Beyond measure!
+It's ok.
+It's terrible.
+```
+
+```
+@pollbot vote 2406 1
+```
+
+```
+@pollbot end
+```
+
+#### From a private message with the pollbot
+
+```
+commands
+```
+
+```
+about
+```
+
+```
+help
+```
+
+```
+new bot+testing testing+here
+How do you like the new poll bot service?
+Beyond measure!
+It's ok.
+It's terrible.
+```
+
+```
+vote bot+testing testing+here 2406 1
+```
+
+```
+@pollbot end
+```
+
+
+## Outstanding Bugs, Limitations & Issues
+
+### Larger issues
+
+* Polls can be explicitly ended, but also end if the bot is terminated or (if
+implemented) if the message cannot be updated. This could occur if the server
+is configured to not allow message editing, or only for a finite time (eg. 1
+hour), which the bot currently has no way of querying. The current state for each
+poll is stored within the bot, but also within the poll messages, which could
+allow for restarting/continuing polls if they could be queried.
+* The help text could be improved?
+
+### Potential areas to develop
+
+* Support changing of votes? (optionally)
+* Show votes only at the end of a vote? (optionally)
+* Support different vote schemes (eg. voting preferences for each option)
+* Support external vote engines?
+
+### Design decisions
+
+* Polls are limited to one per user per topic per stream (thought to be sufficient).
+* Polls are only in streams (this is an intentional initial design).
diff --git a/zulip_bots/zulip_bots/bots/poll/test_poll.py b/zulip_bots/zulip_bots/bots/poll/test_poll.py
new file mode 100644
index 000000000..3d1650063
--- /dev/null
+++ b/zulip_bots/zulip_bots/bots/poll/test_poll.py
@@ -0,0 +1,142 @@
+#!/usr/bin/env python
+
+from __future__ import absolute_import
+from __future__ import print_function
+
+from zulip_bots.test_lib import BotTestCase
+
+from collections import OrderedDict
+
+class TestPollBot(BotTestCase):
+    bot_name = "poll"
+
+    def test_bot(self):
+
+        test_email  = "sender@realm.com"
+        test_stream = "test stream"
+        test_topic  = "test topic"
+
+        private_msg_template = dict(type='private', to=test_email)
+        poll_msg_template = dict(type='stream', to=test_stream, subject=test_topic)
+
+        # Simple commands
+        # These are accepted anywhere, but private messages are sent to the sender
+
+        about_multiline = '''
+        This bot maintains up to one poll per user, per topic, in streams only.
+        It currently keeps a running count of the votes, as they are made, with one
+        mesage in the stream being updated to show the current status.
+        Message the bot privately, appending the stream and topic, or mention it
+        within a topic (for new, vote and end commands); if the stream or topic contain spaces use a
+        '+' where the space would be.
+        '''
+
+        about_txt = " ".join(about_multiline.split())
+        commands_txt = "Commands: about, help, commands, new, vote, end"
+        extended_commands_multiline = '''\n\nIt supports the following commands:\n\n**about** : gives a simple summary of this bot.\n**help** : produces this help.\n**commands** : a concise form of help, listing the supported commands.\n**new** : start a new poll: specify a title on the following line and at least two options on subsequent lines.\n**vote** : vote in an ongoing poll: specify a poll id given in the poll message followed by the number of the option to vote for.\n**end** : end your own ongoing poll.'''
+        help_txt = about_txt + extended_commands_multiline
+        simple_map = OrderedDict([
+            ("", help_txt),
+            ("about", about_txt),
+            ("help", help_txt),
+            ("commands", commands_txt),
+        ])
+        expected_simple = OrderedDict([(k, dict(private_msg_template, content=v))
+                                       for k, v in simple_map.items()])
+        self.check_expected_responses(expected_simple, expected_method='send_message',
+                                      email=test_email, type='all')
+
+        # Complex commands
+
+        error_msg = {
+            'new': ("To start a new poll: specify a title on the following line "
+                    "and at least two options on subsequent lines."),
+            'vote': ("To vote in an ongoing poll: specify a poll id given in the "
+                     "poll message followed by the number of the option to vote for."),
+            'end': ("You do not have a poll in '#test stream' and topic 'test topic'"),
+        }
+        poll_lines = ["Please select from one of the following options:",
+                      "Something great", "Something not so great", "This is bad"]
+        private_msg = "\nPlease specify a stream & topic if messaging the bot privately."
+
+        ## First do tests for cases that don't generate a poll and should fail with just
+        ## a private error message
+
+        def test_fails(src_location):
+            s = t = ""  # These need to be specified if src_location is private
+            recipient = test_stream
+            if src_location == 'private':
+                (s, t) = (" "+test_stream.replace(" ", "+"), " "+test_topic.replace(" ", "+"))
+                recipient = test_email
+            st = (src_location == 'stream')
+            # Note that if stream & topic are "" then some of these entries are overwritten
+            expected_ = OrderedDict([
+                ('new', error_msg['new'] if st else private_msg),
+                ('new' + s, error_msg['new'] if st else private_msg),
+                ('new' + s + t + t, error_msg['new'] if st else private_msg),
+                ('new' + s + t + "\n", error_msg['new']),
+                ('new' + s + t + "\n" + poll_lines[0], error_msg['new']),
+                ('new' + s + t + "\n" + poll_lines[0] + "\n", error_msg['new']),
+                ('new' + s + t + "\n" + poll_lines[0] + "\n" + poll_lines[1], error_msg['new']),
+                ('new' + s + t + "\n" + poll_lines[0] + "\n" + poll_lines[1] + "\n", error_msg['new']),
+                ('vote', error_msg['vote'] if st else error_msg['vote']+private_msg),
+                ('vote' + s, error_msg['vote'] if st else error_msg['vote']+private_msg),
+                ('vote' + s + t + t, error_msg['vote'] if st else error_msg['vote']+private_msg),
+                ('vote' + s + t + " 5", error_msg['vote'] if st else error_msg['vote']+private_msg),
+                ('vote' + s + t + " 5" + " 1", error_msg['vote']),  # no poll, so fails
+                ('end', error_msg['end'] if st else private_msg),
+                ('end' + s, error_msg['end'] if st else private_msg),
+                ('end' + s + t, error_msg['end']),  # no poll, so fails
+            ])
+            # Convert table above into expectation dict with templated messages
+            expected = OrderedDict([(k, dict(private_msg_template, content=v))
+                                    for k, v in expected_.items()])
+            self.check_expected_responses(expected, expected_method='send_message',
+                                          email=test_email, type=src_location,
+                                          recipient=recipient, subject='test topic')
+
+        # Test messages from within a stream/topic
+        test_fails('stream')
+        # Test messages sent privately
+        test_fails('private')
+
+        ## Now test the stateful poll generation, voting and poll ending
+        ## This depends on functional
+        ## - dual-posting support (poll-bot may update the poll message and
+        ##   always sends a private message stating what has happened)
+        ## - message update support
+        ## FIXME Both of these are currently lacking in the test framework
+
+        poll_message = "Poll created in stream '#test stream' with topic 'test topic':\nPoll by Foo Bar (id: 0)\nPlease select from one of the following options:\n1. [0] Something great\n2. [0] Something not so great\n3. [0] This is bad\n"
+
+        def test_polls(src_location):
+            s = t = ""  # These need to be specified if src_location is private
+            recipient = test_stream
+            if src_location == 'private':
+                (s, t) = (" "+test_stream.replace(" ", "+"), " "+test_topic.replace(" ", "+"))
+                recipient = test_email
+            st = (src_location == 'stream')
+            expected_ = OrderedDict([
+                ('new' + s + t + "\n" + "\n".join(poll_lines), poll_message),
+            ])
+            # Convert table above into expectation dict with templated messages
+            expected = OrderedDict([(k, dict(private_msg_template, content=v))
+                                    for k, v in expected_.items()])
+            self.check_expected_responses(expected, expected_method='send_message',
+                                          email=test_email, type=src_location,
+                                          recipient=recipient, subject=test_topic)
+
+            expected_ = OrderedDict([
+                ('new' + s + t + "\n" + "\n".join(poll_lines), poll_message),
+            ])
+            # Convert table above into expectation dict with templated messages
+            expected = OrderedDict([(k, dict(poll_msg_template, content=v))
+                                    for k, v in expected_.items()])
+            self.check_expected_responses(expected, expected_method='send_message',
+                                          email=test_email, type=src_location,
+                                          recipient=recipient, subject=test_topic)
+
+#        if 1: print("Testing poll usage from a stream")
+#        test_polls('stream')
+#        if 1: print("Testing poll usage from a private message")
+#        test_polls('private')

From 2834048e291d5eaef8899b2e88c23937fa36fd3d Mon Sep 17 00:00:00 2001
From: "neiljp (Neil Pilgrim)" <github@kepier.clara.net>
Date: Tue, 25 Jul 2017 15:32:54 -0700
Subject: [PATCH 2/2] Pollbot: Refactor, comment, type-annotate, use
 namedtuples & exceptions.

---
 zulip_bots/zulip_bots/bots/poll/poll.py | 400 +++++++++++++-----------
 1 file changed, 212 insertions(+), 188 deletions(-)

diff --git a/zulip_bots/zulip_bots/bots/poll/poll.py b/zulip_bots/zulip_bots/bots/poll/poll.py
index 1d59bb673..05c7d77df 100644
--- a/zulip_bots/zulip_bots/bots/poll/poll.py
+++ b/zulip_bots/zulip_bots/bots/poll/poll.py
@@ -7,6 +7,8 @@
 from collections import OrderedDict, namedtuple
 
 def input_from_message_content(message_content):
+    # Translate message content into expected input form
+    # (validation is dependent upon the command)
     lines = message_content.split('\n')
     main_line = lines[0].split(' ')
     command = main_line[0]
@@ -21,6 +23,133 @@ def input_from_message_content(message_content):
     Input = namedtuple('Input', ['command', 'options', 'title', 'vote_options'])
     return Input(command, options, title, vote_options)
 
+help_msg = OrderedDict([
+    ('about', "gives a simple summary of this bot."),
+    ('help', "produces this help."),
+    ('commands', "a concise form of help, listing the supported commands."),
+    ('new', ("start a new poll: specify a title on the following line "
+             "and at least two options on subsequent lines.")),
+    ('vote', ("vote in an ongoing poll: specify a poll id given in the poll message "
+              "followed by the number of the option to vote for.")),
+    ('end', ("end your own ongoing poll.")),
+])
+
+PollTuple = namedtuple('PollTuple', ['stream', 'topic', 'id'])
+
+stream_topic_notgiven = "\nPlease specify a stream & topic if messaging the bot privately."
+space_equivalent = "+"
+
+class InvalidInput(Exception):
+    pass
+
+def poll_context(poll_tuple, poll):
+    return (" on stream '#{}' (topic '{}') titled: '{}'"
+            .format(poll_tuple.stream, poll_tuple.topic, poll['title']))
+
+def validate_new_input(inputs, message, active_polls):
+    # (Input, Dict[str, Any], Dict[str, Any]) -> PollTuple
+    # Validate inputs to new command, raising InvalidInput or returning poll_tuple.
+    sender = message["sender_email"]
+    sender_id = message["sender_id"]
+    # Check input.options and set poll_tuple
+    if message['type'] == 'private':
+        if len(inputs.options) != 2:
+            raise InvalidInput(stream_topic_notgiven)
+        else:
+            stream = inputs.options[0].replace(space_equivalent, " ")
+            topic = inputs.options[1].replace(space_equivalent, " ")
+            poll_tuple = PollTuple(stream, topic, sender_id)
+    else:
+        poll_tuple = PollTuple(message['display_recipient'], message['subject'], sender_id)
+    # Check if a poll is already active with this id
+    if poll_tuple in active_polls:
+        raise InvalidInput(("You already have a poll running{}; end it explicitly first."
+                            .format(poll_context(poll_tuple, active_polls[poll_tuple]))))
+    # Check we have at least a poll title and 2(+) vote_options
+    if inputs.title == "" or len(inputs.vote_options) < 2:
+        raise InvalidInput("To " + help_msg['new'])  # FIXME improve message?
+    return poll_tuple
+
+def validate_vote_input(inp, message, active_polls):
+    # (Input, Dict[str, Any], Dict[str, Any]) -> (PollTuple, int)
+    # Validate inputs to vote command, raising InvalidInput,
+    # or returning PollTuple and the index into the votes to increment.
+    sender_id = message["sender_id"]
+    # Use inp.options to make 1st guess of poll_tuple, poll_id and vote_index
+    if message['type'] == 'private':
+        if len(inp.options) != 4:
+            raise InvalidInput("To " + help_msg['vote'] + stream_topic_notgiven)
+        else:
+            poll_id = inp.options[2]
+            vote_index = inp.options[3]
+            stream = inp.options[0].replace(space_equivalent, " ")
+            topic = inp.options[1].replace(space_equivalent, " ")
+            poll_tuple = PollTuple(stream, topic, poll_id)
+    else:
+        if len(inp.options) != 2:
+            raise InvalidInput("To " + help_msg['vote'])
+        else:
+            poll_id = inp.options[0]
+            vote_index = inp.options[1]
+            poll_tuple = PollTuple(message['display_recipient'],
+                                   message['subject'], poll_id)
+    # Update poll_tuple to ensure poll_id is an int
+    try:
+        poll_tuple = PollTuple(poll_tuple.stream, poll_tuple.topic, int(poll_id))
+    except ValueError:
+        raise InvalidInput("To " + help_msg['vote'])  # FIXME Improve message - id is not an int!
+    # Confirm poll_tuple relates to an active poll
+    poll = active_polls.get(poll_tuple)
+    if poll is None:
+        raise InvalidInput("To " + help_msg['vote'])  # FIXME Improve message - no poll exists!
+    # Ensure user has not voted in this poll already
+    for i, tally in enumerate(poll['tallies']):
+        if sender_id in tally:
+            msg = ("You have already voted in the poll{}\n(You voted for: {}. {})"
+                   .format(poll_context(poll_tuple, poll), i+1, poll[i]))
+            raise InvalidInput(msg)
+    # Check that vote_index is an int and within expected bounds
+    try:
+        vote_index = int(vote_index)
+    except ValueError:
+        raise InvalidInput("Please select one **number** to vote for in the poll{}"
+                           .format(poll_context(poll_tuple, poll)))
+    if not (0 < vote_index <= poll['n']):  # Indexed from 1
+        raise InvalidInput("Please select a number to vote for, **between 1-{}**, in the poll{}"
+                       .format(poll['n'], poll_context(poll_tuple, poll)))
+    return (poll_tuple, vote_index)
+
+def validate_end_input(inp, message, active_polls):
+    # (Input, Dict[str, Any], Dict[str, Any]) -> PollTuple
+    # Validate inputs to end command, raising InvalidInput or returning poll_tuple.
+    sender_id = message["sender_id"]
+    # Check input.options and set poll_tuple
+    if message['type'] == 'private':
+        if len(inp.options) != 2:
+            raise InvalidInput(stream_topic_notgiven)
+        else:
+            stream = inp.options[0].replace(space_equivalent, " ")
+            topic = inp.options[1].replace(space_equivalent, " ")
+            poll_tuple = PollTuple(stream, topic, sender_id)
+    else:
+        poll_tuple = PollTuple(message['display_recipient'], message['subject'], sender_id)
+    # Check the poll that the user wants to end exists already
+    if poll_tuple not in active_polls:
+        raise InvalidInput("You do not have a poll in '#{}' and topic '{}'"
+                           .format(poll_tuple.stream, poll_tuple.topic))
+    return poll_tuple
+
+def poll_text_from_poll(poll):
+    # (Dict[str, Any]) -> Text
+    # Given a poll, generate the text which appears in the message.
+    msg = ("Poll by {} (id: {})\n{}\n"
+           .format(poll['creator'], poll['creator_id'], poll['title']))
+    for i in range(poll['n']):
+        msg += ("{}. [{}] {}\n"
+                .format(i+1, len(poll['tallies'][i]), poll[i]))
+    return msg
+
+
 class PollHandler(object):
     def usage(self):
         return '''
@@ -34,66 +163,37 @@ def usage(self):
 
     def handle_message(self, message, bot_handler, state_handler):
 
-        help_msg = OrderedDict([
-            ('about', "gives a simple summary of this bot."),
-            ('help', "produces this help."),
-            ('commands', "a concise form of help, listing the supported commands."),
-            ('new', ("start a new poll: specify a title on the following line "
-                     "and at least two options on subsequent lines.")),
-            ('vote', ("vote in an ongoing poll: specify a poll id given in the poll message "
-                      "followed by the number of the option to vote for.")),
-            ('end', ("end your own ongoing poll.")),
-        ])
-        stream_topic_notgiven = "\nPlease specify a stream & topic if messaging the bot privately."
-        space_equivalent = "+"
-
         sender = message["sender_email"]
-        sender_id = message["sender_id"]
 
         def private_reply(text):
-            bot_handler.send_message(dict(type='private', to=sender, content=text))
+            # (Text) -> None
+            bot_handler.send_message(dict(type='private', to=sender, content=str(text)))
 
-        # Break down the text supplied into potential input sets
-        inp = input_from_message_content(message['content'])
-
-        # Updates the poll text, based on the current data in the poll (may modify poll)
-        # Return True/False, depending if message was successfully posted
-        def update_poll_text(poll, stream_topic, force_end=False):
-            # Construct initial/updated poll text
-            msg = ("Poll by {} (id: {})\n{}\n"
-                   .format(poll['creator'], stream_topic[2], poll['title']))
-            for i in range(poll['n']):
-                msg += ("{}. [{}] {}\n"
-                        .format(i+1, len(poll['tallies'][i]), poll[i]))
-            if force_end:
-                msg += "**This poll has ended**\n"
-            # Set/update poll text
+        def update_poll_text(poll, poll_tuple, force_end=False):
+            # (Dict[str, Any], PollTuple, bool) -> (Text, bool)
+            poll_text = "{}{}".format(poll_text_from_poll(poll),
+                                      "**This poll has ended**\n" if force_end else "")
             if poll['msg_id'] is None:
                 result = bot_handler.send_message(dict(type='stream',
-                                                       to=stream_topic[0],
-                                                       subject=stream_topic[1],
-                                                       content=msg))
+                                                       to=poll_tuple.stream,
+                                                       subject=poll_tuple.topic,
+                                                       content=poll_text))
                 if result['result'] == 'success':
                     poll['msg_id'] = result['id']
-                    return (msg, True)
-                else:
-                    return (msg, False)
             else:
                 result = bot_handler.update_message(dict(message_id = poll['msg_id'],
-                                                         content = msg))
-                if result['result'] == 'success':
-                    return (msg, True)
-                else:
-                    return (msg, False)  # FIXME Cannot always update a message; should poll end?
-                    # FIXME Should add handling of return value
+                                                         content = poll_text))
+            return (poll_text, 'success' == result['result'])
+
+        # Break down the text supplied into potential input
+        inp = input_from_message_content(message['content'])
 
-        # Quickly return for simple help response cases
+        # Simple commands with no state
         if inp.command == "" or inp.command == "help":
-            txt = ("{}\n\nIt supports the following commands:\n"
-                   .format(" ".join(self.usage().split())))
-            for k, v in help_msg.items():
-                txt += "\n**{}** : {}".format(k, v)
-            private_reply(txt)
+            private_reply("{}\n\nIt supports the following commands:\n\n{}"
+                          .format(" ".join(self.usage().split()),
+                                  "\n".join("**{}** : {}".format(k,v)
+                                            for k, v in help_msg.items())))
             return
         elif inp.command == "about":
             private_reply(" ".join(self.usage().split()))
@@ -101,156 +201,80 @@ def update_poll_text(poll, stream_topic, force_end=False):
         elif inp.command == "commands":
             private_reply("Commands: " + ", ".join((k for k in help_msg)))
             return
+        elif inp.command not in ('new', 'vote', 'end'):
+            private_reply("Unsupported command '{}'.".format(inp.command))
+            return
 
-        src_is_private = (message['type'] == 'private')
-
-        # Ensure we have some state
-        active_polls = state_handler.get_state()
-        if active_polls is None:
-            active_polls = {}
+        # We now have commands using state, so ensure we have some
+        with state_handler.state({}) as active_polls:
 
-        if inp.command == "new":
-            # Check where the poll will be -> obtain stream_topic
-            stream_topic = None
-            if src_is_private:
-                if len(inp.options) != 2:
-                    private_reply(stream_topic_notgiven)
-                    return
-                else:
-                    stream = inp.options[0].replace(space_equivalent, " ")
-                    topic = inp.options[1].replace(space_equivalent, " ")
-                    stream_topic = (stream, topic, sender_id)
-            else:
-                stream_topic = (message['display_recipient'], message['subject'], sender_id)
-            # Check if a poll is already active with this id
-            if stream_topic in active_polls:
-                private_reply("Already have a poll with this id; end it explicitly first")
-                return
-            # Check we have at least a poll title and 2(+) options
-            if inp.title == "" or len(inp.vote_options) < 2:
-                private_reply("To " + help_msg['new'])
-                return
-            # Create new poll data
-            new_poll = {
-                'title': inp.title,  # Poll title
-                'tallies': [],   # List of list of sender_id's who voted
-                'msg_id': None,  # Message id containing poll text
-                'n': len(inp.vote_options),  # How many voting options
-                'creator': message['sender_full_name']  # Name of poll creator
-            }
-            for i, v in enumerate(inp.vote_options):      # Set text & tallies for each vote_option
-                new_poll[i] = v
-                new_poll['tallies'].append([])
-            # Try to send initial poll message to stream/topic, and alert user of result
-            (update_msg, success) = update_poll_text(new_poll, stream_topic)
-            if success:
-                # Insert the new poll and update the state
-                active_polls[stream_topic] = new_poll
-                state_handler.set_state(active_polls)
-                # Generate private message to creator, to indicate successful creation
-                msg = ("Poll created in stream '#{}' with topic '{}':\n{}"
-                       .format(stream_topic[0], stream_topic[1], update_msg))
-            else:
-                msg = ("Could not create poll in stream '#{}' with topic '{}'"
-                       .format(stream_topic[0], stream_topic[1]))
-            private_reply(msg)
-        elif inp.command == "vote":
-            # Translate options into requested poll and vote
-            requested_poll_id = 0
-            requested_vote_option = ""
-            stream_topic = None
-            if src_is_private:
-                if len(inp.options) != 4:
-                    private_reply("To " + help_msg['vote'] + stream_topic_notgiven)
-                    return
-                else:
-                    requested_poll_id = inp.options[2]
-                    requested_vote_option = inp.options[3]
-                    stream = inp.options[0].replace(space_equivalent, " ")
-                    topic = inp.options[1].replace(space_equivalent, " ")
-                    stream_topic = (stream, topic, requested_poll_id)
-            else:
-                if len(inp.options) != 2:
-                    private_reply("To " + help_msg['vote'])
+            if inp.command == "new":
+                try:
+                    poll_tuple = validate_new_input(inp, message, active_polls)
+                except InvalidInput as e:
+                    private_reply(e)
                     return
+
+                # Create new poll data
+                new_poll = {
+                    'title': inp.title,  # Poll title
+                    'tallies': [],   # List of list of sender_id's who voted
+                    'msg_id': None,  # Message id with poll text (to update)
+                    'n': len(inp.vote_options),  # How many voting options
+                    'creator': message['sender_full_name'],  # Name of poll creator
+                    'creator_id': poll_tuple.id,
+                }
+                for i, v in enumerate(inp.vote_options):  # Set text & tallies for each vote_option
+                    new_poll[i] = v
+                    new_poll['tallies'].append([])
+
+                # Try to send initial poll message to stream/topic, and alert user of result
+                (update_msg, success) = update_poll_text(new_poll, poll_tuple)
+                if success:
+                    active_polls[poll_tuple] = new_poll
+                    msg = ("Poll created in stream '#{}' with topic '{}':\n{}"
+                           .format(poll_tuple.stream, poll_tuple.topic, update_msg))
                 else:
-                    requested_poll_id = inp.options[0]
-                    requested_vote_option = inp.options[1]
-                    stream_topic = (message['display_recipient'],
-                                    message['subject'], requested_poll_id)
-            # Validate requested_poll_id before setting stream_topic & poll
-            try:
-                requested_poll_id = int(requested_poll_id)
-            except ValueError:
-                private_reply("To " + help_msg['vote'])
-                return
-            stream_topic = (stream_topic[0], stream_topic[1], requested_poll_id)
-            if stream_topic not in active_polls:
-                private_reply("To " + help_msg['vote'])
-                return
-            poll = active_polls[stream_topic]
-            # Check if this user has voted for any option already
-            context_txt = (" in the poll on stream '#{}' (topic '{}') titled: '{}'"
-                           .format(stream_topic[0], stream_topic[1], poll['title']))
-            for i, tally in enumerate(poll['tallies']):
-                if sender_id in tally:
-                    msg = ("You have already voted{}\n(You voted for {}: {})"
-                           .format(context_txt, i+1, poll[i]))
-                    private_reply(msg)
+                    msg = ("Could not create poll in stream '#{}' with topic '{}'"
+                           .format(poll_tuple.stream, poll_tuple.topic))
+
+            elif inp.command == "vote":
+                try:
+                    (poll_tuple, vote_index) = validate_vote_input(inp, message, active_polls)
+                except InvalidInput as e:
+                    private_reply(e)
                     return
-            # Check that vote index is within expected bounds
-            vote_index = 0
-            try:
-                vote_index = int(requested_vote_option)
-            except ValueError:
-                private_reply("Please select one number to vote for{}".format(context_txt))
-                return
-            max_vote_index = poll['n']
-            if 0 < vote_index <= max_vote_index:  # Indexed from 1
+
                 # Use the vote
-                poll['tallies'][vote_index-1].append(sender_id)
-                (update_msg, success) = update_poll_text(poll, stream_topic)
+                poll = active_polls[poll_tuple]
+                poll['tallies'][vote_index-1].append(message['sender_id'])
+
+                # Try to update poll message to stream/topic, and alert user of result
+                (update_msg, success) = update_poll_text(poll, poll_tuple)
                 if success:
-                    state_handler.set_state(active_polls)
-                    msg = ("You just voted{}\n(You voted for {}: {})"
-                           .format(context_txt, vote_index, poll[vote_index-1]))
-                    private_reply(msg)
+                    msg = ("You just voted in the poll{}\n(You voted for: {}. {})"
+                           .format(poll_context(poll_tuple, poll), vote_index, poll[vote_index-1]))
                 else:
-                    private_reply("Could not update the poll with your vote.")
-                    # FIXME Should we end the poll automatically here?
-            else:
-                private_reply(("Please select a number to vote for, between 1-{},{}"
-                               .format(max_vote_index, context_txt)))
-            return
-        elif inp.command == "end":
-            # Translate options into stream & topic
-            stream_topic = None
-            if src_is_private:
-                if len(inp.options) != 2:
-                    private_reply(stream_topic_notgiven)
+                    msg = "**Could not update the poll with your vote.**"
+                    # FIXME Should we end the poll automatically here? roll-back data?
+
+            elif inp.command == "end":
+                try:
+                    poll_tuple = validate_end_input(inp, message, active_polls)
+                except InvalidInput as e:
+                    private_reply(e)
                     return
-                else:
-                    stream = inp.options[0].replace(space_equivalent, " ")
-                    topic = inp.options[1].replace(space_equivalent, " ")
-                    stream_topic = (stream, topic, sender_id)
-            else:
-                stream_topic = (message['display_recipient'], message['subject'], sender_id)
-            # End the poll if present
-            if stream_topic in active_polls:
-                (update_msg, success) = update_poll_text(active_polls[stream_topic],
-                                                         stream_topic, force_end=True)
-                private_reply((("Ending your poll in '#{}' and topic '{}', "
-                                "final results were:\n\n{}")
-                               .format(stream, topic, update_msg)))
+
+                # Try to update poll message to be ended, and end the poll
+                (update_msg, success) = update_poll_text(active_polls[poll_tuple],
+                                                         poll_tuple, force_end=True)
+                msg = (("Ending your poll in '#{}' and topic '{}'; "
+                       "final results were:\n\n{}")
+                       .format(poll_tuple.stream, poll_tuple.topic, update_msg[:-1]))
                 if not success:
-                    private_reply("NOTE: Your poll ended, but the poll message could not be updated.")
-                del active_polls[stream_topic]
-                state_handler.set_state(active_polls)
-            else:
-                private_reply(("You do not have a poll in '#{}' and topic '{}'"
-                               .format(stream_topic[0], stream_topic[1])))
-        else:
-            msg = "Unsupported command."
+                    msg += ("NOTE: Your poll ended, but the poll message could not be updated.")
+                del active_polls[poll_tuple]
+
             private_reply(msg)