Merge pull request #2917 from cisagov/nl/2859-multiple-portfolio-flags-bug

#2859 multiple portfolio flags bug - [KO]

src/registrar/models/user_portfolio_permission.py

@@ -108,16 +108,6 @@ class UserPortfolioPermission(TimeStampedModel):
         """Extends clean method to perform additional validation, which can raise errors in django admin."""
         super().clean()
 
-        # Check if a user is set without accessing the related object.
-        has_user = bool(self.user_id)
-        if self.pk is None and has_user:
-            existing_permissions = UserPortfolioPermission.objects.filter(user=self.user)
-            if not flag_is_active_for_user(self.user, "multiple_portfolios") and existing_permissions.exists():
-                raise ValidationError(
-                    "This user is already assigned to a portfolio. "
-                    "Based on current waffle flag settings, users cannot be assigned to multiple portfolios."
-                )
-
         # Check if portfolio is set without accessing the related object.
         has_portfolio = bool(self.portfolio_id)
         if not has_portfolio and self._get_portfolio_permissions():
@@ -125,3 +115,19 @@ class UserPortfolioPermission(TimeStampedModel):
 
         if has_portfolio and not self._get_portfolio_permissions():
             raise ValidationError("When portfolio is assigned, portfolio roles or additional permissions are required.")
+
+        # Check if a user is set without accessing the related object.
+        has_user = bool(self.user_id)
+        if has_user:
+            existing_permission_pks = UserPortfolioPermission.objects.filter(user=self.user).values_list(
+                "pk", flat=True
+            )
+            if (
+                not flag_is_active_for_user(self.user, "multiple_portfolios")
+                and existing_permission_pks.exists()
+                and self.pk not in existing_permission_pks
+            ):
+                raise ValidationError(
+                    "This user is already assigned to a portfolio. "
+                    "Based on current waffle flag settings, users cannot be assigned to multiple portfolios."
+                )

src/registrar/tests/test_models.py

@@ -299,6 +299,7 @@ class TestUserPortfolioPermission(TestCase):
     @less_console_noise_decorator
     def setUp(self):
         self.user, _ = User.objects.get_or_create(email="mayor@igorville.gov")
+        self.user2, _ = User.objects.get_or_create(email="user2@igorville.gov", username="user2")
         super().setUp()
 
     def tearDown(self):
@@ -336,16 +337,15 @@ class TestUserPortfolioPermission(TestCase):
     @override_flag("multiple_portfolios", active=False)
     def test_clean_on_creates_multiple_portfolios(self):
         """Ensures that a user cannot create multiple portfolio permission objects when the flag is disabled"""
-        # Create an instance of User with a portfolio but no roles or additional permissions
+        # Create an instance of User with a single portfolio
         portfolio, _ = Portfolio.objects.get_or_create(creator=self.user, organization_name="Hotel California")
-        portfolio_2, _ = Portfolio.objects.get_or_create(creator=self.user, organization_name="Motel California")
         portfolio_permission, _ = UserPortfolioPermission.objects.get_or_create(
             portfolio=portfolio, user=self.user, roles=[UserPortfolioRoleChoices.ORGANIZATION_ADMIN]
         )
+        portfolio_2, _ = Portfolio.objects.get_or_create(creator=self.user, organization_name="Motel California")
         portfolio_permission_2 = UserPortfolioPermission(
             portfolio=portfolio_2, user=self.user, roles=[UserPortfolioRoleChoices.ORGANIZATION_ADMIN]
         )
-
         # This should work as intended
         portfolio_permission.clean()
 
@@ -353,7 +353,37 @@ class TestUserPortfolioPermission(TestCase):
         with self.assertRaises(ValidationError) as cm:
             portfolio_permission_2.clean()
 
-        portfolio_permission_2, _ = UserPortfolioPermission.objects.get_or_create(portfolio=portfolio, user=self.user)
+        self.assertEqual(
+            cm.exception.message,
+            (
+                "This user is already assigned to a portfolio. "
+                "Based on current waffle flag settings, users cannot be assigned to multiple portfolios."
+            ),
+        )
+
+    @less_console_noise_decorator
+    @override_flag("multiple_portfolios", active=False)
+    def test_multiple_portfolio_reassignment(self):
+        """Ensures that a user cannot be assigned to multiple portfolios based on reassignment"""
+        # Create an instance of two users with separate portfolios
+        portfolio, _ = Portfolio.objects.get_or_create(creator=self.user, organization_name="Hotel California")
+        portfolio_permission, _ = UserPortfolioPermission.objects.get_or_create(
+            portfolio=portfolio, user=self.user, roles=[UserPortfolioRoleChoices.ORGANIZATION_ADMIN]
+        )
+        portfolio_2, _ = Portfolio.objects.get_or_create(creator=self.user2, organization_name="Motel California")
+        portfolio_permission_2 = UserPortfolioPermission(
+            portfolio=portfolio_2, user=self.user2, roles=[UserPortfolioRoleChoices.ORGANIZATION_ADMIN]
+        )
+
+        # This should work as intended
+        portfolio_permission.clean()
+        portfolio_permission_2.clean()
+
+        # Reassign the portfolio of "user2" to "user" (this should throw an error
+        # preventing "user" from having multiple portfolios)
+        with self.assertRaises(ValidationError) as cm:
+            portfolio_permission_2.user = self.user
+            portfolio_permission_2.clean()
 
         self.assertEqual(
             cm.exception.message,