diff --git a/datahub/dbmaintenance/management/commands/update_investment_project_company.py b/datahub/dbmaintenance/management/commands/update_investment_project_company.py new file mode 100644 index 000000000..5a4a1641b --- /dev/null +++ b/datahub/dbmaintenance/management/commands/update_investment_project_company.py @@ -0,0 +1,95 @@ +import uuid +import reversion + +from datahub.investment.models import InvestmentProject +from ..base import CSVBaseCommand + + +class Command(CSVBaseCommand): + """Command to update investment_project. + + investor_company, intermediate_company, uk_company, uk_company_decided. + """ + + def add_arguments(self, parser): + """Define extra arguments.""" + super().add_arguments(parser) + parser.add_argument( + '--simulate', + action='store_true', + dest='simulate', + default=False, + help='If True it only simulates the command without saving the changes.', + ) + + def _parse_company_id(self, company_id): + """ + :param company_id: string representing uuid of the company + :return: instance of UUID or None + """ + if not company_id or company_id.lower().strip() == 'null': + return None + return uuid.UUID(company_id) + + def _should_update(self, + investment_project, + investor_company_id, + intermediate_company_id, + uk_company_id, + uk_company_decided): + """ + Checks if Investment project should be updated. + + :param investment_project: instance of InvestmentProject + :param investor_company: instance of Company or None + :param intermediate_company: instance of Company or None + :param uk_company: instance of Company or None + :param uk_company_decided: Boolean + :return: True if investment project needs to be updated + """ + return (investment_project.investor_company_id != investor_company_id or + investment_project.intermediate_company_id != intermediate_company_id or + investment_project.uk_company_id != uk_company_id or + investment_project.uk_company_decided != uk_company_decided) + + def get_uk_company_decided(self, uk_company_decided): + """ + :param uk_company_decided: string containing either '1' or '0' + :return: Boolean + """ + translate = { + '0': False, + '1': True, + } + return translate[uk_company_decided.strip()] + + def _process_row(self, row, simulate=False, **options): + """Process one single row.""" + investment_project = InvestmentProject.objects.get(pk=row['id']) + investor_company_id = self._parse_company_id(row['investor_company_id']) + intermediate_company_id = self._parse_company_id(row['intermediate_company_id']) + uk_company_id = self._parse_company_id(row['uk_company_id']) + uk_company_decided = self.get_uk_company_decided(row['uk_company_decided']) + + if self._should_update( + investment_project, + investor_company_id, + intermediate_company_id, + uk_company_id, + uk_company_decided, + ): + investment_project.investor_company_id = investor_company_id + investment_project.intermediate_company_id = intermediate_company_id + investment_project.uk_company_id = uk_company_id + investment_project.uk_company_decided = uk_company_decided + if not simulate: + with reversion.create_revision(): + investment_project.save( + update_fields=( + 'investor_company', + 'intermediate_company', + 'uk_company', + 'uk_company_decided', + ) + ) + reversion.set_comment('Companies data migration.') diff --git a/datahub/dbmaintenance/management/commands/update_investment_project_referral_source_activity_marketing.py b/datahub/dbmaintenance/management/commands/update_investment_project_referral_source_activity_marketing.py new file mode 100644 index 000000000..03b398e14 --- /dev/null +++ b/datahub/dbmaintenance/management/commands/update_investment_project_referral_source_activity_marketing.py @@ -0,0 +1,105 @@ +from functools import lru_cache + +import reversion + +from datahub.investment.models import InvestmentProject +from datahub.metadata.models import ReferralSourceActivity, ReferralSourceMarketing +from ..base import CSVBaseCommand + + +class Command(CSVBaseCommand): + """Command to update investment_project.referral_source_activity/_marketing.""" + + def add_arguments(self, parser): + """Define extra arguments.""" + super().add_arguments(parser) + parser.add_argument( + '--simulate', + action='store_true', + dest='simulate', + default=False, + help='If True it only simulates the command without saving the changes.', + ) + + @lru_cache(maxsize=None) + def get_referral_source_activity(self, referral_source_activity_id): + """ + :param referral_source_activity_id: uuid of the referral source activity + :return: instance of ReferralSourceActivity + with id == referral_source_activity_id if it exists, + None otherwise + """ + if ( + not referral_source_activity_id or + referral_source_activity_id.lower().strip() == 'null' + ): + return None + return ReferralSourceActivity.objects.get(id=referral_source_activity_id) + + @lru_cache(maxsize=None) + def get_referral_source_activity_marketing(self, referral_source_activity_marketing_id): + """ + :param referral_source_activity_marketing_id: uuid of the referral source marketing + :return: instance of ReferralSourceMarketing + with id == referral_source_activity_marketing_id if it exists, + None otherwise + """ + if ( + not referral_source_activity_marketing_id or + referral_source_activity_marketing_id.lower().strip() == 'null' + ): + return None + return ReferralSourceMarketing.objects.get(id=referral_source_activity_marketing_id) + + def _should_update(self, + investment_project, + referral_source_activity, + referral_source_activity_marketing + ): + """ + Checks if Investment project should be updated. + + :param investment_project: instance of InvestmentProject + :param referral_source_activity: instance of ReferralSourceActivity or None + :param referral_source_activity_marketing: + instance of ReferralSourceMarketing or None + :return: True if investment project needs to be updated + """ + return ( + investment_project.referral_source_activity_id != + referral_source_activity.id or + investment_project.referral_source_activity_marketing_id != + referral_source_activity_marketing.id + ) + + def _process_row(self, row, simulate=False, **options): + """Process one single row.""" + investment_project = InvestmentProject.objects.get(pk=row['id']) + + referral_source_activity = self.get_referral_source_activity( + row['referral_source_activity_id'] + ) + + referral_source_activity_marketing = self.get_referral_source_activity_marketing( + row['referral_source_activity_marketing_id'] + ) + + if self._should_update( + investment_project, + referral_source_activity, + referral_source_activity_marketing, + ): + investment_project.referral_source_activity = \ + referral_source_activity + investment_project.referral_source_activity_marketing = \ + referral_source_activity_marketing + + if not simulate: + with reversion.create_revision(): + investment_project.save( + update_fields=( + 'referral_source_activity', + 'referral_source_activity_marketing', + ) + ) + reversion.set_comment('ReferralSourceActivityMarketing migration.') diff --git a/datahub/dbmaintenance/management/commands/update_investment_project_referral_source_activity_website.py b/datahub/dbmaintenance/management/commands/update_investment_project_referral_source_activity_website.py new file mode 100644 index 000000000..33763456f --- /dev/null +++ b/datahub/dbmaintenance/management/commands/update_investment_project_referral_source_activity_website.py @@ -0,0 +1,103 @@ +from functools import lru_cache + +import reversion + +from datahub.investment.models import InvestmentProject +from datahub.metadata.models import ReferralSourceActivity, ReferralSourceWebsite +from ..base import CSVBaseCommand + + +class Command(CSVBaseCommand): + """Command to update investment_project.referral_source_activity/_website.""" + + def add_arguments(self, parser): + """Define extra arguments.""" + super().add_arguments(parser) + parser.add_argument( + '--simulate', + action='store_true', + dest='simulate', + default=False, + help='If True it only simulates the command without saving the changes.', + ) + + @lru_cache(maxsize=None) + def get_referral_source_activity(self, referral_source_activity_id): + """ + :param referral_source_activity_id: uuid of the referral source activity + :return: instance of ReferralSourceActivity + with id == referral_source_activity_id if it exists, + None otherwise + """ + if ( + not referral_source_activity_id or + referral_source_activity_id.lower().strip() == 'null' + ): + return None + return ReferralSourceActivity.objects.get(id=referral_source_activity_id) + + @lru_cache(maxsize=None) + def get_referral_source_activity_website(self, referral_source_activity_website_id): + """ + :param referral_source_activity_website_id: uuid of the referral source website + :return: instance of ReferralSourceWebsite + with id == referral_source_activity_website_id if it exists, + None otherwise + """ + if ( + not referral_source_activity_website_id or + referral_source_activity_website_id.lower().strip() == 'null' + ): + return None + return ReferralSourceWebsite.objects.get(id=referral_source_activity_website_id) + + def _should_update(self, + investment_project, + referral_source_activity, + referral_source_activity_website + ): + """ + Checks if Investment project should be updated. + + :param investment_project: instance of InvestmentProject + :param referral_source_activity: instance of ReferralSourceActivity or None + :param referral_source_activity_website: + instance of ReferralSourceWebsite or None + :return: True if investment project needs to be updated + """ + return ( + investment_project.referral_source_activity_id != + referral_source_activity.id or + investment_project.referral_source_activity_website_id != + referral_source_activity_website.id + ) + + def _process_row(self, row, simulate=False, **options): + """Process one single row.""" + investment_project = InvestmentProject.objects.get(pk=row['id']) + + referral_source_activity = self.get_referral_source_activity( + row['referral_source_activity_id'] + ) + + referral_source_activity_website = self.get_referral_source_activity_website( + row['referral_source_activity_website_id'] + ) + + if self._should_update( + investment_project, + referral_source_activity, + referral_source_activity_website, + ): + investment_project.referral_source_activity = referral_source_activity + investment_project.referral_source_activity_website = referral_source_activity_website + + if not simulate: + with reversion.create_revision(): + investment_project.save( + update_fields=( + 'referral_source_activity', + 'referral_source_activity_website', + ) + ) + reversion.set_comment('ReferralSourceActivityWebsite migration.') diff --git a/datahub/dbmaintenance/management/commands/update_investment_project_sector.py b/datahub/dbmaintenance/management/commands/update_investment_project_sector.py new file mode 100644 index 000000000..1113d9f0b --- /dev/null +++ b/datahub/dbmaintenance/management/commands/update_investment_project_sector.py @@ -0,0 +1,67 @@ +from functools import lru_cache + +import reversion + +from datahub.investment.models import InvestmentProject +from datahub.metadata.models import Sector +from ..base import CSVBaseCommand + + +class Command(CSVBaseCommand): + """Command to update investment_project.sector.""" + + def add_arguments(self, parser): + """Define extra arguments.""" + super().add_arguments(parser) + parser.add_argument( + '--simulate', + action='store_true', + dest='simulate', + default=False, + help='If True it only simulates the command without saving the changes.', + ) + + @lru_cache(maxsize=None) + def get_sector(self, sector_id): + """ + :param company_id: uuid of the company + :return: instance of Company with id == company_id if it exists, + None otherwise + """ + if not sector_id or sector_id.lower().strip() == 'null': + return None + return Sector.objects.get(id=sector_id) + + def _should_update(self, + investment_project, + old_sector, + ): + """ + Checks if Investment project should be updated. + + :param investment_project: instance of InvestmentProject + :param old_sector: instance of Company or None + :return: True if investment project needs to be updated + """ + return (investment_project.sector == old_sector) + + def _process_row(self, row, simulate=False, **options): + """Process one single row.""" + investment_project = InvestmentProject.objects.get(pk=row['id']) + + old_sector = self.get_sector(row['old_sector']) + new_sector = self.get_sector(row['new_sector']) + + if self._should_update( + investment_project, + old_sector, + ): + investment_project.sector = new_sector + if not simulate: + with reversion.create_revision(): + investment_project.save( + update_fields=( + 'sector', + ) + ) + reversion.set_comment('Sector migration.') diff --git a/datahub/dbmaintenance/test/commands/test_update_investment_project_company.py b/datahub/dbmaintenance/test/commands/test_update_investment_project_company.py new file mode 100644 index 000000000..72eb1c3a4 --- /dev/null +++ b/datahub/dbmaintenance/test/commands/test_update_investment_project_company.py @@ -0,0 +1,208 @@ +from io import BytesIO + +import pytest +from django.core.management import call_command +from reversion.models import Version + +from datahub.company.test.factories import CompanyFactory +from datahub.investment.test.factories import InvestmentProjectFactory + +pytestmark = pytest.mark.django_db + + +def test_run(s3_stubber): + """Test that the command updates the relevant records ignoring ones with errors.""" + companies = CompanyFactory.create_batch(18) + + investment_projects = [ + # investment project in CSV doesn't exist so row should fail + + # fields should get updated + InvestmentProjectFactory( + investor_company=None, + intermediate_company=None, + uk_company=None, + uk_company_decided=False, + ), + # should be ignored + InvestmentProjectFactory( + investor_company=companies[0], + intermediate_company=companies[1], + uk_company=companies[2], + uk_company_decided=True, + ), + # investor_company_id is invalid so it should fail + InvestmentProjectFactory( + investor_company=companies[3], + intermediate_company=companies[4], + uk_company=companies[5], + uk_company_decided=True, + ), + # intermediate_company_id is invalid so it should fail + InvestmentProjectFactory( + investor_company=companies[6], + intermediate_company=companies[7], + uk_company=companies[8], + uk_company_decided=True, + ), + # uk_company_id is invalid so it should fail + InvestmentProjectFactory( + investor_company=companies[9], + intermediate_company=companies[10], + uk_company=companies[11], + uk_company_decided=True, + ), + # uk_company_decided is invalid so it should fail + InvestmentProjectFactory( + investor_company=companies[12], + intermediate_company=companies[13], + uk_company=companies[14], + uk_company_decided=True, + ), + # company fields NULL and uk_company_decided 0 + InvestmentProjectFactory( + investor_company=companies[15], + intermediate_company=companies[16], + uk_company=companies[17], + uk_company_decided=True, + ), + + ] + + file_companies = CompanyFactory.create_batch(3) + + bucket = 'test_bucket' + object_key = 'test_key' + csv_content = f"""id,investor_company_id,intermediate_company_id,uk_company_id,uk_company_decided +00000000-0000-0000-0000-000000000000,NULL,NULL,NULL,NULL +{investment_projects[0].id},{file_companies[0].pk},{file_companies[1].pk},{file_companies[2].pk},1 +{investment_projects[2].id},invalid_id,{file_companies[1].pk},{file_companies[2].pk},0 +{investment_projects[3].id},{file_companies[0].pk},invalid_id,{file_companies[2].pk},0 +{investment_projects[4].id},{file_companies[0].pk},{file_companies[1].pk},invalid_id,0 +{investment_projects[5].id},{file_companies[0].pk},{file_companies[1].pk},{file_companies[2].pk},THIS +{investment_projects[6].id},NULL,NULL,NULL,0 +""" + + s3_stubber.add_response( + 'get_object', + { + 'Body': BytesIO(bytes(csv_content, encoding='utf-8')) + }, + expected_params={ + 'Bucket': bucket, + 'Key': object_key + } + ) + + call_command('update_investment_project_company', bucket, object_key) + + for investment_project in investment_projects: + investment_project.refresh_from_db() + + assert investment_projects[0].investor_company == file_companies[0] + assert investment_projects[0].intermediate_company == file_companies[1] + assert investment_projects[0].uk_company == file_companies[2] + assert investment_projects[0].uk_company_decided is True + assert investment_projects[1].investor_company == companies[0] + assert investment_projects[1].intermediate_company == companies[1] + assert investment_projects[1].uk_company == companies[2] + assert investment_projects[1].uk_company_decided is True + assert investment_projects[2].investor_company == companies[3] + assert investment_projects[2].intermediate_company == companies[4] + assert investment_projects[2].uk_company == companies[5] + assert investment_projects[2].uk_company_decided is True + assert investment_projects[3].investor_company == companies[6] + assert investment_projects[3].intermediate_company == companies[7] + assert investment_projects[3].uk_company == companies[8] + assert investment_projects[3].uk_company_decided is True + assert investment_projects[4].investor_company == companies[9] + assert investment_projects[4].intermediate_company == companies[10] + assert investment_projects[4].uk_company == companies[11] + assert investment_projects[4].uk_company_decided is True + assert investment_projects[5].investor_company == companies[12] + assert investment_projects[5].intermediate_company == companies[13] + assert investment_projects[5].uk_company == companies[14] + assert investment_projects[5].uk_company_decided is True + assert investment_projects[6].investor_company is None + assert investment_projects[6].intermediate_company is None + assert investment_projects[6].uk_company is None + assert investment_projects[6].uk_company_decided is False + + +def test_simulate(s3_stubber): + """Test that the command only simulates the actions if --simulate is passed in.""" + companies = CompanyFactory.create_batch(3) + investment_projects = InvestmentProjectFactory.create_batch( + 2, + investor_company=companies[0], + intermediate_company=companies[1], + uk_company=companies[2], + uk_company_decided=True, + ) + file_companies = CompanyFactory.create_batch(3) + + bucket = 'test_bucket' + object_key = 'test_key' + csv_content = f"""id,investor_company_id,intermediate_company_id,uk_company_id,uk_company_decided +{investment_projects[0].id},{file_companies[0].pk},{file_companies[1].pk},{file_companies[2].pk},1 +{investment_projects[1].id},{file_companies[0].pk},{file_companies[1].pk},{file_companies[2].pk},1 +""" + s3_stubber.add_response( + 'get_object', + { + 'Body': BytesIO(bytes(csv_content, encoding='utf-8')) + }, + expected_params={ + 'Bucket': bucket, + 'Key': object_key + } + ) + + call_command('update_investment_project_company', bucket, object_key, simulate=True) + + for investment_project in investment_projects: + investment_project.refresh_from_db() + + assert investment_projects[0].investor_company == companies[0] + assert investment_projects[0].intermediate_company == companies[1] + assert investment_projects[0].uk_company == companies[2] + assert investment_projects[0].uk_company_decided is True + assert investment_projects[1].investor_company == companies[0] + assert investment_projects[1].intermediate_company == companies[1] + assert investment_projects[1].uk_company == companies[2] + assert investment_projects[1].uk_company_decided is True + + +def test_audit_log(s3_stubber): + """Test that the audit log is being created.""" + investment_project = InvestmentProjectFactory() + file_companies = CompanyFactory.create_batch(3) + + bucket = 'test_bucket' + object_key = 'test_key' + csv_content = f"""id,investor_company_id,intermediate_company_id,uk_company_id,uk_company_decided +{investment_project.id},{file_companies[0].pk},{file_companies[1].pk},{file_companies[2].pk},1 +""" + s3_stubber.add_response( + 'get_object', + { + 'Body': BytesIO(bytes(csv_content, encoding='utf-8')) + }, + expected_params={ + 'Bucket': bucket, + 'Key': object_key + } + ) + + call_command('update_investment_project_company', bucket, object_key) + + investment_project.refresh_from_db() + + assert investment_project.investor_company == file_companies[0] + assert investment_project.intermediate_company == file_companies[1] + assert investment_project.uk_company == file_companies[2] + assert investment_project.uk_company_decided is True + + versions = Version.objects.get_for_object(investment_project) + assert len(versions) == 1 + assert versions[0].revision.comment == 'Companies data migration.' diff --git a/datahub/dbmaintenance/test/commands/test_update_investment_project_referral_source_activity_marketing.py b/datahub/dbmaintenance/test/commands/test_update_investment_project_referral_source_activity_marketing.py new file mode 100644 index 000000000..59c4d9c0a --- /dev/null +++ b/datahub/dbmaintenance/test/commands/test_update_investment_project_referral_source_activity_marketing.py @@ -0,0 +1,204 @@ +from io import BytesIO + +import factory +import pytest +from django.core.management import call_command +from reversion.models import Version + +from datahub.investment.test.factories import InvestmentProjectFactory +from datahub.metadata.test.factories import ( + ReferralSourceActivityFactory, + ReferralSourceMarketingFactory +) + +pytestmark = pytest.mark.django_db + + +def test_run(s3_stubber): + """Test that the command updates the relevant records ignoring ones with errors.""" + referral_source_activities = ReferralSourceActivityFactory.create_batch(5) + referral_source_marketings = ReferralSourceMarketingFactory.create_batch(5) + + investment_projects = [ + # investment project in CSV doesn't exist so row should fail + + # referral_source_activity and referral_source_marketing should get updated + InvestmentProjectFactory( + referral_source_activity_id=referral_source_activities[0].id, + referral_source_activity_marketing=referral_source_marketings[0] + ), + # referral_source_activity and referral_source_marketing should get updated + InvestmentProjectFactory( + referral_source_activity_id=referral_source_activities[1].id, + referral_source_activity_marketing=referral_source_marketings[1] + ), + # should be ignored + InvestmentProjectFactory( + referral_source_activity_id=referral_source_activities[2].id, + referral_source_activity_marketing=referral_source_marketings[2] + ), + # referral_source_activity is invalid so it should fail + InvestmentProjectFactory( + referral_source_activity_id=referral_source_activities[3].id, + referral_source_activity_marketing=referral_source_marketings[3] + ), + # referral_source_marketing is invalid so it should fail + InvestmentProjectFactory( + referral_source_activity_id=referral_source_activities[4].id, + referral_source_activity_marketing=referral_source_marketings[4] + ), + ] + + new_referral_activity = ReferralSourceActivityFactory() + new_referral_marketing = ReferralSourceMarketingFactory() + + bucket = 'test_bucket' + object_key = 'test_key' + csv_content = f"""id,referral_source_activity_id,referral_source_activity_marketing_id +00000000-0000-0000-0000-000000000000,{new_referral_activity.id},NULL +{investment_projects[0].id},{new_referral_activity.id},{new_referral_marketing.id} +{investment_projects[1].id},{new_referral_activity.id},NULL +{investment_projects[3].id},00000000-0000-0000-0000-000000000000,NULL +{investment_projects[4].id},NULL,00000000-0000-0000-0000-000000000000 +""" + + s3_stubber.add_response( + 'get_object', + { + 'Body': BytesIO(bytes(csv_content, encoding='utf-8')) + }, + expected_params={ + 'Bucket': bucket, + 'Key': object_key + } + ) + + call_command( + 'update_investment_project_referral_source_activity_marketing', + bucket, + object_key + ) + + for investment_project in investment_projects: + investment_project.refresh_from_db() + + assert ( + investment_projects[0].referral_source_activity == new_referral_activity + ) + assert ( + investment_projects[0].referral_source_activity_marketing == new_referral_marketing + ) + assert ( + investment_projects[1].referral_source_activity == new_referral_activity + ) + assert ( + investment_projects[1].referral_source_activity_marketing is None + ) + assert ( + investment_projects[2].referral_source_activity == referral_source_activities[2] + ) + assert ( + investment_projects[2].referral_source_activity_marketing == referral_source_marketings[2] + ) + assert ( + investment_projects[3].referral_source_activity == referral_source_activities[3] + ) + assert ( + investment_projects[3].referral_source_activity_marketing == referral_source_marketings[3] + ) + assert ( + investment_projects[4].referral_source_activity == referral_source_activities[4] + ) + assert ( + investment_projects[4].referral_source_activity_marketing == referral_source_marketings[4] + ) + + +def test_simulate(s3_stubber): + """Test that the command only simulates the actions if --simulate is passed in.""" + referral_source_activities = ReferralSourceActivityFactory.create_batch(2) + referral_source_marketings = ReferralSourceMarketingFactory.create_batch(2) + + investment_projects = InvestmentProjectFactory.create_batch( + 2, + referral_source_activity_id=factory.Iterator( + referral_source_activity.id for referral_source_activity in referral_source_activities + ), + referral_source_activity_marketing=factory.Iterator(referral_source_marketings), + ) + new_referral_activity = ReferralSourceActivityFactory() + new_referral_marketing = ReferralSourceMarketingFactory() + + bucket = 'test_bucket' + object_key = 'test_key' + csv_content = f"""id,referral_source_activity_id,referral_source_activity_marketing_id +{investment_projects[0].id},{new_referral_activity.id},{new_referral_marketing.id} +{investment_projects[1].id},{new_referral_activity.id},NULL +""" + s3_stubber.add_response( + 'get_object', + { + 'Body': BytesIO(bytes(csv_content, encoding='utf-8')) + }, + expected_params={ + 'Bucket': bucket, + 'Key': object_key + } + ) + + call_command( + 'update_investment_project_referral_source_activity_marketing', + bucket, + object_key, + simulate=True + ) + + for investment_project in investment_projects: + investment_project.refresh_from_db() + + for i in range(2): + assert ( + investment_projects[i].referral_source_activity == referral_source_activities[i] + ) + assert ( + investment_projects[i].referral_source_activity_marketing == + referral_source_marketings[i] + ) + + +def test_audit_log(s3_stubber): + """Test that audit log is being created.""" + investment_project = InvestmentProjectFactory() + new_referral_activity = ReferralSourceActivityFactory() + new_referral_marketing = ReferralSourceMarketingFactory() + + bucket = 'test_bucket' + object_key = 'test_key' + csv_content = f"""id,referral_source_activity_id,referral_source_activity_marketing_id +{investment_project.id},{new_referral_activity.id},{new_referral_marketing.id} +""" + s3_stubber.add_response( + 'get_object', + { + 'Body': BytesIO(bytes(csv_content, encoding='utf-8')) + }, + expected_params={ + 'Bucket': bucket, + 'Key': object_key + } + ) + + call_command( + 'update_investment_project_referral_source_activity_marketing', + bucket, + object_key, + ) + + investment_project.refresh_from_db() + + assert investment_project.referral_source_activity == new_referral_activity + assert investment_project.referral_source_activity_marketing == new_referral_marketing + + versions = Version.objects.get_for_object(investment_project) + assert len(versions) == 1 + assert versions[0].revision.comment == 'ReferralSourceActivityMarketing migration.' diff --git a/datahub/dbmaintenance/test/commands/test_update_investment_project_referral_source_activity_website.py b/datahub/dbmaintenance/test/commands/test_update_investment_project_referral_source_activity_website.py new file mode 100644 index 000000000..839828be2 --- /dev/null +++ b/datahub/dbmaintenance/test/commands/test_update_investment_project_referral_source_activity_website.py @@ -0,0 +1,179 @@ +from io import BytesIO + +import factory +import pytest +from django.core.management import call_command +from reversion.models import Version + +from datahub.investment.test.factories import InvestmentProjectFactory +from datahub.metadata.test.factories import ( + ReferralSourceActivityFactory, + ReferralSourceWebsiteFactory +) + +pytestmark = pytest.mark.django_db + + +def test_run(s3_stubber): + """Test that the command updates the relevant records ignoring ones with errors.""" + referral_source_activities = ReferralSourceActivityFactory.create_batch(5) + referral_source_websites = ReferralSourceWebsiteFactory.create_batch(5) + + investment_projects = [ + # investment project in CSV doesn't exist so row should fail + + # referral_source_activity and referral_source_website should get updated + InvestmentProjectFactory( + referral_source_activity_id=referral_source_activities[0].id, + referral_source_activity_website=referral_source_websites[0] + ), + # referral_source_activity and referral_source_website should get updated + InvestmentProjectFactory( + referral_source_activity_id=referral_source_activities[1].id, + referral_source_activity_website=referral_source_websites[1] + ), + # should be ignored + InvestmentProjectFactory( + referral_source_activity_id=referral_source_activities[2].id, + referral_source_activity_website=referral_source_websites[2] + ), + # referral_source_activity is invalid so it should fail + InvestmentProjectFactory( + referral_source_activity_id=referral_source_activities[3].id, + referral_source_activity_website=referral_source_websites[3] + ), + # referral_source_website is invalid so it should fail + InvestmentProjectFactory( + referral_source_activity_id=referral_source_activities[4].id, + referral_source_activity_website=referral_source_websites[4] + ), + ] + + new_referral_activity = ReferralSourceActivityFactory() + new_referral_website = ReferralSourceWebsiteFactory() + + bucket = 'test_bucket' + object_key = 'test_key' + csv_content = f"""id,referral_source_activity_id,referral_source_activity_website_id +00000000-0000-0000-0000-000000000000,{new_referral_activity.id},NULL +{investment_projects[0].id},{new_referral_activity.id},{new_referral_website.id} +{investment_projects[1].id},{new_referral_activity.id},NULL +{investment_projects[3].id},00000000-0000-0000-0000-000000000000,NULL +{investment_projects[4].id},NULL,00000000-0000-0000-0000-000000000000 +""" + + s3_stubber.add_response( + 'get_object', + { + 'Body': BytesIO(bytes(csv_content, encoding='utf-8')) + }, + expected_params={ + 'Bucket': bucket, + 'Key': object_key + } + ) + + call_command('update_investment_project_referral_source_activity_website', bucket, object_key) + + for investment_project in investment_projects: + investment_project.refresh_from_db() + + assert investment_projects[0].referral_source_activity == new_referral_activity + assert investment_projects[0].referral_source_activity_website == new_referral_website + assert investment_projects[1].referral_source_activity == new_referral_activity + assert investment_projects[1].referral_source_activity_website is None + assert investment_projects[2].referral_source_activity == referral_source_activities[2] + assert investment_projects[2].referral_source_activity_website == referral_source_websites[2] + assert investment_projects[3].referral_source_activity == referral_source_activities[3] + assert investment_projects[3].referral_source_activity_website == referral_source_websites[3] + assert investment_projects[4].referral_source_activity == referral_source_activities[4] + assert investment_projects[4].referral_source_activity_website == referral_source_websites[4] + + +def test_simulate(s3_stubber): + """Test that the command only simulates the actions if --simulate is passed in.""" + referral_source_activities = ReferralSourceActivityFactory.create_batch(2) + referral_source_websites = ReferralSourceWebsiteFactory.create_batch(2) + + investment_projects = InvestmentProjectFactory.create_batch( + 2, + referral_source_activity_id=factory.Iterator( + referral_source_activity.id for referral_source_activity in referral_source_activities + ), + referral_source_activity_website=factory.Iterator(referral_source_websites), + ) + new_referral_activity = ReferralSourceActivityFactory() + new_referral_website = ReferralSourceWebsiteFactory() + + bucket = 'test_bucket' + object_key = 'test_key' + csv_content = f"""id,referral_source_activity_id,referral_source_activity_website_id +{investment_projects[0].id},{new_referral_activity.id},{new_referral_website.id} +{investment_projects[1].id},{new_referral_activity.id},NULL +""" + s3_stubber.add_response( + 'get_object', + { + 'Body': BytesIO(bytes(csv_content, encoding='utf-8')) + }, + expected_params={ + 'Bucket': bucket, + 'Key': object_key + } + ) + + call_command( + 'update_investment_project_referral_source_activity_website', + bucket, + object_key, + simulate=True + ) + + for investment_project in investment_projects: + investment_project.refresh_from_db() + + for i in range(2): + assert ( + investment_projects[i].referral_source_activity == referral_source_activities[i] + ) + assert ( + investment_projects[i].referral_source_activity_website == referral_source_websites[i] + ) + + +def test_audit_log(s3_stubber): + """Test that audit log is being created.""" + investment_project = InvestmentProjectFactory() + new_referral_activity = ReferralSourceActivityFactory() + new_referral_website = ReferralSourceWebsiteFactory() + + bucket = 'test_bucket' + object_key = 'test_key' + csv_content = f"""id,referral_source_activity_id,referral_source_activity_website_id +{investment_project.id},{new_referral_activity.id},{new_referral_website.id} +""" + s3_stubber.add_response( + 'get_object', + { + 'Body': BytesIO(bytes(csv_content, encoding='utf-8')) + }, + expected_params={ + 'Bucket': bucket, + 'Key': object_key + } + ) + + call_command( + 'update_investment_project_referral_source_activity_website', + bucket, + object_key, + ) + + investment_project.refresh_from_db() + + assert investment_project.referral_source_activity == new_referral_activity + assert investment_project.referral_source_activity_website == new_referral_website + + versions = Version.objects.get_for_object(investment_project) + assert len(versions) == 1 + assert versions[0].revision.comment == 'ReferralSourceActivityWebsite migration.' diff --git a/datahub/dbmaintenance/test/commands/test_update_investment_project_sector.py b/datahub/dbmaintenance/test/commands/test_update_investment_project_sector.py new file mode 100644 index 000000000..437872cb5 --- /dev/null +++ b/datahub/dbmaintenance/test/commands/test_update_investment_project_sector.py @@ -0,0 +1,128 @@ +from io import BytesIO + +import pytest +from django.core.management import call_command +from reversion.models import Version + +from datahub.investment.test.factories import InvestmentProjectFactory +from datahub.metadata.test.factories import SectorFactory + +pytestmark = pytest.mark.django_db + + +def test_run(s3_stubber): + """Test that the command updates the relevant records ignoring ones with errors.""" + sectors = SectorFactory.create_batch(5) + + investment_projects = [ + # investment project in CSV doesn't exist so row should fail + + # sector should get updated + InvestmentProjectFactory(sector_id=sectors[0].id), + # sector should get updated + InvestmentProjectFactory(sector_id=None), + # sector should not get updated + InvestmentProjectFactory(sector_id=None), + # should be ignored + InvestmentProjectFactory(sector_id=sectors[3].id), + # should be skipped because of an error + InvestmentProjectFactory(sector_id=sectors[4].id), + ] + + new_sectors = SectorFactory.create_batch(5) + + bucket = 'test_bucket' + object_key = 'test_key' + csv_content = f"""id,old_sector,new_sector +00000000-0000-0000-0000-000000000000,NULL,NULL +{investment_projects[0].id},{sectors[0].id},{new_sectors[0].id} +{investment_projects[1].id},NULL,{new_sectors[1].id} +{investment_projects[2].id},{new_sectors[2].id},{new_sectors[2].id} +{investment_projects[4].id},invalid_id,another_invalid_id +""" + + s3_stubber.add_response( + 'get_object', + { + 'Body': BytesIO(bytes(csv_content, encoding='utf-8')) + }, + expected_params={ + 'Bucket': bucket, + 'Key': object_key + } + ) + + call_command('update_investment_project_sector', bucket, object_key) + + for investment_project in investment_projects: + investment_project.refresh_from_db() + + assert investment_projects[0].sector == new_sectors[0] + assert investment_projects[1].sector == new_sectors[1] + assert investment_projects[2].sector is None + assert investment_projects[3].sector == sectors[3] + assert investment_projects[4].sector == sectors[4] + + +def test_simulate(s3_stubber): + """Test that the command only simulates the actions if --simulate is passed in.""" + new_sectors = SectorFactory.create_batch(5) + investment_projects = InvestmentProjectFactory.create_batch(2) + old_sectors = [ip.sector for ip in investment_projects] + + bucket = 'test_bucket' + object_key = 'test_key' + csv_content = f"""id,old_sector,new_sector +{investment_projects[0].id},{old_sectors[0].id},{new_sectors[0].id} +{investment_projects[1].id},{old_sectors[1].id},{new_sectors[1].id} +""" + s3_stubber.add_response( + 'get_object', + { + 'Body': BytesIO(bytes(csv_content, encoding='utf-8')) + }, + expected_params={ + 'Bucket': bucket, + 'Key': object_key + } + ) + + call_command('update_investment_project_sector', bucket, object_key, simulate=True) + + for investment_project in investment_projects: + investment_project.refresh_from_db() + + assert investment_projects[0].sector == old_sectors[0] + assert investment_projects[1].sector == old_sectors[1] + + +def test_audit_log(s3_stubber): + """Test that audit log is being created.""" + new_sector = SectorFactory() + investment_project = InvestmentProjectFactory() + old_sector = investment_project.sector + + bucket = 'test_bucket' + object_key = 'test_key' + csv_content = f"""id,old_sector,new_sector +{investment_project.id},{old_sector.id},{new_sector.id} +""" + s3_stubber.add_response( + 'get_object', + { + 'Body': BytesIO(bytes(csv_content, encoding='utf-8')) + }, + expected_params={ + 'Bucket': bucket, + 'Key': object_key + } + ) + + call_command('update_investment_project_sector', bucket, object_key) + + investment_project.refresh_from_db() + + assert investment_project.sector == new_sector + versions = Version.objects.get_for_object(investment_project) + assert len(versions) == 1 + assert versions[0].revision.comment == 'Sector migration.' diff --git a/datahub/metadata/test/factories.py b/datahub/metadata/test/factories.py index 9245e4af4..f8cf4861f 100644 --- a/datahub/metadata/test/factories.py +++ b/datahub/metadata/test/factories.py @@ -36,3 +36,43 @@ class TeamFactory(factory.django.DjangoModelFactory): class Meta: model = 'metadata.Team' + + +class ReferralSourceActivityFactory(factory.django.DjangoModelFactory): + """ReferralSourceActivity factory.""" + + id = factory.LazyFunction(uuid.uuid4) + name = factory.Sequence(lambda n: f'name {n}') + + class Meta: + model = 'metadata.ReferralSourceActivity' + + +class ReferralSourceWebsiteFactory(factory.django.DjangoModelFactory): + """ReferralSourceWebsite factory.""" + + id = factory.LazyFunction(uuid.uuid4) + name = factory.Sequence(lambda n: f'name {n}') + + class Meta: + model = 'metadata.ReferralSourceWebsite' + + +class ReferralSourceMarketingFactory(factory.django.DjangoModelFactory): + """ReferralSourceMarketing factory.""" + + id = factory.LazyFunction(uuid.uuid4) + name = factory.Sequence(lambda n: f'name {n}') + + class Meta: + model = 'metadata.ReferralSourceMarketing' + + +class SectorFactory(factory.django.DjangoModelFactory): + """Sector factory.""" + + id = factory.LazyFunction(uuid.uuid4) + name = factory.Sequence(lambda n: f'name {n}') + + class Meta: + model = 'metadata.Sector'