feat: Make top_corporate_members more generic

2 files changed, 41 insertions, 54 deletions

diff --git a/fundraising/templatetags/fundraising_extras.py b/fundraising/templatetags/fundraising_extras.py
index 4936b86c..bed1708a 100644
--- a/fundraising/templatetags/fundraising_extras.py
+++ b/fundraising/templatetags/fundraising_extras.py
@@ -114,14 +114,16 @@ def display_django_heroes():
 
 
 @register.inclusion_tag("fundraising/includes/top_corporate_members.html")
-def top_corporate_members():
-    members = CorporateMember.objects.by_membership_level()
+def top_corporate_members(*levels, header=None):
+    """
+    Usage:
+    {% top_corporate_members "diamond" "platinum" header="Diamond and Platinum Members" %}
+    """
 
-    return {"members": members["diamond"] + members["platinum"]}
+    all_members: dict = CorporateMember.objects.by_membership_level()
 
+    members = []
+    for level in levels:
+        members.extend(all_members.get(level, []))
 
-@register.inclusion_tag("fundraising/includes/top_and_gold_corporate_members.html")
-def top_and_gold_corporate_members():
-    members = CorporateMember.objects.by_membership_level()
-
-    return {"members": members["diamond"] + members["platinum"] + members["gold"]}
+    return {"header": header, "members": members}
diff --git a/fundraising/tests/test_templatetags.py b/fundraising/tests/test_templatetags.py
index b56972ea..10f3958b 100644
--- a/fundraising/tests/test_templatetags.py
+++ b/fundraising/tests/test_templatetags.py
@@ -17,7 +17,6 @@ from ..models import (
 from ..templatetags.fundraising_extras import (
     display_django_heroes,
     donation_form_with_heart,
-    top_and_gold_corporate_members,
     top_corporate_members,
 )
 
@@ -111,12 +110,12 @@ class TestTopCorporateMembers(TestCase):
 
         member_1.invoice_set.create(amount=5, expiration_date=cls.future_date)
         member_2.invoice_set.create(amount=5, expiration_date=cls.future_date)
-        member_3.invoice_set.create(amount=5, expiration_date=cls.future_date)
+        member_3.invoice_set.create(amount=5, expiration_date=cls.past_date)
         member_4.invoice_set.create(amount=5, expiration_date=cls.past_date)
         member_5.invoice_set.create(amount=5, expiration_date=cls.past_date)
 
     def test_with_no_platinum_or_diamond_members(self):
-        members = top_corporate_members()["members"]
+        members = top_corporate_members("diamond", "platinum")["members"]
 
         self.assertEqual(members, [])
 
@@ -129,7 +128,7 @@ class TestTopCorporateMembers(TestCase):
         member_2.invoice_set.create(amount=8, expiration_date=self.future_date)
         member_3.invoice_set.create(amount=2, expiration_date=self.future_date)
 
-        members = top_corporate_members()["members"]
+        members = top_corporate_members("diamond", "platinum")["members"]
 
         self.assertEqual(members, [member_2, member_1, member_3])
 
@@ -142,7 +141,7 @@ class TestTopCorporateMembers(TestCase):
         member_2.invoice_set.create(amount=8, expiration_date=self.future_date)
         member_3.invoice_set.create(amount=2, expiration_date=self.future_date)
 
-        members = top_corporate_members()["members"]
+        members = top_corporate_members("diamond", "platinum")["members"]
 
         self.assertEqual(members, [member_2, member_1, member_3])
 
@@ -163,60 +162,46 @@ class TestTopCorporateMembers(TestCase):
         member_5.invoice_set.create(amount=8, expiration_date=self.future_date)
         member_6.invoice_set.create(amount=2, expiration_date=self.future_date)
 
-        members = top_corporate_members()["members"]
+        members = top_corporate_members("diamond", "platinum")["members"]
 
         expected = [member_5, member_4, member_6, member_2, member_1, member_3]
 
         self.assertEqual(members, expected)
 
+    def test_with_diamond_platinum_and_gold_members(self):
+        member_1 = CorporateMember.objects.create(membership_level=4)
+        member_2 = CorporateMember.objects.create(membership_level=4)
+        member_3 = CorporateMember.objects.create(membership_level=4)
 
-class TestTopAndGoldCorporateMembers(TestCase):
-    past_date = date(2000, 1, 1)
-    future_date = date(3000, 1, 1)
-
-    def test_with_no_gold_members(self):
-        members = top_and_gold_corporate_members()["members"]
-        self.assertEqual(members, [])
+        member_4 = CorporateMember.objects.create(membership_level=5)
+        member_5 = CorporateMember.objects.create(membership_level=5)
+        member_6 = CorporateMember.objects.create(membership_level=5)
 
-    def test_with_gold_members(self):
-        member_1 = CorporateMember.objects.create(membership_level=3)
-        member_2 = CorporateMember.objects.create(membership_level=3)
-        member_3 = CorporateMember.objects.create(membership_level=3)
+        member_7 = CorporateMember.objects.create(membership_level=3)
+        member_8 = CorporateMember.objects.create(membership_level=3)
 
         member_1.invoice_set.create(amount=4, expiration_date=self.future_date)
         member_2.invoice_set.create(amount=8, expiration_date=self.future_date)
         member_3.invoice_set.create(amount=2, expiration_date=self.future_date)
 
-        members = top_and_gold_corporate_members()["members"]
-
-        self.assertEqual(members, [member_2, member_1, member_3])
-
-    def test_with_gold_members_and_other_members(self):
-        member_1 = CorporateMember.objects.create(membership_level=3)
-        member_2 = CorporateMember.objects.create(membership_level=3)
-        member_3 = CorporateMember.objects.create(membership_level=4)
-        member_4 = CorporateMember.objects.create(membership_level=5)
-        member_5 = CorporateMember.objects.create(membership_level=2)
-
-        member_1.invoice_set.create(amount=4, expiration_date=self.future_date)
-        member_2.invoice_set.create(amount=8, expiration_date=self.future_date)
-        member_3.invoice_set.create(amount=2, expiration_date=self.future_date)
-        member_4.invoice_set.create(amount=2, expiration_date=self.future_date)
-        member_5.invoice_set.create(amount=2, expiration_date=self.future_date)
+        member_4.invoice_set.create(amount=4, expiration_date=self.future_date)
+        member_5.invoice_set.create(amount=8, expiration_date=self.future_date)
+        member_6.invoice_set.create(amount=2, expiration_date=self.future_date)
 
-        members = top_and_gold_corporate_members()["members"]
-        self.assertEqual(members, [member_4, member_3, member_2, member_1])
-        self.assertNotIn(member_5, members)
+        member_7.invoice_set.create(amount=8, expiration_date=self.future_date)
+        member_8.invoice_set.create(amount=2, expiration_date=self.future_date)
 
-    def test_with_gold_members_and_expired_invoice(self):
-        member_1 = CorporateMember.objects.create(membership_level=3)
-        member_2 = CorporateMember.objects.create(membership_level=3)
-        member_3 = CorporateMember.objects.create(membership_level=4)
+        members = top_corporate_members("diamond", "platinum", "gold")["members"]
 
-        member_1.invoice_set.create(amount=4, expiration_date=self.future_date)
-        member_2.invoice_set.create(amount=8, expiration_date=self.future_date)
-        member_3.invoice_set.create(amount=2, expiration_date=self.past_date)
+        expected = [
+            member_5,
+            member_4,
+            member_6,
+            member_2,
+            member_1,
+            member_3,
+            member_7,
+            member_8,
+        ]
 
-        members = top_and_gold_corporate_members()["members"]
-        self.assertEqual(members, [member_2, member_1])
-        self.assertNotIn(member_3, members)
+        self.assertEqual(members, expected)