Add test for default Race votes and add_result method that updates unique candidate instance

Author: zstumgoren

elex4/lib/models.py

@@ -4,8 +4,24 @@ class Race(object):
     def __init__(self, date, raw_office):
         self.date = date
         self.office, self.district = self.__clean_office(raw_office)
+        self.total_votes = 0
+        self.candidates = {}
+
+    def add_result(self, result):
+        self.total_votes += result['votes']
+        candidate = self.__get_or_create_candidate(result)
+        candidate.add_votes(result['county'], result['votes'])
 
     # Private methods
+    def __get_or_create_candidate(self, result):
+        key = (result['party'], result['candidate'])
+        try:
+            candidate = self.candidates[key]
+        except KeyError:
+            candidate = Candidate(result['candidate'], result['party'])
+            self.candidates[key] = candidate
+        return candidate
+
     def __clean_office(self, office):
         if 'Rep' in office:
             office_clean = 'U.S. House of Representatives'

elex4/tests/test_models.py

@@ -77,3 +77,13 @@ def test_total_votes_update(self):
         race = Race("2012-11-06", "President")
         race.add_result(self.smith_result)
         self.assertEquals(race.total_votes, 2000)
+
+    def test_add_result_to_candidate(self):
+        "Race.add_result update unique candidate instance"
+        race = Race("2012-11-06", "President")
+        # Add a vote twice. If it's the same candidate, vote total should be sum of results
+        race.add_result(self.smith_result)
+        race.add_result(self.smith_result)
+        cand_key = (self.smith_result['party'], self.smith_result['candidate'])
+        candidate = race.candidates[cand_key]
+        self.assertEquals(candidate.votes, 4000)