diff --git a/src/registrar/models/portfolio_invitation.py b/src/registrar/models/portfolio_invitation.py index a76f0f9f9..68ce69fe3 100644 --- a/src/registrar/models/portfolio_invitation.py +++ b/src/registrar/models/portfolio_invitation.py @@ -6,6 +6,7 @@ from django.contrib.auth import get_user_model from django.db import models from django_fsm import FSMField, transition +from waffle import flag_is_active from registrar.models.user_portfolio_permission import UserPortfolioPermission from .utility.portfolio_helper import UserPortfolioPermissionChoices, UserPortfolioRoleChoices # type: ignore @@ -95,3 +96,12 @@ class PortfolioInvitation(TimeStampedModel): if self.portfolio_additional_permissions and len(self.portfolio_additional_permissions) > 0: user_portfolio_permission.portfolio_additional_permissions = self.portfolio_additional_permissions user_portfolio_permission.save() + + # The multiple_portfolios flag ensures that many portfolios can exist per user. + # This is (or will be) exposed as a multiselect, which they can choose from. + # Without it, we should just add the user on invitation as they'd have no way to select it. + # For now, the caller should ensure that this only occurs when no portfolio is selected. + # NOTE: Model enforcement of this occurs elsewhere. + if not flag_is_active(None, "multiple_portfolios"): + user.last_selected_portfolio = user_portfolio_permission.portfolio + user.save() diff --git a/src/registrar/models/user.py b/src/registrar/models/user.py index 2f173548b..29e0bc90d 100644 --- a/src/registrar/models/user.py +++ b/src/registrar/models/user.py @@ -14,7 +14,6 @@ from .transition_domain import TransitionDomain from .verified_by_staff import VerifiedByStaff from .domain import Domain from .domain_request import DomainRequest -from django.contrib.postgres.fields import ArrayField from waffle.decorators import flag_is_active from phonenumber_field.modelfields import PhoneNumberField # type: ignore @@ -378,7 +377,8 @@ class User(AbstractUser): for invitation in PortfolioInvitation.objects.filter( email__iexact=self.email, status=PortfolioInvitation.PortfolioInvitationStatus.INVITED ): - if self.last_selected_portfolio is None: + only_single_portfolio = not flag_is_active(None, "multiple_portfolios") and self.last_selected_portfolio is None + if only_single_portfolio or flag_is_active(None, "multiple_portfolios"): try: invitation.retrieve() invitation.save()