Refactor to move logic into save + helper

src/registrar/models/domain_information.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 from django.db import transaction
 
 from registrar.models.utility.domain_helper import DomainHelper
+from registrar.models.utility.generic_helper import CreateOrUpdateOrganizationTypeHelper
 from .domain_request import DomainRequest
 from .utility.time_stamped_model import TimeStampedModel
 
@@ -235,6 +236,34 @@ class DomainInformation(TimeStampedModel):
         except Exception:
             return ""
 
+    def save(self, *args, **kwargs):
+        """Save override for custom properties"""
+
+        # Define mappings between generic org and election org.
+        # These have to be defined here, as you'd get a cyclical import error
+        # otherwise.
+
+        # For any given organization type, return the "_election" variant.
+        # For example: STATE_OR_TERRITORY => STATE_OR_TERRITORY_ELECTION
+        generic_org_map = DomainRequest.OrgChoicesElectionOffice.get_org_generic_to_org_election()
+
+        # For any given "_election" variant, return the base org type.
+        # For example: STATE_OR_TERRITORY_ELECTION => STATE_OR_TERRITORY
+        election_org_map = DomainRequest.OrgChoicesElectionOffice.get_org_election_to_org_generic()
+
+        # Manages the "organization_type" variable and keeps in sync with
+        # "is_election_office" and "generic_organization_type"
+        org_type_helper = CreateOrUpdateOrganizationTypeHelper(
+            sender=self.__class__,
+            instance=self,
+            generic_org_to_org_map=generic_org_map,
+            election_org_to_generic_org_map=election_org_map,
+        )
+
+        # Actually updates the organization_type field
+        org_type_helper.create_or_update_organization_type()
+        super().save(*args, **kwargs)
+
     @classmethod
     def create_from_da(cls, domain_request: DomainRequest, domain=None):
         """Takes in a DomainRequest and converts it into DomainInformation"""

src/registrar/models/domain_request.py

@@ -9,6 +9,7 @@ from django.db import models
 from django_fsm import FSMField, transition  # type: ignore
 from django.utils import timezone
 from registrar.models.domain import Domain
+from registrar.models.utility.generic_helper import CreateOrUpdateOrganizationTypeHelper
 from registrar.utility.errors import FSMDomainRequestError, FSMErrorCodes
 
 from .utility.time_stamped_model import TimeStampedModel
@@ -665,6 +666,34 @@ class DomainRequest(TimeStampedModel):
         help_text="Notes about this request",
     )
 
+    def save(self, *args, **kwargs):
+        """Save override for custom properties"""
+
+        # Define mappings between generic org and election org.
+        # These have to be defined here, as you'd get a cyclical import error
+        # otherwise.
+
+        # For any given organization type, return the "_election" variant.
+        # For example: STATE_OR_TERRITORY => STATE_OR_TERRITORY_ELECTION
+        generic_org_map = self.OrgChoicesElectionOffice.get_org_generic_to_org_election()
+
+        # For any given "_election" variant, return the base org type.
+        # For example: STATE_OR_TERRITORY_ELECTION => STATE_OR_TERRITORY
+        election_org_map = self.OrgChoicesElectionOffice.get_org_election_to_org_generic()
+
+        # Manages the "organization_type" variable and keeps in sync with
+        # "is_election_office" and "generic_organization_type"
+        org_type_helper = CreateOrUpdateOrganizationTypeHelper(
+            sender=self.__class__,
+            instance=self,
+            generic_org_to_org_map=generic_org_map,
+            election_org_to_generic_org_map=election_org_map,
+        )
+
+        # Actually updates the organization_type field
+        org_type_helper.create_or_update_organization_type()
+        super().save(*args, **kwargs)
+
     def __str__(self):
         try:
             if self.requested_domain and self.requested_domain.name:

src/registrar/models/utility/generic_helper.py

@@ -35,3 +35,211 @@ class Timer:
         self.end = time.time()
         self.duration = self.end - self.start
         logger.info(f"Execution time: {self.duration} seconds")
+
+
+class CreateOrUpdateOrganizationTypeHelper:
+    """
+    A helper that manages the "organization_type" field in DomainRequest and DomainInformation
+    """
+    def __init__(self, sender, instance, generic_org_to_org_map, election_org_to_generic_org_map):
+        # The "model type"
+        self.sender = sender
+        self.instance = instance
+        self.generic_org_to_org_map = generic_org_to_org_map
+        self.election_org_to_generic_org_map = election_org_to_generic_org_map
+
+    def create_or_update_organization_type(self):
+        """The organization_type field on DomainRequest and DomainInformation is consituted from the
+        generic_org_type and is_election_board fields. To keep the organization_type
+        field up to date, we need to update it before save based off of those field
+        values.
+
+        If the instance is marked as an election board and the generic_org_type is not
+        one of the excluded types (FEDERAL, INTERSTATE, SCHOOL_DISTRICT), the
+        organization_type is set to a corresponding election variant. Otherwise, it directly
+        mirrors the generic_org_type value.
+        """
+
+        # A new record is added with organization_type not defined.
+        # This happens from the regular domain request flow.
+        is_new_instance = self.instance.id is None
+        if is_new_instance:
+            self._handle_new_instance()
+        else:
+            self._handle_existing_instance()
+        
+        return self.instance
+
+    def _handle_new_instance(self):
+        # == Check for invalid conditions before proceeding == #
+        should_proceed = self._validate_new_instance()
+        if not should_proceed:
+            return None
+        # == Program flow will halt here if there is no reason to update == #
+
+        # == Update the linked values == #
+        organization_type_needs_update = self.instance.organization_type is None
+        generic_org_type_needs_update = self.instance.generic_org_type is None
+
+        # If a field is none, it indicates (per prior checks) that the
+        # related field (generic org type <-> org type) has data and we should update according to that.
+        if organization_type_needs_update:
+            self._update_org_type_from_generic_org_and_election()
+        elif generic_org_type_needs_update:
+            self._update_generic_org_and_election_from_org_type()
+        
+        # Update the field
+        self._update_fields(organization_type_needs_update, generic_org_type_needs_update)
+
+    def _handle_existing_instance(self):
+        # == Init variables == #
+        # Instance is already in the database, fetch its current state
+        current_instance = self.sender.objects.get(id=self.instance.id)
+
+        # Check the new and old values
+        generic_org_type_changed = self.instance.generic_org_type != current_instance.generic_org_type
+        is_election_board_changed = self.instance.is_election_board != current_instance.is_election_board
+        organization_type_changed = self.instance.organization_type != current_instance.organization_type
+
+        # == Check for invalid conditions before proceeding == #
+        if organization_type_changed and (generic_org_type_changed or is_election_board_changed):
+            # Since organization type is linked with generic_org_type and election board,
+            # we have to update one or the other, not both.
+            # This will not happen in normal flow as it is not possible otherwise.
+            raise ValueError("Cannot update organization_type and generic_org_type simultaneously.")
+        elif not organization_type_changed and (not generic_org_type_changed and not is_election_board_changed):
+            # No values to update - do nothing
+            return None
+        # == Program flow will halt here if there is no reason to update == #
+
+        # == Update the linked values == #
+        # Find out which field needs updating
+        organization_type_needs_update = generic_org_type_changed or is_election_board_changed
+        generic_org_type_needs_update = organization_type_changed
+
+        # Update the field
+        self._update_fields(organization_type_needs_update, generic_org_type_needs_update)
+
+    def _update_fields(self, organization_type_needs_update, generic_org_type_needs_update):
+        """
+        Validates the conditions for updating organization and generic organization types.
+        
+        Raises:
+            ValueError: If both organization_type_needs_update and generic_org_type_needs_update are True,
+                        indicating an attempt to update both fields simultaneously, which is not allowed.
+        """
+        # We shouldn't update both of these at the same time.
+        # It is more useful to have these as seperate variables, but it does impose
+        # this restraint.
+        if organization_type_needs_update and generic_org_type_needs_update:
+            raise ValueError("Cannot update both org type and generic org type at the same time.")
+        
+        if organization_type_needs_update:
+            self._update_org_type_from_generic_org_and_election()
+        elif generic_org_type_needs_update:
+            self._update_generic_org_and_election_from_org_type()
+
+    def _update_org_type_from_generic_org_and_election(self):
+        """Given a field values for generic_org_type and is_election_board, update the
+        organization_type field."""
+
+        # We convert to a string because the enum types are different.
+        generic_org_type = str(self.instance.generic_org_type)
+        if generic_org_type not in self.generic_org_to_org_map:
+            # Election board should always be reset to None if the record
+            # can't have one. For example, federal.
+            if self.instance.is_election_board is not None:
+                # This maintains data consistency.
+                # There is no avenue for this to occur in the UI,
+                # as such - this can only occur if the object is initialized in this way.
+                # Or if there are pre-existing data.
+                logger.warning(
+                    "create_or_update_organization_type() -> is_election_board "
+                    f"cannot exist for {generic_org_type}. Setting to None."
+                )
+                self.instance.is_election_board = None
+            self.instance.organization_type = generic_org_type
+        else:
+            # This can only happen with manual data tinkering, which causes these to be out of sync.
+            if self.instance.is_election_board is None:
+                logger.warning("create_or_update_organization_type() -> is_election_board is out of sync. Updating value.")
+                self.instance.is_election_board = False
+
+            if self.instance.is_election_board:
+                self.instance.organization_type = self.generic_org_to_org_map[generic_org_type]  
+            else: 
+                self.instance.organization_type = generic_org_type
+
+    def _update_generic_org_and_election_from_org_type(self):
+        """Given the field value for organization_type, update the
+        generic_org_type and is_election_board field."""
+
+        # We convert to a string because the enum types are different
+        # between OrgChoicesElectionOffice and OrganizationChoices.
+        # But their names are the same (for the most part).
+        current_org_type = str(self.instance.organization_type)
+        election_org_map = self.election_org_to_generic_org_map
+        generic_org_map = self.generic_org_to_org_map
+
+        # This essentially means: "_election" in current_org_type.
+        if current_org_type in election_org_map:
+            new_org = election_org_map[current_org_type]
+            self.instance.generic_org_type = new_org
+            self.instance.is_election_board = True
+        else:
+            self.instance.generic_org_type = current_org_type
+
+            # This basically checks if the given org type
+            # can even have an election board in the first place.
+            # For instance, federal cannot so is_election_board = None
+            if current_org_type in generic_org_map:
+                self.instance.is_election_board = False
+            else:
+                # This maintains data consistency.
+                # There is no avenue for this to occur in the UI,
+                # as such - this can only occur if the object is initialized in this way.
+                # Or if there are pre-existing data.
+                logger.warning(
+                    "create_or_update_organization_type() -> is_election_board "
+                    f"cannot exist for {current_org_type}. Setting to None."
+                )
+                self.instance.is_election_board = None
+
+    def _validate_new_instance(self):
+        """
+        Validates whether a new instance of DomainRequest or DomainInformation can proceed with the update
+        based on the consistency between organization_type, generic_org_type, and is_election_board.
+
+        Returns a boolean determining if execution should proceed or not.
+        """
+
+        # We conditionally accept both of these values to exist simultaneously, as long as
+        # those values do not intefere with eachother.
+        # Because this condition can only be triggered through a dev (no user flow),
+        # we throw an error if an invalid state is found here.
+        if self.instance.organization_type and self.instance.generic_org_type:
+            generic_org_type = str(self.instance.generic_org_type)
+            organization_type = str(self.instance.organization_type)
+
+            # Strip "_election" if it exists
+            mapped_org_type = self.election_org_to_generic_org_map.get(organization_type)
+
+            # Do tests on the org update for election board changes.
+            is_election_type = "_election" in organization_type
+            can_have_election_board = organization_type in self.generic_org_to_org_map
+
+            election_board_mismatch = (is_election_type != self.instance.is_election_board) and can_have_election_board
+            org_type_mismatch = mapped_org_type is not None and (generic_org_type != mapped_org_type)
+            if election_board_mismatch or org_type_mismatch:
+                message = (
+                    "Cannot add organization_type and generic_org_type simultaneously "
+                    "when generic_org_type, is_election_board, and organization_type values do not match."
+                )
+                raise ValueError(message)
+
+            return True
+        elif not self.instance.organization_type and not self.instance.generic_org_type:
+            return False
+        else:
+            return True
+

src/registrar/signals.py

@@ -9,193 +9,6 @@ from .models import User, Contact, DomainRequest, DomainInformation
 logger = logging.getLogger(__name__)
 
 
-@receiver(pre_save, sender=DomainRequest)
-@receiver(pre_save, sender=DomainInformation)
-def create_or_update_organization_type(sender: DomainRequest | DomainInformation, instance, **kwargs):
-    """The organization_type field on DomainRequest and DomainInformation is consituted from the
-    generic_org_type and is_election_board fields. To keep the organization_type
-    field up to date, we need to update it before save based off of those field
-    values.
-
-    If the instance is marked as an election board and the generic_org_type is not
-    one of the excluded types (FEDERAL, INTERSTATE, SCHOOL_DISTRICT), the
-    organization_type is set to a corresponding election variant. Otherwise, it directly
-    mirrors the generic_org_type value.
-    """
-
-    # == Init variables == #
-    election_org_choices = DomainRequest.OrgChoicesElectionOffice
-
-    # For any given organization type, return the "_election" variant.
-    # For example: STATE_OR_TERRITORY => STATE_OR_TERRITORY_ELECTION
-    generic_org_to_org_map = election_org_choices.get_org_generic_to_org_election()
-
-    # For any given "_election" variant, return the base org type.
-    # For example: STATE_OR_TERRITORY_ELECTION => STATE_OR_TERRITORY
-    election_org_to_generic_org_map = election_org_choices.get_org_election_to_org_generic()
-
-    # A new record is added with organization_type not defined.
-    # This happens from the regular domain request flow.
-    is_new_instance = instance.id is None
-
-    if is_new_instance:
-
-        # == Check for invalid conditions before proceeding == #
-        should_proceed = _validate_new_instance(instance, election_org_to_generic_org_map, generic_org_to_org_map)
-        if not should_proceed:
-            return None
-        # == Program flow will halt here if there is no reason to update == #
-
-        # == Update the linked values == #
-        organization_type_needs_update = instance.organization_type is None
-        generic_org_type_needs_update = instance.generic_org_type is None
-
-        # If a field is none, it indicates (per prior checks) that the
-        # related field (generic org type <-> org type) has data and we should update according to that.
-        if organization_type_needs_update:
-            _update_org_type_from_generic_org_and_election(instance, generic_org_to_org_map)
-        elif generic_org_type_needs_update:
-            _update_generic_org_and_election_from_org_type(
-                instance, election_org_to_generic_org_map, generic_org_to_org_map
-            )
-    else:
-
-        # == Init variables == #
-        # Instance is already in the database, fetch its current state
-        current_instance = sender.objects.get(id=instance.id)
-
-        # Check the new and old values
-        generic_org_type_changed = instance.generic_org_type != current_instance.generic_org_type
-        is_election_board_changed = instance.is_election_board != current_instance.is_election_board
-        organization_type_changed = instance.organization_type != current_instance.organization_type
-
-        # == Check for invalid conditions before proceeding == #
-        if organization_type_changed and (generic_org_type_changed or is_election_board_changed):
-            # Since organization type is linked with generic_org_type and election board,
-            # we have to update one or the other, not both.
-            # This will not happen in normal flow as it is not possible otherwise.
-            raise ValueError("Cannot update organization_type and generic_org_type simultaneously.")
-        elif not organization_type_changed and (not generic_org_type_changed and not is_election_board_changed):
-            # No values to update - do nothing
-            return None
-        # == Program flow will halt here if there is no reason to update == #
-
-        # == Update the linked values == #
-        # Find out which field needs updating
-        organization_type_needs_update = generic_org_type_changed or is_election_board_changed
-        generic_org_type_needs_update = organization_type_changed
-
-        # Update the field
-        if organization_type_needs_update:
-            _update_org_type_from_generic_org_and_election(instance, generic_org_to_org_map)
-        elif generic_org_type_needs_update:
-            _update_generic_org_and_election_from_org_type(
-                instance, election_org_to_generic_org_map, generic_org_to_org_map
-            )
-
-
-def _update_org_type_from_generic_org_and_election(instance, org_map):
-    """Given a field values for generic_org_type and is_election_board, update the
-    organization_type field."""
-
-    # We convert to a string because the enum types are different.
-    generic_org_type = str(instance.generic_org_type)
-    if generic_org_type not in org_map:
-        # Election board should always be reset to None if the record
-        # can't have one. For example, federal.
-        if instance.is_election_board is not None:
-            # This maintains data consistency.
-            # There is no avenue for this to occur in the UI,
-            # as such - this can only occur if the object is initialized in this way.
-            # Or if there are pre-existing data.
-            logger.warning(
-                "create_or_update_organization_type() -> is_election_board "
-                f"cannot exist for {generic_org_type}. Setting to None."
-            )
-            instance.is_election_board = None
-        instance.organization_type = generic_org_type
-    else:
-        # This can only happen with manual data tinkering, which causes these to be out of sync.
-        if instance.is_election_board is None:
-            logger.warning("create_or_update_organization_type() -> is_election_board is out of sync. Updating value.")
-            instance.is_election_board = False
-
-        instance.organization_type = org_map[generic_org_type] if instance.is_election_board else generic_org_type
-
-
-def _update_generic_org_and_election_from_org_type(instance, election_org_map, generic_org_map):
-    """Given the field value for organization_type, update the
-    generic_org_type and is_election_board field."""
-
-    # We convert to a string because the enum types are different
-    # between OrgChoicesElectionOffice and OrganizationChoices.
-    # But their names are the same (for the most part).
-    current_org_type = str(instance.organization_type)
-
-    # This essentially means: "_election" in current_org_type.
-    if current_org_type in election_org_map:
-        new_org = election_org_map[current_org_type]
-        instance.generic_org_type = new_org
-        instance.is_election_board = True
-    else:
-        instance.generic_org_type = current_org_type
-
-        # This basically checks if the given org type
-        # can even have an election board in the first place.
-        # For instance, federal cannot so is_election_board = None
-        if current_org_type in generic_org_map:
-            instance.is_election_board = False
-        else:
-            # This maintains data consistency.
-            # There is no avenue for this to occur in the UI,
-            # as such - this can only occur if the object is initialized in this way.
-            # Or if there are pre-existing data.
-            logger.warning(
-                "create_or_update_organization_type() -> is_election_board "
-                f"cannot exist for {current_org_type}. Setting to None."
-            )
-            instance.is_election_board = None
-
-
-def _validate_new_instance(instance, election_org_to_generic_org_map, generic_org_to_org_map):
-    """
-    Validates whether a new instance of DomainRequest or DomainInformation can proceed with the update
-    based on the consistency between organization_type, generic_org_type, and is_election_board.
-
-    Returns a boolean determining if execution should proceed or not.
-    """
-
-    # We conditionally accept both of these values to exist simultaneously, as long as
-    # those values do not intefere with eachother.
-    # Because this condition can only be triggered through a dev (no user flow),
-    # we throw an error if an invalid state is found here.
-    if instance.organization_type and instance.generic_org_type:
-        generic_org_type = str(instance.generic_org_type)
-        organization_type = str(instance.organization_type)
-
-        # Strip "_election" if it exists
-        mapped_org_type = election_org_to_generic_org_map.get(organization_type)
-
-        # Do tests on the org update for election board changes.
-        is_election_type = "_election" in organization_type
-        can_have_election_board = organization_type in generic_org_to_org_map
-
-        election_board_mismatch = (is_election_type != instance.is_election_board) and can_have_election_board
-        org_type_mismatch = mapped_org_type is not None and (generic_org_type != mapped_org_type)
-        if election_board_mismatch or org_type_mismatch:
-            message = (
-                "Cannot add organization_type and generic_org_type simultaneously "
-                "when generic_org_type, is_election_board, and organization_type values do not match."
-            )
-            raise ValueError(message)
-
-        return True
-    elif not instance.organization_type and not instance.generic_org_type:
-        return False
-    else:
-        return True
-
-
 @receiver(post_save, sender=User)
 def handle_profile(sender, instance, **kwargs):
     """Method for when a User is saved.