Merge pull request #5709 from uktrade/TET-830-add-omis-to-activity-st…

…ream

Add omis order to company activity

datahub/cleanup/management/commands/delete_old_records.py

@@ -122,9 +122,7 @@ class Command(BaseCleanupCommand):
                 DatetimeLessThanCleanupFilter('created_on', COMPANY_EXPIRY_PERIOD),
                 DatetimeLessThanCleanupFilter('modified_on', COMPANY_EXPIRY_PERIOD),
             ),
-            excluded_relations=(
-                CompanyReferral._meta.get_field('activity'),
-            ),
+            excluded_relations=(CompanyReferral._meta.get_field('activity'),),
         ),
         'interaction.Interaction': ModelCleanupConfig(
             (DatetimeLessThanCleanupFilter('date', INTERACTION_EXPIRY_PERIOD),),
@@ -250,6 +248,7 @@ class Command(BaseCleanupCommand):
             excluded_relations=(
                 Order._meta.get_field('assignees'),
                 Order._meta.get_field('subscribers'),
+                Order._meta.get_field('activity'),
             ),
         ),
         'company.CompanyExport': ModelCleanupConfig(

datahub/cleanup/test/commands/test_delete_old_records.py

@@ -34,8 +34,11 @@
     OneListCoreTeamMemberFactory,
     SubsidiaryFactory,
 )
-from datahub.company_activity.tests.factories import CompanyActivityInteractionFactory
-from datahub.company_activity.tests.factories import CompanyActivityInvestmentProjectFactory
+from datahub.company_activity.tests.factories import (
+    CompanyActivityInteractionFactory,
+    CompanyActivityInvestmentProjectFactory,
+    CompanyActivityOmisOrderFactory,
+)
 from datahub.company_referral.test.factories import (
     CompanyReferralFactory,
     CompleteCompanyReferralFactory,
@@ -212,6 +215,12 @@
                 'expired_objects_kwargs': [],
                 'unexpired_objects_kwargs': [],
             },
+            {
+                'factory': CompanyActivityOmisOrderFactory,
+                'field': 'company',
+                'expired_objects_kwargs': [],
+                'unexpired_objects_kwargs': [],
+            },
             {
                 'factory': EYBLeadFactory,
                 'field': 'company',
@@ -681,6 +690,7 @@
             'omis_payment.Payment',
             'omis_payment.Refund',
             'omis_payment.PaymentGatewaySession',
+            'company_activity.CompanyActivity',
         },
         'expired_objects_kwargs': [
             {

datahub/cleanup/test/commands/test_delete_orphaned_versions.py

@@ -110,6 +110,7 @@
     CompanyReferralFactory,
     CompanyInteractionFactory,
     InvestmentProjectFactory,
+    OrderFactory,
 ]
 
 

datahub/cleanup/test/commands/test_delete_orphans.py

@@ -94,7 +94,7 @@
         ),
         'implicit_related_models': (),
         'ignored_models': (
-            # Ignored as deleted with interactions, investments and referrals
+            # Ignored as deleted with interactions, investments, orders and referrals
             ('company_activity.CompanyActivity', 'company'),
         ),
     },

datahub/company/merge_company.py

@@ -163,7 +163,7 @@ def merge_companies(source_company: Company, target_company: Company, user):
         results[CompanyActivity] = {
             'company': results[Interaction]['company']
             + results[InvestmentProject]['investor_company']
-            + results[CompanyReferral]['company'],
+            + results[CompanyReferral]['company'] + results[Order]['company'],
         }
 
         source_company.mark_as_transferred(

datahub/company/test/test_merge_company.py

@@ -221,6 +221,7 @@ class TestDuplicateCompanyMerger:
                 company_with_orders_factory,
                 {
                     **base_expected_results,
+                    CompanyActivity: {'company': 3},
                     Contact: {'company': 3},
                     Order: {'company': 3},
 

datahub/company_activity/migrations/0007_companyactivity_order_and_more.py

@@ -0,0 +1,25 @@
+# Generated by Django 4.2.16 on 2024-10-14 08:22
+
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('order', '0014_alter_order_company_alter_order_contact_and_more'),
+        ('company_activity', '0006_alter_great_data_comment_and_more'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='companyactivity',
+            name='order',
+            field=models.ForeignKey(blank=True, help_text='If related to an omis order, must not have relations to any other activity (referral, event etc)', null=True, on_delete=django.db.models.deletion.CASCADE, related_name='activity', to='order.order', unique=True),
+        ),
+        migrations.AlterField(
+            model_name='companyactivity',
+            name='activity_source',
+            field=models.CharField(choices=[('interaction', 'interaction'), ('referral', 'referral'), ('event', 'event'), ('investment', 'investment'), ('order', 'order')], help_text='The type of company activity, such as an interaction, event, referral etc.', max_length=255),
+        ),
+    ]

datahub/company_activity/models/company_activity.py

@@ -28,6 +28,7 @@ class ActivitySource(models.TextChoices):
         referral = ('referral', 'referral')
         event = ('event', 'event')
         investment = ('investment', 'investment')
+        order = ('order', 'order')
 
     id = models.UUIDField(primary_key=True, default=uuid.uuid4)
     company = models.ForeignKey(
@@ -85,6 +86,18 @@ class ActivitySource(models.TextChoices):
             'investment project for company'
         ),
     )
+    order = models.ForeignKey(
+        'order.Order',
+        unique=True,
+        null=True,
+        blank=True,
+        related_name='activity',
+        on_delete=models.CASCADE,
+        help_text=(
+            'If related to an omis order, must not have relations to any other activity '
+            '(referral, event etc)'
+        ),
+    )
 
     def __str__(self):
         """Readable name for CompanyActivity"""

datahub/company_activity/tasks/sync.py

@@ -8,6 +8,7 @@
 from datahub.core.queues.scheduler import LONG_RUNNING_QUEUE
 from datahub.interaction.models import Interaction
 from datahub.investment.project.models import InvestmentProject
+from datahub.omis.order.models import Order
 
 
 logger = logging.getLogger(__name__)
@@ -189,3 +190,59 @@ def relate_company_activity_to_investment_projects(batch_size=500):
             f'Creating in batches of: {batch_size} CompanyActivities. {total} remaining.')
         CompanyActivity.objects.bulk_create(objs=batch, batch_size=batch_size)
         total -= batch_size
+
+
+def schedule_sync_order_to_company_activity():
+    """
+    Schedules a task to relate all `Omis Order`s to `CompanyActivity`s
+
+    Can be used to populate the CompanyActivity with missing orders
+    or to initially populate the model.
+    """
+    job = job_scheduler(
+        queue_name=LONG_RUNNING_QUEUE,
+        function=relate_company_activity_to_orders,
+        job_timeout=HALF_DAY_IN_SECONDS,
+        max_retries=5,
+        retry_backoff=True,
+    )
+    logger.info(
+        f'Task {job.id} schedule_sync_order_to_company_activity scheduled.',
+    )
+    return job
+
+
+def relate_company_activity_to_orders(batch_size=500):
+    """
+    Grabs all omis orders so they can be related to in the
+    `CompanyActivity` model with a bulk_create. Excludes any
+    order projects already associated in the CompanyActivity model.
+    """
+    activity_orders = set(
+        CompanyActivity.objects.filter(
+            order__isnull=False,
+        ).values_list('order_id', flat=True),
+    )
+
+    orders = Order.objects.filter(
+        company_id__isnull=False,
+    ).values('id', 'created_on', 'company_id')
+    objs = (
+        CompanyActivity(
+            order_id=order['id'],
+            date=order['created_on'],
+            company_id=order['company_id'],
+            activity_source=CompanyActivity.ActivitySource.order,
+        )
+        for order in orders
+        if order['id'] not in activity_orders
+    )
+    total = orders.count()
+    while True:
+        batch = list(islice(objs, batch_size))
+        if not batch:
+            logger.info('Finished bulk creating CompanyActivities.')
+            break
+        logger.info(f'Creating in batches of: {batch_size} CompanyActivities. {total} remaining.')
+        CompanyActivity.objects.bulk_create(objs=batch, batch_size=batch_size)
+        total -= batch_size

datahub/company_activity/tests/factories.py

@@ -7,6 +7,7 @@
 from datahub.interaction.test.factories import CompanyInteractionFactory
 from datahub.investment.project.test.factories import InvestmentProjectFactory
 from datahub.metadata.test.factories import CountryFactory
+from datahub.omis.order.test.factories import OrderFactory
 
 
 class CompanyActivityInteractionFactory(factory.django.DjangoModelFactory):
@@ -23,6 +24,7 @@ class CompanyActivityInteractionFactory(factory.django.DjangoModelFactory):
     interaction = factory.SubFactory(CompanyInteractionFactory)
     referral = None
     investment = None
+    order = None
 
     class Meta:
         model = 'company_activity.CompanyActivity'
@@ -49,6 +51,7 @@ class CompanyActivityReferralFactory(factory.django.DjangoModelFactory):
     referral = factory.SubFactory(CompanyReferralFactory)
     interaction = None
     investment = None
+    order = None
 
     class Meta:
         model = 'company_activity.CompanyActivity'
@@ -75,6 +78,7 @@ class CompanyActivityInvestmentProjectFactory(factory.django.DjangoModelFactory)
     investment = factory.SubFactory(InvestmentProjectFactory)
     interaction = None
     referral = None
+    order = None
 
     class Meta:
         model = 'company_activity.CompanyActivity'
@@ -90,6 +94,33 @@ def _create(cls, model_class, *args, **kwargs):
         return CompanyActivity.objects.get(investment_id=obj.investment_id)
 
 
+class CompanyActivityOmisOrderFactory(factory.django.DjangoModelFactory):
+    """
+    CompanyActivity factory with an omis order.
+    """
+
+    date = now()
+    activity_source = CompanyActivity.ActivitySource.order
+    company = factory.SubFactory(CompanyFactory)
+    investment = None
+    interaction = None
+    referral = None
+    order = factory.SubFactory(OrderFactory)
+
+    class Meta:
+        model = 'company_activity.CompanyActivity'
+
+    @classmethod
+    def _create(cls, model_class, *args, **kwargs):
+        """
+        Overwrite the _create to prevent the CompanyActivity from saving to the database.
+        This is due to the Omis Order already creating the CompanyActivity inside its
+        model save.
+        """
+        obj = model_class(*args, **kwargs)
+        return CompanyActivity.objects.get(order_id=obj.order_id)
+
+
 class CompanyActivityIngestedFileFactory(factory.django.DjangoModelFactory):
     """
     CompanyActivity ingested file factory

datahub/company_activity/tests/test_tasks/test_order_task.py

@@ -0,0 +1,76 @@
+from unittest import mock
+
+import pytest
+
+from datahub.company_activity.models import CompanyActivity
+from datahub.company_activity.tasks.sync import (
+    relate_company_activity_to_orders,
+    schedule_sync_order_to_company_activity,
+)
+from datahub.omis.order.test.factories import OrderFactory
+
+
+@pytest.mark.django_db
+class TestCompanyActivityOrderTasks:
+    """
+    Tests for the schedule_sync_investments_to_company_activity task.
+    """
+
+    def test_orders_are_copied_to_company_activity(self):
+        """
+        Test that investments are added to the CompanyActivity model.
+        """
+        orders = OrderFactory.create_batch(5)
+
+        # Remove the created CompanyActivities added by the omis order `save` method
+        # to mimic already existing data in staging and prod database.
+        CompanyActivity.objects.all().delete()
+        assert CompanyActivity.objects.count() == 0
+
+        # Check the "existing" orders are added to the company activity model
+        schedule_sync_order_to_company_activity()
+        assert CompanyActivity.objects.count() == len(orders)
+
+        company_activity = CompanyActivity.objects.get(order_id=orders[0].id)
+        assert company_activity.date == orders[0].created_on
+        assert company_activity.activity_source == CompanyActivity.ActivitySource.order
+        assert company_activity.company_id == orders[0].company.id
+
+    @mock.patch('datahub.company_activity.models.CompanyActivity.objects.bulk_create')
+    def test_order_are_bulk_created_in_batches(self, mocked_bulk_create, caplog):
+        """
+        Test that omis orders are bulk created in batches.
+        """
+        caplog.set_level('INFO')
+        batch_size = 5
+
+        OrderFactory.create_batch(10)
+
+        # Delete any activity created through the investments save method.
+        CompanyActivity.objects.all().delete()
+        assert CompanyActivity.objects.count() == 0
+
+        # Ensure omis orders are bulk_created
+        relate_company_activity_to_orders(batch_size)
+        assert mocked_bulk_create.call_count == 2
+
+        assert (
+            f'Creating in batches of: {batch_size} CompanyActivities. 10 remaining.' in caplog.text
+        )
+        assert (
+            f'Creating in batches of: {batch_size} CompanyActivities. 5 remaining.' in caplog.text
+        )
+        assert 'Finished bulk creating CompanyActivities.' in caplog.text
+
+    def test_order_with_a_company_activity_are_not_added_again(self):
+        """
+        Test that investment projects which are already part of the `CompanyActivity` model
+        are not added again.
+        """
+        OrderFactory.create_batch(4)
+
+        assert CompanyActivity.objects.count() == 4
+
+        # Check count remains unchanged.
+        schedule_sync_order_to_company_activity()
+        assert CompanyActivity.objects.count() == 4

datahub/interaction/test/test_models.py

@@ -61,7 +61,7 @@ def test_save(self):
         Test save also saves to the `CompanyActivity` model.
         Test save does not save to the `CompanyActivity` model if it already exists.
         """
-        assert CompanyActivity.objects.all().count() == 0
+        assert not CompanyActivity.objects.all().exists()
         interaction = CompanyInteractionFactory()
         assert CompanyActivity.objects.all().count() == 1
 

datahub/omis/order/models.py

@@ -11,6 +11,7 @@
 from mptt.fields import TreeForeignKey
 
 from datahub.company.models import Advisor, Company, Contact
+from datahub.company_activity.models import CompanyActivity
 from datahub.core import reversion
 from datahub.core.models import (
     BaseConstantModel,
@@ -355,7 +356,19 @@ def save(self, *args, **kwargs):
             self.reference = self.generate_reference()
         if not self.public_token:
             self.public_token = self.generate_public_token()
-        return super().save(*args, **kwargs)
+
+        with transaction.atomic():
+            super().save(*args, **kwargs)
+            if not self.company_id:
+                return
+            CompanyActivity.objects.update_or_create(
+                order_id=self.id,
+                activity_source=CompanyActivity.ActivitySource.order,
+                defaults={
+                    'date': self.created_on,
+                    'company_id': self.company_id,
+                },
+            )
 
     def get_lead_assignee(self):
         """