Add legacy Export wins migration commands.

datahub/dbmaintenance/management/commands/update_legacy_export_wins_company_link.py

@@ -0,0 +1,39 @@
+import reversion
+
+from django.core.management.base import CommandError
+
+from datahub.company.models import Company
+from datahub.dbmaintenance.management.base import CSVBaseCommand
+from datahub.dbmaintenance.utils import parse_uuid
+from datahub.export_win.models import Win
+
+
+class Command(CSVBaseCommand):
+    """Command to update legacy Export Win mapping to Data Hub Company."""
+
+    def _process_row(self, row, simulate=False, **options):
+        """Process one single row."""
+        export_win_id = parse_uuid(row['export_win_id'])
+        company_id = parse_uuid(row['data_hub_id']) if row['data_hub_id'] else None
+
+        if company_id:
+            try:
+                Company.objects.get(pk=company_id)
+            except Company.DoesNotExist:
+                raise CommandError(f'Company with ID {company_id} does not exist')
+
+            export_win = Win.objects.get(id=export_win_id)
+
+            if export_win.company_id == company_id:
+                return
+
+            export_win.company_id = company_id
+
+            if not simulate:
+                with reversion.create_revision():
+                    export_win.save(
+                        update_fields=(
+                            'company_id',
+                        ),
+                    )
+                    reversion.set_comment('Legacy export wins company migration.')

datahub/dbmaintenance/management/commands/update_legacy_export_wins_data.py

@@ -0,0 +1,41 @@
+import reversion
+
+from datahub.dbmaintenance.management.base import CSVBaseCommand
+from datahub.dbmaintenance.utils import parse_uuid
+from datahub.export_win.models import Win
+
+
+class Command(CSVBaseCommand):
+    """Command to update legacy Export Win data."""
+
+    def _process_row(self, row, simulate=False, **options):
+        """Process one single row."""
+        export_win_id = parse_uuid(row['id'])
+
+        export_win = Win.objects.get(id=export_win_id)
+        export_win.company_name = row['company_name']
+        export_win.lead_officer_name = row['lead_officer_name']
+        export_win.lead_officer_email_address = row['lead_officer_email_address']
+        export_win.adviser_name = row['user_name']
+        export_win.adviser_email_address = row['user_email']
+        export_win.line_manager_name = row['line_manager_name']
+        export_win.customer_name = row['customer_name']
+        export_win.customer_job_title = row['customer_job_title']
+        export_win.customer_email_address = row['customer_email_address']
+
+        if not simulate:
+            with reversion.create_revision():
+                export_win.save(
+                    update_fields=(
+                        'company_name',
+                        'lead_officer_name',
+                        'lead_officer_email_address',
+                        'adviser_name',
+                        'adviser_email_address',
+                        'line_manager_name',
+                        'customer_name',
+                        'customer_job_title',
+                        'customer_email_address',
+                    ),
+                )
+                reversion.set_comment('Legacy export wins data migration.')

datahub/dbmaintenance/test/commands/test_update_legacy_export_wins_company_link.py

@@ -0,0 +1,102 @@
+from io import BytesIO
+from uuid import uuid4
+
+import pytest
+from django.core.management import call_command
+
+from reversion.models import Version
+
+from datahub.company.test.factories import CompanyFactory
+from datahub.export_win.models import Win
+from datahub.export_win.test.factories import WinFactory
+
+
+pytestmark = pytest.mark.django_db
+
+
+def test_run(s3_stubber, caplog):
+    """Test that the command updates the specified records (ignoring ones with errors)."""
+    caplog.set_level('ERROR')
+    companies = CompanyFactory.create_batch(4)
+    wins = WinFactory.create_batch(4, company=None)
+
+    uuids = [win.id for win in wins]
+    company_uuids = [company.id for company in companies]
+
+    bucket = 'test_bucket'
+    object_key = 'test_key'
+    csv_contents = ['export_win_id,data_hub_id']
+    for uuid, company_uuid in zip(uuids, company_uuids):
+        csv_contents.append(f'{uuid},{company_uuid}')
+
+    bad_company_id = uuid4()
+    csv_contents.append(f'{uuid4()},{bad_company_id}')
+
+    csv_content = '\n'.join(csv_contents)
+
+    s3_stubber.add_response(
+        'get_object',
+        {
+            'Body': BytesIO(csv_content.encode(encoding='utf-8')),
+        },
+        expected_params={
+            'Bucket': bucket,
+            'Key': object_key,
+        },
+    )
+
+    call_command('update_legacy_export_wins_company_link', bucket, object_key)
+
+    for uuid, company in zip(uuids, companies):
+        win = Win.objects.get(id=uuid)
+        assert win.company_id == company.id
+
+        versions = Version.objects.get_for_object(win)
+        assert versions.count() == 1
+        comment = versions[0].revision.get_comment()
+        assert comment == 'Legacy export wins company migration.'
+
+    assert f'Company with ID {bad_company_id} does not exist' in caplog.text
+
+
+def test_simulate(s3_stubber, caplog):
+    """Test that the command simulates updates if --simulate is passed in."""
+    caplog.set_level('ERROR')
+    companies = CompanyFactory.create_batch(4)
+    wins = WinFactory.create_batch(4, company=None)
+
+    uuids = [win.id for win in wins]
+    company_uuids = [company.id for company in companies]
+
+    bucket = 'test_bucket'
+    object_key = 'test_key'
+    csv_contents = ['export_win_id,data_hub_id']
+    for uuid, company_uuid in zip(uuids, company_uuids):
+        csv_contents.append(f'{uuid},{company_uuid}')
+
+    bad_company_id = uuid4()
+    csv_contents.append(f'{uuid4()},{bad_company_id}')
+
+    csv_content = '\n'.join(csv_contents)
+
+    s3_stubber.add_response(
+        'get_object',
+        {
+            'Body': BytesIO(csv_content.encode(encoding='utf-8')),
+        },
+        expected_params={
+            'Bucket': bucket,
+            'Key': object_key,
+        },
+    )
+
+    call_command('update_legacy_export_wins_company_link', bucket, object_key, simulate=True)
+
+    for uuid in uuids:
+        win = Win.objects.get(id=uuid)
+        assert win.company_id is None
+
+        versions = Version.objects.get_for_object(win)
+        assert versions.count() == 0
+
+    assert f'Company with ID {bad_company_id} does not exist' in caplog.text