deal with empty ballots (and test correct behaviour)

fsr-de · Apr 2, 2024 · 54036d5 · 54036d5
1 parent 7258835
commit 54036d5
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/myhpi/polls/models.py b/myhpi/polls/models.py
@@ -234,7 +234,8 @@ def calculate_ranking(self):
             names[option.pk] = option.name
 
         for ballot in ballots:
-            current_votes[heapq.heappop(ballot)[1].option.pk].append(ballot)
+            if ballot:
+                current_votes[heapq.heappop(ballot)[1].option.pk].append(ballot)
 
         def find_loosers():
             threshold = min(map(lambda key: len(current_votes[key]), current_votes.keys()))

diff --git a/myhpi/tests/polls/test_ranked_choice_algorithm.py b/myhpi/tests/polls/test_ranked_choice_algorithm.py
@@ -111,6 +111,13 @@ def test_no_ballots(self):
             [(1, "Alice", 0), (1, "Bob", 0), (1, "Charlie", 0), (1, "Dora", 0)],
         )
 
+    def test_empty_ballots(self):
+        self.cast_ballots([[]])
+        self.assertEqual(
+            self.poll.calculate_ranking(),
+            [(1, "Alice", 0), (1, "Bob", 0), (1, "Charlie", 0), (1, "Dora", 0)],
+        )
+
     def test_fast(self):
         self.cast_ballots(([["alice", "bob"]] * 1000) + ([["bob", "alice", "charlie"]] * 900))