manage.get.gov/src/registrar/tests/test_reports.py

from django.test import TestCase
from django.core.files import File
from io import StringIO
import csv
from registrar.models.domain_information import DomainInformation
from registrar.models.domain import Domain
from registrar.models.user import User
from django.contrib.auth import get_user_model
from registrar.utility.csv_export import (
    export_domains_to_writer,
    export_data_type_to_csv,
    export_data_full_to_csv,
    export_data_federal_to_csv,
)

class ExportDataTest(TestCase):
    def setUp(self):
        
        username = "test_user"
        first_name = "First"
        last_name = "Last"
        email = "info@example.com"
        self.user = get_user_model().objects.create(
            username=username, first_name=first_name, last_name=last_name, email=email
        )

        self.domain_1, _ = Domain.objects.get_or_create(name="cdomain1.gov", state=Domain.State.READY)
        self.domain_2, _ = Domain.objects.get_or_create(name="adomain2.gov", state=Domain.State.READY)
        self.domain_3, _ = Domain.objects.get_or_create(name="ddomain3.gov", state=Domain.State.READY)
        self.domain_4, _ = Domain.objects.get_or_create(name="bdomain4.gov", state=Domain.State.UNKNOWN)
        self.domain_4, _ = Domain.objects.get_or_create(name="bdomain4.gov", state=Domain.State.UNKNOWN)
        
        self.domain_information_1, _ = DomainInformation.objects.get_or_create(
            creator=self.user,
            domain=self.domain_1,
            organization_type="federal",
            federal_agency="World War I Centennial Commission",
        )
        self.domain_information_2, _ = DomainInformation.objects.get_or_create(
            creator=self.user,
            domain=self.domain_2,
            organization_type="interstate",
        )
        self.domain_information_3, _ = DomainInformation.objects.get_or_create(
            creator=self.user,
            domain=self.domain_3,
            organization_type="federal",
            federal_agency="Armed Forces Retirement Home",
        )
        self.domain_information_4, _ = DomainInformation.objects.get_or_create(
            creator=self.user,
            domain=self.domain_4,
            organization_type="federal",
            federal_agency="Armed Forces Retirement Home",
        )
        
    def tearDown(self):
        Domain.objects.all().delete()
        DomainInformation.objects.all().delete()
        User.objects.all().delete()
        super().tearDown()


    def test_export_domains_to_writer(self):
        """Test that export_domains_to_writer returns the
        existing domain, test that sort by domain name works,
        test that filter works"""
        # Create a CSV file in memory
        csv_file = StringIO()
        writer = csv.writer(csv_file)
        
        # Define columns, sort fields, and filter condition
        columns = [
            "Domain name",
            "Domain type",
            "Federal agency",
            "Organization name",
            "City",
            "State",
            "AO",
            "AO email",
            "Submitter",
            "Submitter title",
            "Submitter email",
            "Submitter phone",
            "Security Contact Email",
            "Status",
        ]
        sort_fields = ["domain__name"]
        filter_condition = {"domain__state": Domain.State.READY}

        # Call the export function
        export_domains_to_writer(writer, columns, sort_fields, filter_condition)

        # Reset the CSV file's position to the beginning
        csv_file.seek(0)
        
        # Read the content into a variable
        csv_content = csv_file.read()
        
        # We expect READY domains,
        # sorted alphabetially by domain name
        expected_content = """\
Domain name,Domain type,Federal agency,Organization name,City,State,AO,AO email,Submitter,Submitter title,Submitter email,Submitter phone,Security Contact Email,Status
adomain2.gov,interstate,,,,, , , , , , , ,ready
cdomain1.gov,federal,World War I Centennial Commission,,,, , , , , , , ,ready
ddomain3.gov,federal,Armed Forces Retirement Home,,,, , , , , , , ,ready
"""
        # print(csv_content)
        # self.maxDiff = None
        
        # Normalize line endings and remove leading/trailing whitespace
        csv_content = csv_content.replace('\r\n', '\n').strip()
        expected_content = expected_content.strip()
    
        self.assertEqual(csv_content, expected_content)
        
    def test_export_domains_to_writer_2(self):
        """An additional test for filters and multi-column sort"""
        # Create a CSV file in memory
        csv_file = StringIO()
        writer = csv.writer(csv_file)
        
        # Define columns, sort fields, and filter condition
        columns = [
            "Domain name",
            "Domain type",
            "Federal agency",
            "Organization name",
            "City",
            "State",
            "Security Contact Email",
        ]
        sort_fields = ["domain__name", "federal_agency", "organization_type"]
        filter_condition = {"organization_type__icontains": "federal", "domain__state": Domain.State.READY}

        # Call the export function
        export_domains_to_writer(writer, columns, sort_fields, filter_condition)

        # Reset the CSV file's position to the beginning
        csv_file.seek(0)
        
        # Read the content into a variable
        csv_content = csv_file.read()
        
        # We expect READY domains,
        # federal only
        # sorted alphabetially by domain name
        expected_content = """\
Domain name,Domain type,Federal agency,Organization name,City,State,Security Contact Email
cdomain1.gov,federal,World War I Centennial Commission,,,, 
ddomain3.gov,federal,Armed Forces Retirement Home,,,,
"""
        # print(csv_content)
        # self.maxDiff = None
        
        # Normalize line endings and remove leading/trailing whitespace
        csv_content = csv_content.replace('\r\n', '\n').strip()
        expected_content = expected_content.strip()
    
        self.assertEqual(csv_content, expected_content)