From 1e6560a365d02fa077831efe4a08cff99e468f0b Mon Sep 17 00:00:00 2001
From: Brandon Rodriguez <brodriguez8774@gmail.com>
Date: Tue, 11 Jul 2023 21:30:08 -0400
Subject: [PATCH] Update add_permission/add_group functions to be more
 versatile based on how they're called

---
 .../mixins/core_mixin.py                      | 32 ++++++++++++++++---
 1 file changed, 28 insertions(+), 4 deletions(-)

diff --git a/django_expanded_test_cases/mixins/core_mixin.py b/django_expanded_test_cases/mixins/core_mixin.py
index f68718f..69ccb8c 100644
--- a/django_expanded_test_cases/mixins/core_mixin.py
+++ b/django_expanded_test_cases/mixins/core_mixin.py
@@ -581,7 +581,7 @@ class CoreTestCaseMixin:
 
         return user
 
-    def add_user_permission(self, user_permission, user=ETC_DEFAULT_STANDARD_USER_IDENTIFIER):
+    def add_user_permission(self, user_permission, user=None):
         """Adds Permission to given user.
 
         :param user_permission: Permission to add.
@@ -609,13 +609,25 @@ class CoreTestCaseMixin:
                     raise ValueError('Failed to find permission of "{0}".'.format(user_permission)) from pde
 
         # If we made it this far, then valid Permission was found. Apply to user.
-        user = self.get_user(user)
+        # But first determine what user we're applying to.
+        if user is not None:
+            # Actual user arg provided to function. Use that.
+            user = self.get_user(user)
+        elif self.user:
+            # No arg provided to function. Fall back to set class user, if provided.
+            user = self.get_user(self.user)
+        else:
+            # No arg provided to function AND no class user defined.
+            # Resort to "default standard user" as final fallback.
+            user = self.get_user(ETC_DEFAULT_STANDARD_USER_IDENTIFIER)
+
+        # Actually add permission.
         user.user_permissions.add(permission)
 
         # Return user object in case user wants to run additional checks.
         return user
 
-    def add_user_group(self, user_group, user=ETC_DEFAULT_STANDARD_USER_IDENTIFIER):
+    def add_user_group(self, user_group, user=None):
         """Adds Group to given user.
 
         :param user_group: Group to add.
@@ -639,7 +651,19 @@ class CoreTestCaseMixin:
                 raise ValueError('Failed to find Group of "{0}".'.format(user_group)) from gde
 
         # If we made it this far, then valid Group was found. Apply to user.
-        user = self.get_user(user)
+        # But first determine what user we're applying to.
+        if user is not None:
+            # Actual user arg provided to function. Use that.
+            user = self.get_user(user)
+        elif self.user:
+            # No arg provided to function. Fall back to set class user, if provided.
+            user = self.get_user(self.user)
+        else:
+            # No arg provided to function AND no class user defined.
+            # Resort to "default standard user" as final fallback.
+            user = self.get_user(ETC_DEFAULT_STANDARD_USER_IDENTIFIER)
+
+        # Actually add groups.
         user.groups.add(group)
 
         # Return user object in case user wants to run additional checks.
-- 
GitLab