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 faker import Faker
 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 @@ class DomainRequestFixture:
         }
 
     @classmethod
-    def fake_dot_gov(cls):
+    def fake_dot_gov(cls, max_attempts=100):
-        return f"{fake.slug()}.gov"
+        """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 @@ class DomainRequestFixture:
         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 @@ class DomainRequestFixture:
         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 @@ class DomainRequestFixture:
             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()]
+            # Filter Suborganizations by the request's portfolio
-            return random.choice(suborg_options)  # nosec
+            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 @@ class DomainRequestFixture:
 
         # 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,54 +267,24 @@ class UserFixture:
         """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
-        user_identifiers = [(user.get("username"), user.get("id")) for user in users]
+        existing_usernames, existing_user_ids = cls._get_existing_users(users)
-        existing_users = User.objects.filter(
+        new_users = cls._prepare_new_users(users, existing_usernames, existing_user_ids, are_superusers)
-            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)
+        # Step 3: Create new users
-        existing_user_ids = set(user[1] for user in existing_users)
+        cls._create_new_users(new_users)
-
-        # Filter out users with existing IDs or usernames
-        new_users = [
-            User(
-                id=user_data.get("id"),
-                first_name=user_data.get("first_name"),
-                last_name=user_data.get("last_name"),
-                username=user_data.get("username"),
-                email=user_data.get("email", ""),
-                title=user_data.get("title", "Peon"),
-                phone=user_data.get("phone", "2022222222"),
-                is_active=user_data.get("is_active", True),
-                is_staff=True,
-                is_superuser=are_superusers,
-            )
-            for user_data in users
-            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
-        if new_users:
-            try:
-                User.objects.bulk_create(new_users)
-                logger.info(f"Created {len(new_users)} new users.")
-            except Exception as e:
-                logger.error(f"Unexpected error during user bulk creation: {e}")
-        else:
-            logger.info("No new users to create.")
 
+        # 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)
 
-        # Filter out users who are already in the group
+        # Step 5: Assign users to the group
-        users_not_in_group = created_or_existing_users.exclude(groups__id=group.id)
+        cls._assign_users_to_group(group, created_or_existing_users)
-
-        # Add only users who are not already in the group
-        if users_not_in_group.exists():
-            group.user_set.add(*users_not_in_group)
 
         logger.info(f"Users loaded for group {group_name}.")
 
@@ -346,6 +316,76 @@ class UserFixture:
         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
+
+    @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"),
+                last_name=user_data.get("last_name"),
+                username=user_data.get("username"),
+                email=user_data.get("email", ""),
+                title=user_data.get("title", "Peon"),
+                phone=user_data.get("phone", "2022222222"),
+                is_active=user_data.get("is_active", True),
+                is_staff=True,
+                is_superuser=are_superusers,
+            )
+            for user_data in users
+            if user_data.get("username") not in existing_usernames and user_data.get("id") not in existing_user_ids
+        ]
+
+    @staticmethod
+    def _create_new_users(new_users):
+        if new_users:
+            try:
+                User.objects.bulk_create(new_users)
+                logger.info(f"Created {len(new_users)} new users.")
+            except Exception as e:
+                logger.error(f"Unexpected error during user bulk creation: {e}")
+        else:
+            logger.info("No new users to create.")
+
+    @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
+
+    @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.")
+
+    @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)
+
     @classmethod
     def load(cls):
         with transaction.atomic():