Merge pull request #3089 from cisagov/el/3012-fixtures-bug

#3012: Fixtures bug - [EL]

.github/workflows/load-fixtures.yaml

@@ -0,0 +1,50 @@
+# Manually load fixtures to an environment of choice.
+
+name: Load fixtures
+run-name: Manually load fixtures to sandbox of choice
+
+on:
+  workflow_dispatch:
+    inputs:
+      environment:
+        description: Which environment should we load data for?
+        type: 'choice'
+        options:
+          - ab
+          - backup
+          - el
+          - cb
+          - dk
+          - es
+          - gd
+          - ko
+          - ky
+          - nl
+          - rb
+          - rh
+          - rjm
+          - meoward
+          - bob
+          - hotgov
+          - litterbox
+          - ms
+          - ad
+          - ag
+
+jobs:
+  load-fixtures:
+    runs-on: ubuntu-latest
+    env:
+      CF_USERNAME: CF_${{ github.event.inputs.environment }}_USERNAME
+      CF_PASSWORD: CF_${{ github.event.inputs.environment }}_PASSWORD
+    steps:
+      - uses: GitHubSecurityLab/actions-permissions/monitor@v1
+      - name: Load fake data for ${{ github.event.inputs.environment }}
+        uses: cloud-gov/cg-cli-tools@main
+        with:
+          cf_username: ${{ secrets[env.CF_USERNAME] }}
+          cf_password: ${{ secrets[env.CF_PASSWORD] }}
+          cf_org: cisa-dotgov
+          cf_space: ${{ github.event.inputs.environment }}
+          cf_command: "run-task getgov-${{ github.event.inputs.environment }} --command 'python manage.py load' --name loaddata"
+

src/registrar/fixtures/fixtures_requests.py

@@ -6,8 +6,10 @@
 from django.db import transaction
 
 from registrar.fixtures.fixtures_portfolios import PortfolioFixture
+from registrar.fixtures.fixtures_suborganizations import SuborganizationFixture
 from registrar.fixtures.fixtures_users import UserFixture
 from registrar.models import User, DomainRequest, DraftDomain, Contact, Website, FederalAgency
+from registrar.models.domain import Domain
 from registrar.models.portfolio import Portfolio
 from registrar.models.suborganization import Suborganization
 
@@ -101,8 +103,13 @@ def fake_contact(cls):
         }
 
     @classmethod
-    def fake_dot_gov(cls):
-        return f"{fake.slug()}.gov"
+    def fake_dot_gov(cls, max_attempts=100):
+        """Generate a unique .gov domain name without using an infinite loop."""
+        for _ in range(max_attempts):
+            fake_name = f"{fake.slug()}.gov"
+            if not Domain.objects.filter(name=fake_name).exists():
+                return DraftDomain.objects.create(name=fake_name)
+        raise RuntimeError(f"Failed to generate a unique .gov domain after {max_attempts} attempts")
 
     @classmethod
     def fake_expiration_date(cls):
@@ -189,7 +196,9 @@ def _get_requested_domain(cls, request: DomainRequest, request_dict: dict):
         if not request.requested_domain:
             if "requested_domain" in request_dict and request_dict["requested_domain"] is not None:
                 return DraftDomain.objects.get_or_create(name=request_dict["requested_domain"])[0]
-            return DraftDomain.objects.create(name=cls.fake_dot_gov())
+
+            # Generate a unique fake domain
+            return cls.fake_dot_gov()
         return request.requested_domain
 
     @classmethod
@@ -213,7 +222,7 @@ def _get_sub_organization(cls, request: DomainRequest, request_dict: dict):
         if not request.sub_organization:
             if "sub_organization" in request_dict and request_dict["sub_organization"] is not None:
                 return Suborganization.objects.get_or_create(name=request_dict["sub_organization"])[0]
-            return cls._get_random_sub_organization()
+            return cls._get_random_sub_organization(request)
         return request.sub_organization
 
     @classmethod
@@ -228,10 +237,19 @@ def _get_random_portfolio(cls):
             return None
 
     @classmethod
-    def _get_random_sub_organization(cls):
+    def _get_random_sub_organization(cls, request):
         try:
-            suborg_options = [Suborganization.objects.first(), Suborganization.objects.last()]
-            return random.choice(suborg_options)  # nosec
+            # Filter Suborganizations by the request's portfolio
+            portfolio_suborganizations = Suborganization.objects.filter(portfolio=request.portfolio)
+
+            # Select a suborg that's defined in the fixtures
+            suborganization_names = [suborg["name"] for suborg in SuborganizationFixture.SUBORGS]
+
+            # Further filter by names in suborganization_names
+            suborganization_options = portfolio_suborganizations.filter(name__in=suborganization_names)
+
+            # Randomly choose one if any exist
+            return random.choice(suborganization_options) if suborganization_options.exists() else None  # nosec
         except Exception as e:
             logger.warning(f"Expected fixture sub_organization, did not find it: {e}")
             return None
@@ -273,6 +291,9 @@ def load(cls):
 
         # Lumped under .atomic to ensure we don't make redundant DB calls.
         # This bundles them all together, and then saves it in a single call.
+        # The atomic block will cause the code to stop executing if one instance in the
+        # nested iteration fails, which will cause an early exit and make it hard to debug.
+        # Comment out with transaction.atomic() when debugging.
         with transaction.atomic():
             try:
                 # Get the usernames of users created in the UserFixture

src/registrar/fixtures/fixtures_users.py

@@ -267,19 +267,68 @@ def load_users(cls, users, group_name, are_superusers=False):
         """Loads the users into the database and assigns them to the specified group."""
         logger.info(f"Going to load {len(users)} users for group {group_name}")
 
+        # Step 1: Fetch the group
         group = UserGroup.objects.get(name=group_name)
 
-        # Prepare sets of existing usernames and IDs in one query
+        # Step 2: Identify new and existing users
+        existing_usernames, existing_user_ids = cls._get_existing_users(users)
+        new_users = cls._prepare_new_users(users, existing_usernames, existing_user_ids, are_superusers)
+
+        # Step 3: Create new users
+        cls._create_new_users(new_users)
+
+        # Step 4: Update existing users
+        # Get all users to be updated (both new and existing)
+        created_or_existing_users = User.objects.filter(username__in=[user.get("username") for user in users])
+        users_to_update = cls._get_users_to_update(created_or_existing_users)
+        cls._update_existing_users(users_to_update)
+
+        # Step 5: Assign users to the group
+        cls._assign_users_to_group(group, created_or_existing_users)
+
+        logger.info(f"Users loaded for group {group_name}.")
+
+    def load_allowed_emails(cls, users, additional_emails):
+        """Populates a whitelist of allowed emails (as defined in this list)"""
+        logger.info(f"Going to load allowed emails for {len(users)} users")
+        if additional_emails:
+            logger.info(f"Going to load {len(additional_emails)} additional allowed emails")
+
+        existing_emails = set(AllowedEmail.objects.values_list("email", flat=True))
+        new_allowed_emails = []
+
+        for user_data in users:
+            user_email = user_data.get("email")
+            if user_email and user_email not in existing_emails:
+                new_allowed_emails.append(AllowedEmail(email=user_email))
+
+        # Load additional emails, only if they don't exist already
+        for email in additional_emails:
+            if email not in existing_emails:
+                new_allowed_emails.append(AllowedEmail(email=email))
+
+        if new_allowed_emails:
+            try:
+                AllowedEmail.objects.bulk_create(new_allowed_emails)
+                logger.info(f"Loaded {len(new_allowed_emails)} allowed emails")
+            except Exception as e:
+                logger.error(f"Unexpected error during allowed emails bulk creation: {e}")
+        else:
+            logger.info("No allowed emails to load")
+
+    @staticmethod
+    def _get_existing_users(users):
         user_identifiers = [(user.get("username"), user.get("id")) for user in users]
         existing_users = User.objects.filter(
             username__in=[user[0] for user in user_identifiers] + [user[1] for user in user_identifiers]
         ).values_list("username", "id")
-
         existing_usernames = set(user[0] for user in existing_users)
         existing_user_ids = set(user[1] for user in existing_users)
+        return existing_usernames, existing_user_ids
 
-        # Filter out users with existing IDs or usernames
-        new_users = [
+    @staticmethod
+    def _prepare_new_users(users, existing_usernames, existing_user_ids, are_superusers):
+        return [
             User(
                 id=user_data.get("id"),
                 first_name=user_data.get("first_name"),
@@ -296,7 +345,8 @@ def load_users(cls, users, group_name, are_superusers=False):
             if user_data.get("username") not in existing_usernames and user_data.get("id") not in existing_user_ids
         ]
 
-        # Perform bulk creation for new users
+    @staticmethod
+    def _create_new_users(new_users):
         if new_users:
             try:
                 User.objects.bulk_create(new_users)
@@ -306,46 +356,36 @@ def load_users(cls, users, group_name, are_superusers=False):
         else:
             logger.info("No new users to create.")
 
-        # Get all users to be updated (both new and existing)
-        created_or_existing_users = User.objects.filter(username__in=[user.get("username") for user in users])
+    @staticmethod
+    def _get_users_to_update(users):
+        users_to_update = []
+        for user in users:
+            updated = False
+            if not user.title:
+                user.title = "Peon"
+                updated = True
+            if not user.phone:
+                user.phone = "2022222222"
+                updated = True
+            if not user.is_staff:
+                user.is_staff = True
+                updated = True
+            if updated:
+                users_to_update.append(user)
+        return users_to_update
 
-        # Filter out users who are already in the group
-        users_not_in_group = created_or_existing_users.exclude(groups__id=group.id)
+    @staticmethod
+    def _update_existing_users(users_to_update):
+        if users_to_update:
+            User.objects.bulk_update(users_to_update, ["is_staff", "title", "phone"])
+            logger.info(f"Updated {len(users_to_update)} existing users.")
 
-        # Add only users who are not already in the group
+    @staticmethod
+    def _assign_users_to_group(group, users):
+        users_not_in_group = users.exclude(groups__id=group.id)
         if users_not_in_group.exists():
             group.user_set.add(*users_not_in_group)
 
-        logger.info(f"Users loaded for group {group_name}.")
-
-    def load_allowed_emails(cls, users, additional_emails):
-        """Populates a whitelist of allowed emails (as defined in this list)"""
-        logger.info(f"Going to load allowed emails for {len(users)} users")
-        if additional_emails:
-            logger.info(f"Going to load {len(additional_emails)} additional allowed emails")
-
-        existing_emails = set(AllowedEmail.objects.values_list("email", flat=True))
-        new_allowed_emails = []
-
-        for user_data in users:
-            user_email = user_data.get("email")
-            if user_email and user_email not in existing_emails:
-                new_allowed_emails.append(AllowedEmail(email=user_email))
-
-        # Load additional emails, only if they don't exist already
-        for email in additional_emails:
-            if email not in existing_emails:
-                new_allowed_emails.append(AllowedEmail(email=email))
-
-        if new_allowed_emails:
-            try:
-                AllowedEmail.objects.bulk_create(new_allowed_emails)
-                logger.info(f"Loaded {len(new_allowed_emails)} allowed emails")
-            except Exception as e:
-                logger.error(f"Unexpected error during allowed emails bulk creation: {e}")
-        else:
-            logger.info("No allowed emails to load")
-
     @classmethod
     def load(cls):
         with transaction.atomic():