From 02bc0cba1a942c90cec953d499fcc31db67698c5 Mon Sep 17 00:00:00 2001
From: zandercymatics <141044360+zandercymatics@users.noreply.github.com>
Date: Fri, 30 Aug 2024 09:49:43 -0600
Subject: [PATCH] Add suborg info to domains and requests

---
 .../commands/create_federal_portfolio.py      | 34 ++++++++++++-------
 .../tests/test_management_scripts.py          |  2 ++
 2 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/src/registrar/management/commands/create_federal_portfolio.py b/src/registrar/management/commands/create_federal_portfolio.py
index d02762817..cad318b57 100644
--- a/src/registrar/management/commands/create_federal_portfolio.py
+++ b/src/registrar/management/commands/create_federal_portfolio.py
@@ -185,13 +185,18 @@ class Command(BaseCommand):
         if not domain_requests.exists():
             message = "Portfolios not added to domain requests: no valid records found"
             TerminalHelper.colorful_logger(logger.info, TerminalColors.YELLOW, message)
-        else:
-            for domain_request in domain_requests:
-                domain_request.portfolio = portfolio
+            return None
+        
+        # Get all suborg information and store it in a dict to avoid doing a db call
+        suborgs = Suborganization.objects.filter(portfolio=portfolio).in_bulk(field_name="name")
+        for domain_request in domain_requests:
+            domain_request.portfolio = portfolio
+            if domain_request.organization_name in suborgs:
+                domain_request.sub_organization = suborgs.get(domain_request.organization_name)
 
-            DomainRequest.objects.bulk_update(domain_requests, ["portfolio"])
-            message = f"Added portfolio '{portfolio}' to {len(domain_requests)} domain requests"
-            TerminalHelper.colorful_logger(logger.info, TerminalColors.OKGREEN, message)
+        DomainRequest.objects.bulk_update(domain_requests, ["portfolio", "sub_organization"])
+        message = f"Added portfolio '{portfolio}' to {len(domain_requests)} domain requests"
+        TerminalHelper.colorful_logger(logger.info, TerminalColors.OKGREEN, message)
 
     def handle_portfolio_domains(self, portfolio: Portfolio, federal_agency: FederalAgency):
         """
@@ -202,10 +207,15 @@ class Command(BaseCommand):
         if not domain_infos.exists():
             message = "Portfolios not added to domains: no valid records found"
             TerminalHelper.colorful_logger(logger.info, TerminalColors.YELLOW, message)
-        else:
-            for domain_info in domain_infos:
-                domain_info.portfolio = portfolio
+            return None
 
-            DomainInformation.objects.bulk_update(domain_infos, ["portfolio"])
-            message = f"Added portfolio '{portfolio}' to {len(domain_infos)} domains"
-            TerminalHelper.colorful_logger(logger.info, TerminalColors.OKGREEN, message)
+        # Get all suborg information and store it in a dict to avoid doing a db call
+        suborgs = Suborganization.objects.filter(portfolio=portfolio).in_bulk(field_name="name")
+        for domain_info in domain_infos:
+            domain_info.portfolio = portfolio
+            if domain_info.organization_name in suborgs:
+                domain_info.sub_organization = suborgs.get(domain_info.organization_name)
+
+        DomainInformation.objects.bulk_update(domain_infos, ["portfolio", "sub_organization"])
+        message = f"Added portfolio '{portfolio}' to {len(domain_infos)} domains"
+        TerminalHelper.colorful_logger(logger.info, TerminalColors.OKGREEN, message)
diff --git a/src/registrar/tests/test_management_scripts.py b/src/registrar/tests/test_management_scripts.py
index 29db8a4c5..cbdc2c034 100644
--- a/src/registrar/tests/test_management_scripts.py
+++ b/src/registrar/tests/test_management_scripts.py
@@ -1490,6 +1490,7 @@ class TestCreateFederalPortfolio(TestCase):
         self.domain_request.refresh_from_db()
         self.assertIsNotNone(self.domain_request.portfolio)
         self.assertEqual(self.domain_request.portfolio.federal_agency, self.federal_agency)
+        self.assertEqual(self.domain_request.sub_organization.name, "Testorg")
 
     def test_handle_portfolio_domains(self):
         """Check portfolio association with domain information."""
@@ -1498,6 +1499,7 @@ class TestCreateFederalPortfolio(TestCase):
         self.domain_info.refresh_from_db()
         self.assertIsNotNone(self.domain_info.portfolio)
         self.assertEqual(self.domain_info.portfolio.federal_agency, self.federal_agency)
+        self.assertEqual(self.domain_info.sub_organization.name, "Testorg")
 
     def test_handle_parse_both(self):
         """Ensure correct parsing of both requests and domains."""