Merge pull request #2917 from cisagov/nl/2859-multiple-portfolio-flag…

…s-bug

#2859 multiple portfolio flags bug - [KO]

src/registrar/models/user_portfolio_permission.py

@@ -108,20 +108,26 @@ def clean(self):
         """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():
             raise ValidationError("When portfolio roles or additional permissions are assigned, portfolio is required.")
 
         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="[email protected]")
+        self.user2, _ = User.objects.get_or_create(email="[email protected]", username="user2")
         super().setUp()
 
     def tearDown(self):
@@ -336,24 +337,53 @@ def test_clean_on_multiple_portfolios_when_flag_active(self):
     @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()
 
         # Test if the ValidationError is raised with the correct message
         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,