Pollster list and assignment to polls

1 files changed, 22 insertions, 0 deletions

diff --git a/webapp/polls/tests/poll_tests.py b/webapp/polls/tests/poll_tests.py
index 27f2e2e..2e10416 100644
--- a/webapp/polls/tests/poll_tests.py
+++ b/webapp/polls/tests/poll_tests.py
@@ -115,6 +115,28 @@ class PollTests(MongoTestCase):
 
         self.assertFalse(poll.is_open())
 
+    def test_pollster_assignment(self):
+
+        poll1 = Poll(data={'name': "poll #1"})
+        poll1_id = poll1.save()
+        poll2 = Poll(data={'name': "poll #2"})
+        poll2_id = poll2.save()
+        poll3 = Poll(data={'name': "poll #3"})
+        poll3_id = poll3.save()
+
+        pollster = Pollster.create(username="test", password="test")
+
+        ids = [str(poll1_id), str(poll2_id)]
+
+        Poll.pollster_assignment(str(pollster.id), ids)
+
+        poll1 = Poll.get(poll1_id)
+        poll2 = Poll.get(poll2_id)
+        poll3 = Poll.get(poll3_id)
+        self.assertEqual(pollster.id, poll1.get_pollster().id)
+        self.assertEqual(pollster.id, poll2.get_pollster().id)
+        self.assertIsNone(poll3.get_pollster())
+
 
 class PollFormTests(MongoTestCase):