From aba92eb2a4b15b96ad5bc8efd3285f5f852eadbd Mon Sep 17 00:00:00 2001
From: Jay Zeng <astrozzc@gmail.com>
Date: Tue, 25 Jun 2024 11:57:18 -0400
Subject: [PATCH] Fix unit test

---
 rbac/migration_tool/migrate.py        | 12 +++++-------
 tests/migration_tool/tests_migrate.py |  3 ++-
 2 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/rbac/migration_tool/migrate.py b/rbac/migration_tool/migrate.py
index ffd4ad92..ade98737 100644
--- a/rbac/migration_tool/migrate.py
+++ b/rbac/migration_tool/migrate.py
@@ -109,13 +109,11 @@ def migrate_role(role: Role):
     # With the replicated role bindings algorithm, role bindings are scoped by group, so we need to add groups
     # TODO: replace the hard coded groups
     policies = role.policies.all()
-    groups = frozenset(
-        {
-            V1group(str(policy.group.uuid), frozenset(policy.group.principals.values_list("uuid", flat=True)))
-            for policy in policies
-        }
-    )
-    v1_role = dataclasses.replace(v1_role, groups=groups)
+    groups = set()
+    for policy in policies:
+        principals = [str(principal) for principal in policy.group.principals.values_list("uuid", flat=True)]
+        groups.add(V1group(str(policy.group.uuid), frozenset(principals)))
+    v1_role = dataclasses.replace(v1_role, groups=frozenset(groups))
 
     # This is where we wire in the implementation we're using into the Migrator
     v1_to_v2_mapping = shared_system_role_replicated_role_bindings_v1_to_v2_mapping
diff --git a/tests/migration_tool/tests_migrate.py b/tests/migration_tool/tests_migrate.py
index fd97cca7..8b4e9527 100644
--- a/tests/migration_tool/tests_migrate.py
+++ b/tests/migration_tool/tests_migrate.py
@@ -56,6 +56,7 @@ def setUp(self):
         self.groupA21 = Group.objects.create(name="groupA21", tenant=self.tenant)
         self.principal1 = Principal.objects.create(username="principal1", tenant=self.tenant)
         self.principal2 = Principal.objects.create(username="principal2", tenant=self.tenant)
+        self.groupA21.principals.add(self.principal1, self.principal2)
         self.policyA21 = Policy.objects.create(name="System PolicyA21", group=self.groupA21, tenant=self.tenant)
         self.policyA21.roles.add(self.roleA2)
         self.policyA21.save()
@@ -82,5 +83,5 @@ def test_migration_of_roles(self, logger_mock):
         migrate_roles(**kwargs)
         self.assertEqual(
             len(logger_mock.info.call_args_list),
-            16,
+            18,
         )