Add tests

src/registrar/tests/data/fake_election_domains.csv

@@ -0,0 +1 @@
+manualtransmission.gov

src/registrar/tests/test_management_scripts.py

@@ -1,5 +1,6 @@
 import copy
 from datetime import date, datetime, time
+from unittest import skip
 from django.utils import timezone
 
 from django.test import TestCase
@@ -7,6 +8,9 @@ from django.test import TestCase
 from registrar.models import (
     User,
     Domain,
+    DomainRequest,
+    Contact,
+    Website,
     DomainInvitation,
     TransitionDomain,
     DomainInformation,
@@ -18,7 +22,224 @@ from django.core.management import call_command
 from unittest.mock import patch, call
 from epplibwrapper import commands, common
 
-from .common import MockEppLib, less_console_noise
+from .common import MockEppLib, less_console_noise, completed_domain_request
+from api.tests.common import less_console_noise_decorator
+
+class TestPopulateOrganizationType(MockEppLib):
+    """Tests for the populate_organization_type script"""
+
+    def setUp(self):
+        """Creates a fake domain object"""
+        super().setUp()
+
+        # Get the domain requests
+        self.domain_request_1 = completed_domain_request(
+            name="lasers.gov",
+            generic_org_type=DomainRequest.OrganizationChoices.FEDERAL,
+            is_election_board=True,
+        )
+        self.domain_request_2 = completed_domain_request(
+            name="readysetgo.gov",
+            generic_org_type=DomainRequest.OrganizationChoices.CITY,
+        )
+        self.domain_request_3 = completed_domain_request(
+            name="manualtransmission.gov",
+            generic_org_type=DomainRequest.OrganizationChoices.TRIBAL,
+        )
+        self.domain_request_4 = completed_domain_request(
+            name="saladandfries.gov",
+            generic_org_type=DomainRequest.OrganizationChoices.TRIBAL,
+            is_election_board=True,
+        )
+
+        # Approve all three requests
+        self.domain_request_1.approve()
+        self.domain_request_2.approve()
+        self.domain_request_3.approve()
+        self.domain_request_4.approve()
+
+        # Get the domains
+        self.domain_1 = Domain.objects.get(name="lasers.gov")
+        self.domain_2 = Domain.objects.get(name="readysetgo.gov")
+        self.domain_3 = Domain.objects.get(name="manualtransmission.gov")
+        self.domain_4 = Domain.objects.get(name="saladandfries.gov")
+
+        # Get the domain infos
+        self.domain_info_1 = DomainInformation.objects.get(domain=self.domain_1)
+        self.domain_info_2 = DomainInformation.objects.get(domain=self.domain_2)
+        self.domain_info_3 = DomainInformation.objects.get(domain=self.domain_3)
+        self.domain_info_4 = DomainInformation.objects.get(domain=self.domain_4)
+
+    def tearDown(self):
+        """Deletes all DB objects related to migrations"""
+        super().tearDown()
+
+        # Delete domains and related information
+        Domain.objects.all().delete()
+        DomainInformation.objects.all().delete()
+        DomainRequest.objects.all().delete()
+        User.objects.all().delete()
+        Contact.objects.all().delete()
+        Website.objects.all().delete()
+
+    @less_console_noise_decorator
+    def run_populate_organization_type(self):
+        """
+        This method executes the populate_organization_type command.
+
+        The 'call_command' function from Django's management framework is then used to
+        execute the populate_organization_type command with the specified arguments.
+        """
+        with patch(
+            "registrar.management.commands.utility.terminal_helper.TerminalHelper.query_yes_no_exit",  # noqa
+            return_value=True,
+        ):
+            call_command("populate_organization_type", "registrar/tests/data/fake_election_domains.csv", debug=True)
+    
+    def assert_expected_org_values_on_request_and_info(
+        self,
+        domain_request: DomainRequest,
+        domain_info: DomainInformation,
+        expected_values: dict,
+    ):
+        """
+        This is a a helper function that ensures that:
+        1. DomainRequest and DomainInformation (on given objects) are equivalent
+        2. That generic_org_type, is_election_board, and organization_type are equal to passed in values
+        """
+
+        # Test domain request
+        self.assertEqual(domain_request.generic_org_type, expected_values['generic_org_type'])
+        self.assertEqual(domain_request.is_election_board, expected_values['is_election_board'])
+        self.assertEqual(domain_request.organization_type, expected_values['organization_type'])
+
+        # Test domain info
+        self.assertEqual(domain_info.generic_org_type, expected_values['generic_org_type'])
+        self.assertEqual(domain_info.is_election_board, expected_values['is_election_board'])
+        self.assertEqual(domain_info.organization_type, expected_values['organization_type'])
+
+    def test_request_and_info_city_not_in_csv(self):
+        """Tests what happens to a city domain that is not defined in the CSV"""
+        city_request = self.domain_request_2
+        city_info = self.domain_request_2
+
+        # Make sure that all data is correct before proceeding.
+        # Since the presave fixture is in effect, we should expect that
+        # is_election_board is equal to none, even though we tried to define it as "True"
+        expected_values = {
+            'is_election_board': False,
+            'generic_org_type': DomainRequest.OrganizationChoices.CITY,
+            'organization_type': DomainRequest.OrgChoicesElectionOffice.CITY,
+        }
+        self.assert_expected_org_values_on_request_and_info(city_request, city_info, expected_values)
+
+        # Run the populate script
+        try:
+            self.run_populate_organization_type()
+        except Exception as e:
+            self.fail(f"Could not run populate_organization_type script. Failed with exception: {e}")
+        
+        # All values should be the same
+        self.assert_expected_org_values_on_request_and_info(city_request, city_info, expected_values)
+
+    def test_request_and_info_federal(self):
+        """Tests what happens to a federal domain with is_election_board=True (which should be reverted to None)"""
+        federal_request = self.domain_request_1
+        federal_info = self.domain_info_1
+
+        # Make sure that all data is correct before proceeding.
+        # Since the presave fixture is in effect, we should expect that
+        # is_election_board is equal to none, even though we tried to define it as "True"
+        # TODO - either add some logging if this happens or a ValueError in the original ticket.
+        # Or FSM? Probably FSM. 
+        expected_values = {
+            'is_election_board': None,
+            'generic_org_type': DomainRequest.OrganizationChoices.FEDERAL,
+            'organization_type': DomainRequest.OrgChoicesElectionOffice.FEDERAL,
+        }
+        self.assert_expected_org_values_on_request_and_info(federal_request, federal_info, expected_values)
+
+        # Run the populate script
+        try:
+            self.run_populate_organization_type()
+        except Exception as e:
+            self.fail(f"Could not run populate_organization_type script. Failed with exception: {e}")
+        
+        # All values should be the same
+        self.assert_expected_org_values_on_request_and_info(federal_request, federal_info, expected_values)
+
+    def test_request_and_info_tribal_add_election_office(self):
+        """
+        Tests if a tribal domain in the election csv changes organization_type to TRIBAL - ELECTION
+        for the domain request and the domain info
+        """
+
+        tribal_request = self.domain_request_3
+        tribal_info = self.domain_info_3
+
+        # Make sure that all data is correct before proceeding.
+        # Because the presave fixture is in place when creating this, we should expect that the
+        # organization_type variable is already pre-populated. We will test what happens when
+        # it is not in another test.
+
+        expected_values = {
+            'is_election_board': False,
+            'generic_org_type': DomainRequest.OrganizationChoices.TRIBAL,
+            'organization_type': DomainRequest.OrgChoicesElectionOffice.TRIBAL
+        }
+        self.assert_expected_org_values_on_request_and_info(tribal_request, tribal_info, expected_values)
+
+        # Run the populate script
+        try:
+            self.run_populate_organization_type()
+        except Exception as e:
+            self.fail(f"Could not run populate_organization_type script. Failed with exception: {e}")
+        
+        # Because we define this in the "csv", we expect that is election board will switch to True,
+        # and organization_type will now be tribal_election
+        expected_values["is_election_board"] = True
+        expected_values["organization_type"] = DomainRequest.OrgChoicesElectionOffice.TRIBAL_ELECTION
+
+        self.assert_expected_org_values_on_request_and_info(tribal_request, tribal_info, expected_values)
+
+    def test_request_and_info_tribal_remove_election_office(self):
+        """
+        Tests if a tribal domain in the election csv changes organization_type to TRIBAL
+        when it used to be TRIBAL - ELECTION
+        for the domain request and the domain info
+        """
+
+        tribal_election_request = self.domain_request_4
+        tribal_election_info = self.domain_info_4
+
+        # Make sure that all data is correct before proceeding.
+        # Because the presave fixture is in place when creating this, we should expect that the
+        # organization_type variable is already pre-populated. We will test what happens when
+        # it is not in another test.
+        expected_values = {
+            'is_election_board': True,
+            'generic_org_type': DomainRequest.OrganizationChoices.TRIBAL,
+            'organization_type': DomainRequest.OrgChoicesElectionOffice.TRIBAL_ELECTION
+        }
+        self.assert_expected_org_values_on_request_and_info(tribal_election_request, tribal_election_info, expected_values)
+
+        # Run the populate script
+        try:
+            self.run_populate_organization_type()
+        except Exception as e:
+            self.fail(f"Could not run populate_organization_type script. Failed with exception: {e}")
+        
+        # Because we don't define this in the "csv", we expect that is election board will switch to False,
+        # and organization_type will now be tribal
+        expected_values["is_election_board"] = False
+        expected_values["organization_type"] = DomainRequest.OrgChoicesElectionOffice.TRIBAL
+        self.assert_expected_org_values_on_request_and_info(tribal_election_request, tribal_election_info, expected_values)
+
+    @skip("TODO")
+    def test_transition_data(self):
+        """Tests for how this script interacts with prexisting data (for instance, stable)"""
+        # Make instead of mocking we can literally just run the transition domain scripts?
+        pass
 
 
 class TestPopulateFirstReady(TestCase):