Fix nested realm groups issue with Santa Voter

zentralopensource · Sep 13, 2024 · 7faddef · 7faddef
1 parent 5596feb
commit 7faddef
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 3 deletions.
diff --git a/tests/santa/test_ballot_box.py b/tests/santa/test_ballot_box.py
@@ -5,7 +5,7 @@
                                               ResetNotAllowedError, Voter, VotingError, VotingNotAllowedError)
 from zentral.contrib.santa.models import Rule, Target, TargetState
 from .utils import (add_file_to_test_class, force_ballot, force_configuration, force_enrolled_machine,
-                    force_realm_user, force_target, force_voting_group)
+                    force_realm_group, force_realm_user, force_target, force_voting_group)
 
 
 class SantaBallotBoxTestCase(TestCase):
@@ -22,6 +22,17 @@ def test_voter_realm_groups(self):
         voter = Voter(realm_user)
         self.assertEqual(voter.realm_groups, [])
 
+    def test_voter_realm_nested_groups(self):
+        _, realm_user = force_realm_user()
+        parent = force_realm_group(realm=realm_user.realm)
+        child = force_realm_group(realm=parent.realm, parent=parent)
+        realm_user.groups.add(child)
+        voter = Voter(realm_user)
+        self.assertEqual(
+            sorted(voter.realm_groups, key=lambda rg: rg.created_at),
+            [parent, child]
+        )
+
     def test_voter_enrolled_machines(self):
         _, realm_user = force_realm_user()
         now = datetime.utcnow()

diff --git a/tests/santa/utils.py b/tests/santa/utils.py
@@ -39,6 +39,14 @@ def force_realm_user(realm=None, username=None, email=None):
     return realm, realm_user
 
 
+def force_realm_group(realm=None, parent=None):
+    return RealmGroup.objects.create(
+        realm=realm or force_realm(),
+        display_name=get_random_string(12),
+        parent=parent,
+    )
+
+
 def force_voting_group(
     configuration,
     realm_user,
@@ -48,7 +56,7 @@ def force_voting_group(
     can_unflag_target=False,
     can_reset_target=False,
 ):
-    realm_group = RealmGroup.objects.create(realm=realm_user.realm, display_name=get_random_string(12))
+    realm_group = force_realm_group(realm=realm_user.realm)
     realm_user.groups.add(realm_group)
     if ballot_target_types is None:
         ballot_target_types = [Target.Type.METABUNDLE, Target.Type.SIGNING_ID]

diff --git a/zentral/contrib/santa/ballot_box.py b/zentral/contrib/santa/ballot_box.py
@@ -40,7 +40,7 @@ def __init__(self, realm_user, max_machine_age_days=90, all_configurations=False
 
     @cached_property
     def realm_groups(self):
-        return list(self.realm_user.groups.all())
+        return [g for g, _ in self.realm_user.groups_with_types()]
 
     @cached_property
     def enrolled_machines(self):