Unit tests

src/registrar/management/commands/create_federal_portfolio.py

@@ -72,8 +72,7 @@ class Command(BaseCommand):
             portfolio_args["senior_official"] = federal_agency.so_federal_agency.first()
 
         portfolio, created = Portfolio.objects.get_or_create(
-            organization_name=portfolio_args.get("organization_name"),
-            defaults=portfolio_args
+            organization_name=portfolio_args.get("organization_name"), defaults=portfolio_args
         )
 
         if created:
@@ -88,6 +87,15 @@ class Command(BaseCommand):
                 prompt_title="Do you wish to modify this record?",
             )
             if proceed:
+
+                # Don't override the creator and notes fields
+                if portfolio.creator:
+                    portfolio_args.pop("creator")
+
+                if portfolio.notes:
+                    portfolio_args.pop("notes")
+
+                # Update everything else
                 for key, value in portfolio_args.items():
                     setattr(portfolio, key, value)
                 portfolio.save()
@@ -98,11 +106,15 @@ class Command(BaseCommand):
 
     def create_suborganizations(self, portfolio: Portfolio, federal_agency: FederalAgency):
         """Create Suborganizations tied to the given portfolio based on DomainInformation objects"""
-        valid_agencies = DomainInformation.objects.filter(federal_agency=federal_agency, organization_name__isnull=False)
+        valid_agencies = DomainInformation.objects.filter(
+            federal_agency=federal_agency, organization_name__isnull=False
+        )
         org_names = set(valid_agencies.values_list("organization_name", flat=True))
 
         if not org_names:
-            TerminalHelper.colorful_logger(logger.warning, TerminalColors.YELLOW, f"No suborganizations found for {federal_agency}")
+            TerminalHelper.colorful_logger(
+                logger.warning, TerminalColors.YELLOW, f"No suborganizations found for {federal_agency}"
+            )
             return
 
         # Check if we need to update any existing suborgs first. This step is optional.
@@ -116,13 +128,17 @@ class Command(BaseCommand):
             if name.lower() == portfolio.organization_name.lower():
                 # If the suborg name is a portfolio name that currently exists, thats not a suborg - thats the portfolio itself!
                 # In this case, we can use this as an opportunity to update address information.
-                self._update_portfolio_location_details(portfolio, valid_agencies.filter(organization_name=name).first())
+                self._update_portfolio_location_details(
+                    portfolio, valid_agencies.filter(organization_name=name).first()
+                )
             else:
                 new_suborgs.append(Suborganization(name=name, portfolio=portfolio))
 
         if new_suborgs:
             Suborganization.objects.bulk_create(new_suborgs)
-            TerminalHelper.colorful_logger(logger.info, TerminalColors.OKGREEN, f"Added {len(new_suborgs)} suborganizations")
+            TerminalHelper.colorful_logger(
+                logger.info, TerminalColors.OKGREEN, f"Added {len(new_suborgs)} suborganizations"
+            )
         else:
             TerminalHelper.colorful_logger(logger.warning, TerminalColors.YELLOW, "No suborganizations added")
 

src/registrar/tests/common.py

@@ -911,6 +911,8 @@ def completed_domain_request(  # noqa
     federal_type=None,
     action_needed_reason=None,
     portfolio=None,
+    organization_name=None,
+    city=None,
 ):
     """A completed domain request."""
     if not user:
@@ -954,7 +956,7 @@ def completed_domain_request(  # noqa
         federal_type="executive",
         purpose="Purpose of the site",
         is_policy_acknowledged=True,
-        organization_name="Testorg",
+        organization_name=organization_name if organization_name else "Testorg",
         address_line1="address 1",
         address_line2="address 2",
         state_territory="NY",
@@ -984,6 +986,9 @@ def completed_domain_request(  # noqa
     if portfolio:
         domain_request_kwargs["portfolio"] = portfolio
 
+    if city:
+        domain_request_kwargs["city"] = city
+
     domain_request, _ = DomainRequest.objects.get_or_create(**domain_request_kwargs)
 
     if has_other_contacts:

src/registrar/tests/test_management_scripts.py

@@ -1,4 +1,5 @@
 import copy
+import boto3_mocking  # type: ignore
 from datetime import date, datetime, time
 from django.core.management import call_command
 from django.test import TestCase, override_settings
@@ -8,6 +9,7 @@ from django.utils import timezone
 from django.utils.module_loading import import_string
 import logging
 import pyzipper
+from django.core.management.base import CommandError
 from registrar.management.commands.clean_tables import Command as CleanTablesCommand
 from registrar.management.commands.export_tables import Command as ExportTablesCommand
 from registrar.models import (
@@ -23,14 +25,17 @@ from registrar.models import (
     VerifiedByStaff,
     PublicContact,
     FederalAgency,
+    Portfolio,
+    Suborganization,
 )
 import tablib
 from unittest.mock import patch, call, MagicMock, mock_open
 from epplibwrapper import commands, common
 
-from .common import MockEppLib, less_console_noise, completed_domain_request
+from .common import MockEppLib, less_console_noise, completed_domain_request, MockSESClient
 from api.tests.common import less_console_noise_decorator
 
+
 logger = logging.getLogger(__name__)
 
 
@@ -1411,23 +1416,125 @@ class TestPopulateFederalAgencyInitialsAndFceb(TestCase):
 
 
 class TestCreateFederalPortfolio(TestCase):
-    def setUp(self):
-        self.csv_path = "registrar/tests/data/fake_federal_cio.csv"
 
-        # Create test FederalAgency objects
-        self.agency1, _ = FederalAgency.objects.get_or_create(agency="American Battle Monuments Commission")
-        self.agency2, _ = FederalAgency.objects.get_or_create(agency="Advisory Council on Historic Preservation")
-        self.agency3, _ = FederalAgency.objects.get_or_create(agency="AMTRAK")
-        self.agency4, _ = FederalAgency.objects.get_or_create(agency="John F. Kennedy Center for Performing Arts")
+    @less_console_noise_decorator
+    def setUp(self):
+        self.mock_client = MockSESClient()
+        self.user = User.objects.create(username="testuser")
+        self.federal_agency = FederalAgency.objects.create(agency="Test Federal Agency")
+        with boto3_mocking.clients.handler_for("sesv2", self.mock_client):
+            self.domain_request = completed_domain_request(
+                status=DomainRequest.DomainRequestStatus.IN_REVIEW,
+                generic_org_type=DomainRequest.OrganizationChoices.FEDERAL,
+                federal_agency=self.federal_agency,
+                user=self.user,
+                city="WrongCity",
+            )
+            self.domain_request.approve()
+            self.domain_info = DomainInformation.objects.filter(domain_request=self.domain_request).get()
+
+            self.domain_request_2 = completed_domain_request(
+                name="sock@igorville.org",
+                status=DomainRequest.DomainRequestStatus.IN_REVIEW,
+                generic_org_type=DomainRequest.OrganizationChoices.FEDERAL,
+                federal_agency=self.federal_agency,
+                user=self.user,
+                organization_name="Test Federal Agency",
+                city="Block",
+            )
+            self.domain_request_2.approve()
+            self.domain_info_2 = DomainInformation.objects.filter(domain_request=self.domain_request_2).get()
 
     def tearDown(self):
-        SeniorOfficial.objects.all().delete()
+        DomainInformation.objects.all().delete()
+        DomainRequest.objects.all().delete()
+        Suborganization.objects.all().delete()
+        Portfolio.objects.all().delete()
         FederalAgency.objects.all().delete()
-    
-    # == create_or_modify_portfolio tests == #
-    # == create_suborganizations tests == #
-    # == handle_portfolio_requests tests == #
-    # == handle_portfolio_domains tests == #
-    # test for parse_requests
-    # test for parse_domains
-    # test for both
+        User.objects.all().delete()
+
+    @less_console_noise_decorator
+    def run_create_federal_portfolio(self, agency_name, parse_requests=False, parse_domains=False):
+        with patch(
+            "registrar.management.commands.utility.terminal_helper.TerminalHelper.query_yes_no_exit",
+            return_value=True,
+        ):
+            call_command(
+                "create_federal_portfolio", agency_name, parse_requests=parse_requests, parse_domains=parse_domains
+            )
+
+    def test_create_or_modify_portfolio(self):
+        self.run_create_federal_portfolio("Test Federal Agency", parse_requests=True)
+
+        portfolio = Portfolio.objects.get(federal_agency=self.federal_agency)
+        self.assertEqual(portfolio.organization_name, self.federal_agency.agency)
+        self.assertEqual(portfolio.organization_type, DomainRequest.OrganizationChoices.FEDERAL)
+        self.assertEqual(portfolio.creator, User.get_default_user())
+        self.assertEqual(portfolio.notes, "Auto-generated record")
+
+        # Test the suborgs
+        suborganizations = Suborganization.objects.filter(portfolio__federal_agency=self.federal_agency)
+        self.assertEqual(suborganizations.count(), 1)
+        self.assertEqual(suborganizations.first().name, "Testorg")
+
+        # Test other address information
+        self.assertEqual(portfolio.address_line1, "address 1")
+        self.assertEqual(portfolio.city, "Block")
+        self.assertEqual(portfolio.state_territory, "NY")
+        self.assertEqual(portfolio.zipcode, "10002")
+
+    def test_handle_portfolio_requests(self):
+        self.run_create_federal_portfolio("Test Federal Agency", parse_requests=True)
+
+        self.domain_request.refresh_from_db()
+        self.assertIsNotNone(self.domain_request.portfolio)
+        self.assertEqual(self.domain_request.portfolio.federal_agency, self.federal_agency)
+
+    def test_handle_portfolio_domains(self):
+        self.run_create_federal_portfolio("Test Federal Agency", parse_domains=True)
+
+        self.domain_info.refresh_from_db()
+        self.assertIsNotNone(self.domain_info.portfolio)
+        self.assertEqual(self.domain_info.portfolio.federal_agency, self.federal_agency)
+
+    def test_handle_parse_both(self):
+        self.run_create_federal_portfolio("Test Federal Agency", parse_requests=True, parse_domains=True)
+
+        self.domain_request.refresh_from_db()
+        self.domain_info.refresh_from_db()
+        self.assertIsNotNone(self.domain_request.portfolio)
+        self.assertIsNotNone(self.domain_info.portfolio)
+        self.assertEqual(self.domain_request.portfolio, self.domain_info.portfolio)
+
+    def test_command_error_no_parse_options(self):
+        with self.assertRaisesRegex(
+            CommandError, "You must specify at least one of --parse_requests or --parse_domains."
+        ):
+            self.run_create_federal_portfolio("Test Federal Agency")
+
+    def test_command_error_agency_not_found(self):
+        expected_message = (
+            "Cannot find the federal agency 'Non-existent Agency' in our database. "
+            "The value you enter for `agency_name` must be prepopulated in the FederalAgency table before proceeding."
+        )
+        with self.assertRaisesRegex(ValueError, expected_message):
+            self.run_create_federal_portfolio("Non-existent Agency", parse_requests=True)
+
+    def test_update_existing_portfolio(self):
+        # Create an existing portfolio
+        existing_portfolio = Portfolio.objects.create(
+            federal_agency=self.federal_agency,
+            organization_name="Test Federal Agency",
+            organization_type=DomainRequest.OrganizationChoices.CITY,
+            creator=self.user,
+            notes="Old notes",
+        )
+
+        self.run_create_federal_portfolio("Test Federal Agency", parse_requests=True)
+
+        existing_portfolio.refresh_from_db()
+        self.assertEqual(existing_portfolio.organization_name, self.federal_agency.agency)
+        self.assertEqual(existing_portfolio.organization_type, DomainRequest.OrganizationChoices.FEDERAL)
+        # Notes and creator should be untouched
+        self.assertEqual(existing_portfolio.notes, "Old notes")
+        self.assertEqual(existing_portfolio.creator, self.user)