From 922b1683fc2ba2ed20e809aabe2acf5381d963dd Mon Sep 17 00:00:00 2001
From: Steve Howell <showell@zulipchat.com>
Date: Wed, 22 Aug 2018 18:38:34 +0000
Subject: [PATCH] tests: Add test for /poll message.

This goes through the /messages endpoint
to get more full-stack testing, and it
focuses on how SubMessage rows get
created.
---
 zerver/tests/test_widgets.py | 61 ++++++++++++++++++++++++++++++++++++
 1 file changed, 61 insertions(+)

diff --git a/zerver/tests/test_widgets.py b/zerver/tests/test_widgets.py
index dc6d8c110a..24da658a27 100644
--- a/zerver/tests/test_widgets.py
+++ b/zerver/tests/test_widgets.py
@@ -2,6 +2,8 @@ import ujson
 
 from typing import Dict, Any
 
+from zerver.models import SubMessage
+
 from zerver.lib.test_classes import ZulipTestCase
 
 from zerver.lib.validator import check_widget_content
@@ -78,3 +80,62 @@ class WidgetContentTestCase(ZulipTestCase):
         payload['widget_content'] = ujson.dumps(bogus_data)
         result = self.api_post(sender_email, "/api/v1/messages", payload)
         self.assert_json_error_contains(result, 'Widgets: widget_type is not in widget_content')
+
+    def test_poll_command_extra_data(self) -> None:
+        sender_email = self.example_email('cordelia')
+        stream_name = 'Verona'
+        content = '/poll What is your favorite color?'
+
+        payload = dict(
+            type="stream",
+            to=stream_name,
+            sender=sender_email,
+            client='test suite',
+            subject='whatever',
+            content=content,
+        )
+        result = self.api_post(sender_email, "/api/v1/messages", payload)
+        self.assert_json_success(result)
+
+        message = self.get_last_message()
+        self.assertEqual(message.content, content)
+
+        expected_submessage_content = dict(
+            widget_type="poll",
+            extra_data=dict(
+                question="What is your favorite color?",
+            ),
+        )
+
+        submessage = SubMessage.objects.get(message_id=message.id)
+        self.assertEqual(submessage.msg_type, 'widget')
+        self.assertEqual(ujson.loads(submessage.content), expected_submessage_content)
+
+        # Now don't supply a question.
+
+        content = '/poll'
+        payload['content'] = content
+        result = self.api_post(sender_email, "/api/v1/messages", payload)
+        self.assert_json_success(result)
+
+        expected_submessage_content = dict(
+            widget_type="poll",
+            extra_data=dict(
+                question="",
+            ),
+        )
+
+        message = self.get_last_message()
+        self.assertEqual(message.content, content)
+        submessage = SubMessage.objects.get(message_id=message.id)
+        self.assertEqual(submessage.msg_type, 'widget')
+        self.assertEqual(ujson.loads(submessage.content), expected_submessage_content)
+
+        # Now test the feature flag.
+        with self.settings(ALLOW_SUB_MESSAGES=False):
+            result = self.api_post(sender_email, "/api/v1/messages", payload)
+        self.assert_json_success(result)
+
+        message = self.get_last_message()
+        self.assertEqual(message.content, content)
+        self.assertFalse(SubMessage.objects.filter(message_id=message.id).exists())