From 7faddef6de60c632965ffb930390038352ffa3ac Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=89ric=20Falconnier?= <eric@zentral.com>
Date: Thu, 12 Sep 2024 15:07:01 +0200
Subject: [PATCH] Fix nested realm groups issue with Santa Voter

---
 tests/santa/test_ballot_box.py      | 13 ++++++++++++-
 tests/santa/utils.py                | 10 +++++++++-
 zentral/contrib/santa/ballot_box.py |  2 +-
 3 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/tests/santa/test_ballot_box.py b/tests/santa/test_ballot_box.py
index 622a8bea9c..47dc03cca4 100644
--- 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
index 663d9f61af..781a4cd5a8 100644
--- 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
index 65f901095c..5129e3f6df 100644
--- 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):