fix: update parent_content_key and is_assigned_course_run on assignme…

…nt re-allocation

enterprise_access/apps/content_assignments/api.py

@@ -55,6 +55,7 @@
     'lms_user_id', 'learner_email', 'allocation_batch_id',
     'content_quantity', 'state', 'preferred_course_run_key',
     'allocated_at', 'cancelled_at', 'expired_at', 'errored_at',
+    'parent_content_key', 'is_assigned_course_run',
 ]
 
 
@@ -329,9 +330,16 @@ def allocate_assignments(assignment_configuration, learner_emails, content_key,
     existing_assignments_needs_update = set()
 
     # This step to find and update the preferred_course_run_key is required in order
-    # for nudge emails to target the start date of the new run.
+    # for nudge emails to target the start date of the new run. For run-based assignments,
+    # the preferred_course_run_key is the same as the assignment's content_key.
     preferred_course_run_key = _get_preferred_course_run_key(assignment_configuration, content_key)
 
+    # Determine if the assignment's content_key is a course run or a course key based
+    # on an associated parent content key. If the parent content key is None, then the
+    # assignment is for a course; otherwise, it's an assignment for a course run.
+    parent_content_key = _get_parent_content_key(assignment_configuration, content_key)
+    is_assigned_course_run = bool(parent_content_key)
+
     # Split up the existing assignment records by state
     for assignment in existing_assignments:
         if not assignment.lms_user_id:
@@ -349,14 +357,28 @@ def allocate_assignments(assignment_configuration, learner_emails, content_key,
                 existing_assignments_needs_update.add(assignment)
 
         if assignment.state in LearnerContentAssignmentStateChoices.REALLOCATE_STATES:
-            _reallocate_assignment(assignment, content_quantity, allocation_batch_id, preferred_course_run_key)
+            _reallocate_assignment(
+                assignment,
+                content_quantity,
+                allocation_batch_id,
+                preferred_course_run_key,
+                parent_content_key,
+                is_assigned_course_run,
+            )
             existing_assignments_needs_update.add(assignment)
         elif assignment.state == LearnerContentAssignmentStateChoices.ALLOCATED:
             # For some already-allocated assignments being re-assigned, we might still need to update the preferred
             # course run for nudge email purposes.
             if assignment.preferred_course_run_key != preferred_course_run_key:
                 assignment.preferred_course_run_key = preferred_course_run_key
                 existing_assignments_needs_update.add(assignment)
+            # Update the parent_content_key and is_assigned_course_run fields if they have changed.
+            if assignment.parent_content_key != parent_content_key:
+                assignment.parent_content_key = parent_content_key
+                existing_assignments_needs_update.add(assignment)
+            if assignment.is_assigned_course_run != is_assigned_course_run:
+                assignment.is_assigned_course_run = is_assigned_course_run
+                existing_assignments_needs_update.add(assignment)
 
         learner_emails_with_existing_assignments.add(assignment.learner_email.lower())
 
@@ -505,7 +527,14 @@ def _get_existing_assignments_for_allocation(
     return existing_assignments
 
 
-def _reallocate_assignment(assignment, content_quantity, allocation_batch_id, preferred_course_run_key):
+def _reallocate_assignment(
+        assignment,
+        content_quantity,
+        allocation_batch_id,
+        preferred_course_run_key,
+        parent_content_key,
+        is_assigned_course_run,
+    ):
     """
     Modifies a ``LearnerContentAssignment`` record during the allocation flow.  The record
     is **not** saved.
@@ -519,6 +548,8 @@ def _reallocate_assignment(assignment, content_quantity, allocation_batch_id, pr
     assignment.expired_at = None
     assignment.errored_at = None
     assignment.preferred_course_run_key = preferred_course_run_key
+    assignment.parent_content_key = parent_content_key
+    assignment.is_assigned_course_run = is_assigned_course_run
     # Prevent invalid data from entering the database by calling the low-level full_clean() function manually.
     assignment.full_clean()
     return assignment
@@ -571,6 +602,8 @@ def _get_parent_content_key(assignment_configuration, content_key):
     course_content_metadata = _get_content_summary(assignment_configuration, content_key)
     metadata_content_key = course_content_metadata.get('content_key')
 
+    print('_get_parent_content_key?!: ', metadata_content_key, content_key)
+
     # Check if the assignment's content_key matches the returned content_key. If so, this is a course key
     # which has no parent key.
     if content_key == metadata_content_key:

enterprise_access/apps/content_assignments/tests/factories.py

@@ -38,5 +38,6 @@ class Meta:
     lms_user_id = factory.LazyAttribute(lambda _: FAKER.pyint())
     content_key = factory.LazyAttribute(lambda _: random_content_key())
     parent_content_key = factory.LazyAttribute(lambda _: random_parent_content_key())
+    is_assigned_course_run = factory.LazyAttribute(lambda _: FAKER.pybool())
     content_title = factory.LazyAttribute(lambda _: f'{FAKER.word()}: a master class')
     content_quantity = factory.LazyAttribute(lambda _: FAKER.pyfloat(positive=False, right_digits=0))

enterprise_access/apps/content_assignments/tests/test_api.py

@@ -338,17 +338,36 @@ def test_allocate_assignments_negative_quantity(self):
         'enterprise_access.apps.content_assignments.api.get_and_cache_content_metadata',
         return_value=mock.MagicMock(),
     )
+    @ddt.data(
+        {'is_assigned_course_run': False},
+        # {'is_assigned_course_run': True},
+    )
+    @ddt.unpack
     def test_allocate_assignments_happy_path(
-        self, mock_get_and_cache_content_metadata, mock_pending_learner_task, mock_new_assignment_email_task
+        self,
+        mock_get_and_cache_content_metadata,
+        mock_pending_learner_task,
+        mock_new_assignment_email_task,
+        is_assigned_course_run,
     ):
         """
         Tests the allocation of new assignments against a given configuration.
         """
-        content_key = 'demoX'
-        course_run_key_old = 'demoX+1T2022'
-        course_run_key = 'demoX+2T2023'
+        course_key = 'edX+DemoX'
         content_title = 'edx: Demo 101'
         content_price_cents = 100
+
+        # Course-level assignments
+        content_key = course_key
+        parent_content_key = None
+        course_run_key_old = 'course-v1:edX+DemoX+1T2022'
+        course_run_key = 'course-v1:edX+DemoX+2T2023'
+
+        # Course run-level assignments
+        if is_assigned_course_run:
+            content_key = course_run_key
+            parent_content_key = course_key
+
         # add a duplicate email to the input list to ensure only one
         # allocation occurs.
         # We throw a couple ALL UPPER CASE emails into the requested emails to allocate
@@ -365,18 +384,25 @@ def test_allocate_assignments_happy_path(
                 'eugene',
                 'faythe',
                 'erin',
+                'mary',
                 'grace',
+                'xavier',
+                'steve',
+                'kian',
             )
         ]
         mock_get_and_cache_content_metadata.return_value = {
             'content_title': content_title,
+            'content_key': course_key,
             'course_run_key': course_run_key,
         }
 
         default_factory_options = {
             'assignment_configuration': self.assignment_configuration,
             'content_key': content_key,
-            'preferred_course_run_key': course_run_key,
+            'parent_content_key': parent_content_key,
+            'is_assigned_course_run': is_assigned_course_run,
+            'preferred_course_run_key': content_key if is_assigned_course_run else course_run_key,
             'content_title': content_title,
             'content_quantity': -content_price_cents,
         }
@@ -395,6 +421,22 @@ def test_allocate_assignments_happy_path(
             # outdated run should trigger update.
             'preferred_course_run_key': course_run_key_old,
         })
+        # Allocated assignment SHOULD be updated because it had an outdated parent_content_key.
+        allocated_assignment_old_parent_content_key = LearnerContentAssignmentFactory.create(**{
+            **default_factory_options,
+            'learner_email': 'mary@foo.com',
+            'state': LearnerContentAssignmentStateChoices.ALLOCATED,
+            # outdated parent_content_key should trigger update.
+            'parent_content_key': None,
+        })
+        # Allocated assignment SHOULD be updated because it had an outdated is_assigned_course_run.
+        allocated_assignment_old_run_based = LearnerContentAssignmentFactory.create(**{
+            **default_factory_options,
+            'learner_email': 'xavier@foo.com',
+            'state': LearnerContentAssignmentStateChoices.ALLOCATED,
+            # outdated is_assigned_course_run should trigger update.
+            'is_assigned_course_run': False,
+        })
         # Allocated assignment SHOULD be updated because it had a NULL run.
         allocated_assignment_null_run = LearnerContentAssignmentFactory.create(**{
             **default_factory_options,
@@ -423,13 +465,41 @@ def test_allocate_assignments_happy_path(
             # outdated run should trigger update.
             'preferred_course_run_key': course_run_key_old,
         })
+        # Cancelled assignment should be updated/re-allocated (including recieving new parent_content_key).
+        cancelled_assignment_old_parent_content_key = LearnerContentAssignmentFactory.create(**{
+            **default_factory_options,
+            'learner_email': 'steve@foo.com',
+            'state': LearnerContentAssignmentStateChoices.CANCELLED,
+            # outdated parent_content_key should trigger update.
+            'parent_content_key': None,
+        })
+        # Cancelled assignment should be updated/re-allocated (including recieving new is_assigned_course_run).
+        cancelled_assignment_old_run_based = LearnerContentAssignmentFactory.create(**{
+            **default_factory_options,
+            'learner_email': 'kian@foo.com',
+            'state': LearnerContentAssignmentStateChoices.CANCELLED,
+            # outdated is_assigned_course_run should trigger update.
+            'is_assigned_course_run': False,
+        })
         # Errored assignment should be updated/re-allocated.
         errored_assignment = LearnerContentAssignmentFactory.create(**{
             **default_factory_options,
             'learner_email': 'david@foo.com',
             'state': LearnerContentAssignmentStateChoices.ERRORED,
         })
 
+        for assignment in LearnerContentAssignment.objects.all():
+            print(
+                'original assignment:',
+                str(assignment.uuid),
+                assignment.learner_email,
+                assignment.state,
+                assignment.content_key,
+                assignment.parent_content_key,
+                assignment.is_assigned_course_run,
+                assignment.preferred_course_run_key,
+            )
+
         allocation_results = allocate_assignments(
             self.assignment_configuration,
             learners_to_assign,
@@ -438,26 +508,69 @@ def test_allocate_assignments_happy_path(
         )
 
         # Refresh from db to get any updates reflected in the python objects.
-        for record in (
+        assignments_to_refresh = (
             allocated_assignment,
             allocated_assignment_old_run,
             allocated_assignment_null_run,
             accepted_assignment,
             cancelled_assignment,
             cancelled_assignment_old_run,
+            cancelled_assignment_old_parent_content_key,
+            cancelled_assignment_old_run_based,
             errored_assignment,
-        ):
+        )
+        if is_assigned_course_run:
+            assignments_to_refresh += (
+                allocated_assignment_old_parent_content_key,
+                allocated_assignment_old_run_based,
+            )
+
+        for record in assignments_to_refresh:
             record.refresh_from_db()
 
-        # The allocated assignments with outdated runs, errored assignments, and cancelled assignments should be the
-        # only things updated.
+        for assignment in LearnerContentAssignment.objects.all():
+            print(
+                'assignment after allocation:',
+                str(assignment.uuid),
+                assignment.learner_email,
+                assignment.state,
+                assignment.content_key,
+                assignment.parent_content_key,
+                assignment.is_assigned_course_run,
+                assignment.preferred_course_run_key,
+            )
+
+
+        # Create list of assignments expected to be updated, including:
+        #   - The allocated assignments with outdated: preferred course run, parent_content_key, is_assigned_course_run.
+        #   - Errored assignments
+        #   - Cancelled assignments
         expected_updated_assignments = (
             allocated_assignment_old_run,
             allocated_assignment_null_run,
             cancelled_assignment,
             cancelled_assignment_old_run,
+            cancelled_assignment_old_parent_content_key,
+            cancelled_assignment_old_run_based,
             errored_assignment,
         )
+
+        if is_assigned_course_run:
+            expected_updated_assignments += (
+                allocated_assignment_old_parent_content_key,
+                allocated_assignment_old_run_based,
+            )
+
+        print('actual updated uuids:', {str(record.uuid) for record in allocation_results['updated']})
+
+        if is_assigned_course_run:
+            expected_updated_assignments += (
+                allocated_assignment_old_parent_content_key,
+                allocated_assignment_old_run_based,
+            )
+
+        print('expected updated uuids:', {str(record.uuid) for record in expected_updated_assignments})
+
         self.assertEqual(
             {record.uuid for record in allocation_results['updated']},
             {assignment.uuid for assignment in expected_updated_assignments},
@@ -470,6 +583,11 @@ def test_allocate_assignments_happy_path(
             allocated_assignment,
             accepted_assignment,
         )
+        if not is_assigned_course_run:
+            expected_no_change_assignments += (
+                allocated_assignment_old_parent_content_key,
+                allocated_assignment_old_run_based,
+            )
         self.assertEqual(
             {record.uuid for record in allocation_results['no_change']},
             {assignment.uuid for assignment in expected_no_change_assignments},
@@ -480,10 +598,18 @@ def test_allocate_assignments_happy_path(
         # The existing assignments should be 'allocated' now, except for the already accepted one
         self.assertEqual(cancelled_assignment.state, LearnerContentAssignmentStateChoices.ALLOCATED)
         self.assertEqual(cancelled_assignment_old_run.state, LearnerContentAssignmentStateChoices.ALLOCATED)
+        self.assertEqual(
+            cancelled_assignment_old_parent_content_key.state, LearnerContentAssignmentStateChoices.ALLOCATED
+        )
+        self.assertEqual(cancelled_assignment_old_run_based.state, LearnerContentAssignmentStateChoices.ALLOCATED)
         self.assertEqual(errored_assignment.state, LearnerContentAssignmentStateChoices.ALLOCATED)
         self.assertEqual(allocated_assignment.state, LearnerContentAssignmentStateChoices.ALLOCATED)
         self.assertEqual(allocated_assignment_old_run.state, LearnerContentAssignmentStateChoices.ALLOCATED)
         self.assertEqual(allocated_assignment_null_run.state, LearnerContentAssignmentStateChoices.ALLOCATED)
+        self.assertEqual(
+            allocated_assignment_old_parent_content_key.state, LearnerContentAssignmentStateChoices.ALLOCATED
+        )
+        self.assertEqual(allocated_assignment_old_run_based.state, LearnerContentAssignmentStateChoices.ALLOCATED)
         self.assertEqual(accepted_assignment.state, LearnerContentAssignmentStateChoices.ACCEPTED)
 
         # We should have created only one new, allocated assignment for eugene@foo.com