feat: add allocate_assignments method to content_assignments app.

ENT-7531

enterprise_access/apps/content_assignments/api.py

@@ -21,28 +21,161 @@ def get_assignment_configuration(uuid):
 
 def get_assignments_for_configuration(
     assignment_configuration,
-    state=LearnerContentAssignmentStateChoices.ALLOCATED,
+    **additional_filters,
 ):
     """
     Returns a queryset of all ``LearnerContentAssignment`` records
-    for the given assignment configuration.
+    for the given assignment configuration, optionally filtered
+    further by the provided ``additional_filters``.
     """
     queryset = LearnerContentAssignment.objects.select_related(
         'assignment_configuration',
     ).filter(
         assignment_configuration=assignment_configuration,
-        state=state,
+        **additional_filters,
     )
     return queryset
 
 
+def get_assignments_by_learner_email_and_content(
+    assignment_configuration,
+    learner_emails,
+    content_key,
+):
+    """
+    Returns a queryset of all ``LearnerContentAssignment`` records
+    in the given assignment configuration for the provided list
+    of learner_emails and the given content_key.
+    """
+    return get_assignments_for_configuration(
+        assignment_configuration,
+        learner_email__in=learner_emails,
+        content_key=content_key,
+    )
+
+
 def get_allocated_quantity_for_configuration(assignment_configuration):
     """
     Returns a float representing the total quantity, in USD cents, currently allocated
     via Assignments for the given configuration.
     """
-    assignments_queryset = get_assignments_for_configuration(assignment_configuration)
+    assignments_queryset = get_assignments_for_configuration(
+        assignment_configuration,
+        state=LearnerContentAssignmentStateChoices.ALLOCATED,
+    )
     aggregate = assignments_queryset.aggregate(
         total_quantity=Sum('content_quantity'),
     )
     return aggregate['total_quantity']
+
+
+def allocate_assignments(assignment_configuration, learner_emails, content_key, content_price_cents):
+    """
+    Creates or updates an allocated assignment record
+    for the given ``content_key`` in the given ``assignment_configuration``,
+    for each email in the list of ``learner_emails``.
+
+    For existing assignment records with a (config, learner, content) combination, this function
+    does the following:
+      * If the existing record is cancelled or errored, update the existing record state to allocated
+      * If the existing record is allocated or accepted, don't do anything with the record
+
+    Params:
+      - ``assignment_configuration``: The AssignmentConfiguration record under which assignments should be allocated.
+      - ``learner_emails``: A list of learner email addresses to whom assignments should be allocated.
+      - ``content_key``: Typically a *course* key to which the learner is assigned.
+      - ``content_price_cents``: The cost of redeeming the content, in USD cents, at the time of allocation.
+
+    Returns: A dictionary of updated, created, and unchanged assignment records. e.g.
+      ```
+      {
+        'updated': [Updated LearnerContentAssignment records],
+        'created': [Newly-created LearnerContentAssignment records],
+        'no-change': [LearnerContentAssignment records that matched
+                      the provided (config, learner, content) combination,
+                      but were already in an 'allocated' or 'accepted' state],
+      }
+      ```
+
+    """
+    # Fetch any existing assignments for all pairs of (learner, content) in this assignment config.
+    existing_assignments = get_assignments_by_learner_email_and_content(
+        assignment_configuration,
+        learner_emails,
+        content_key,
+    )
+
+    # Existing Assignments in consideration by state
+    already_allocated_or_accepted = []
+    cancelled_or_errored_to_update = []
+
+    # Maintain a set of emails with existing records - we know we don't have to create
+    # new assignments for these.
+    learner_emails_with_existing_assignments = set()
+
+    # Split up the existing assignment records by state
+    for assignment in existing_assignments:
+        learner_emails_with_existing_assignments.add(assignment.learner_email)
+        if assignment.state in LearnerContentAssignmentStateChoices.REALLOCATE_STATES:
+            assignment.content_quantity = content_price_cents
+            assignment.state = LearnerContentAssignmentStateChoices.ALLOCATED
+            cancelled_or_errored_to_update.append(assignment)
+        else:
+            already_allocated_or_accepted.append(assignment)
+
+    # Bulk update and get a list of refreshed objects
+    updated_assignments = _update_and_refresh_assignments(
+        cancelled_or_errored_to_update,
+        ['content_quantity', 'state']
+    )
+
+    # Narrow down creation list of learner emails
+    learner_emails_for_assignment_creation = set(learner_emails) - learner_emails_with_existing_assignments
+
+    # Initialize and save LearnerContentAssignment instances for each of them
+    created_assignments = _create_new_assignments(
+        assignment_configuration,
+        learner_emails_for_assignment_creation,
+        content_key,
+        content_price_cents,
+    )
+
+    # Return a mapping of the action we took to lists of relevant assignment records.
+    return {
+        'updated': updated_assignments,
+        'created': created_assignments,
+        'no_change': already_allocated_or_accepted,
+    }
+
+
+def _update_and_refresh_assignments(assignment_records, fields_changed):
+    """
+    Helper to bulk save the given assignment_records
+    and refresh their state from the DB.
+    """
+    # Save the assignments to update
+    LearnerContentAssignment.bulk_update(assignment_records, fields_changed)
+
+    # Get a list of refreshed objects that we just updated
+    return LearnerContentAssignment.objects.filter(
+        uuid__in=[record.uuid for record in assignment_records],
+    )
+
+
+def _create_new_assignments(assignment_configuration, learner_emails, content_key, content_price_cents):
+    """
+    Helper to bulk save new LearnerContentAssignment instances.
+    """
+    assignments_to_create = [
+        LearnerContentAssignment(
+            assignment_configuration=assignment_configuration,
+            learner_email=learner_email,
+            content_key=content_key,
+            content_quantity=content_price_cents,
+            state=LearnerContentAssignmentStateChoices.ALLOCATED,
+        )
+        for learner_email in learner_emails
+    ]
+
+    # Do the bulk creation to save these records
+    return LearnerContentAssignment.bulk_create(assignments_to_create)

enterprise_access/apps/content_assignments/constants.py

@@ -17,3 +17,5 @@ class LearnerContentAssignmentStateChoices:
         (CANCELLED, 'Cancelled'),
         (ERRORED, 'Errored'),
     )
+
+    REALLOCATE_STATES = (CANCELLED, ERRORED)

enterprise_access/apps/content_assignments/models.py

@@ -5,11 +5,15 @@
 
 from django.core.exceptions import ObjectDoesNotExist
 from django.db import models
+from django.utils import timezone
 from django_extensions.db.models import TimeStampedModel
 from simple_history.models import HistoricalRecords
+from simple_history.utils import bulk_create_with_history, bulk_update_with_history
 
 from .constants import LearnerContentAssignmentStateChoices
 
+BULK_OPERATION_BATCH_SIZE = 50
+
 
 class AssignmentConfiguration(TimeStampedModel):
     """
@@ -46,6 +50,9 @@ class AssignmentConfiguration(TimeStampedModel):
 
     history = HistoricalRecords()
 
+    def __str__(self):
+        return f'uuid={self.uuid}, customer={self.enterprise_customer_uuid}'
+
     def delete(self, *args, **kwargs):
         """
         Perform a soft-delete, overriding the standard delete() method to prevent hard-deletes.
@@ -160,3 +167,44 @@ class Meta:
         ),
     )
     history = HistoricalRecords()
+
+    def __str__(self):
+        return (
+            f'uuid={self.uuid}, state={self.state}, learner_email={self.learner_email}, content_key={self.content_key}'
+        )
+
+    @classmethod
+    def bulk_create(cls, assignment_records):
+        """
+        Creates new ``LearnerContentAssignment`` records in bulk,
+        while saving their history:
+        https://django-simple-history.readthedocs.io/en/latest/common_issues.html#bulk-creating-a-model-with-history
+        """
+        return bulk_create_with_history(
+            assignment_records,
+            cls,
+            batch_size=BULK_OPERATION_BATCH_SIZE,
+        )
+
+    @classmethod
+    def bulk_update(cls, assignment_records, updated_field_names):
+        """
+        Updates and saves the given ``assignment_records`` in bulk,
+        while saving their history:
+        https://django-simple-history.readthedocs.io/en/latest/common_issues.html#bulk-updating-a-model-with-history-new
+
+        Note that the simple-history utility function uses Django's bulk_update() under the hood:
+        https://docs.djangoproject.com/en/3.2/ref/models/querysets/#bulk-update
+
+        which does *not* call save(), so we have to manually update the `modified` field
+        during this bulk operation in order for that field's value to be updated.
+        """
+        for record in assignment_records:
+            record.modified = timezone.now()
+
+        return bulk_update_with_history(
+            assignment_records,
+            cls,
+            updated_field_names + ['modified'],
+            batch_size=BULK_OPERATION_BATCH_SIZE,
+        )

enterprise_access/apps/content_assignments/tests/test_api.py

@@ -3,7 +3,7 @@
 """
 from django.test import TestCase
 
-from ..api import get_allocated_quantity_for_configuration, get_assignments_for_configuration
+from ..api import allocate_assignments, get_allocated_quantity_for_configuration, get_assignments_for_configuration
 from ..constants import LearnerContentAssignmentStateChoices
 from ..models import AssignmentConfiguration
 from .factories import LearnerContentAssignmentFactory
@@ -65,7 +65,10 @@ def test_get_assignments_for_configuration_different_states(self):
         ):
             with self.assertNumQueries(1):
                 actual_assignments = list(
-                    get_assignments_for_configuration(self.assignment_configuration, filter_state)
+                    get_assignments_for_configuration(
+                        self.assignment_configuration,
+                        state=filter_state
+                    )
                 )
 
             self.assertEqual(
@@ -87,3 +90,84 @@ def test_get_allocated_quantity_for_configuration(self):
         with self.assertNumQueries(1):
             actual_amount = get_allocated_quantity_for_configuration(self.assignment_configuration)
             self.assertEqual(actual_amount, 6000)
+
+    def test_allocate_assignments_happy_path(self):
+        """
+        Tests the allocation of new assignments against a given configuration.
+        """
+        content_key = 'demoX'
+        content_price_cents = 100
+        learners_to_assign = [
+            f'{name}@foo.com' for name in ('alice', 'bob', 'carol', 'david', 'eugene')
+        ]
+
+        allocated_assignment = LearnerContentAssignmentFactory.create(
+            assignment_configuration=self.assignment_configuration,
+            learner_email='[email protected]',
+            content_key=content_key,
+            content_quantity=content_price_cents,
+            state=LearnerContentAssignmentStateChoices.ALLOCATED,
+        )
+        accepted_assignment = LearnerContentAssignmentFactory.create(
+            assignment_configuration=self.assignment_configuration,
+            learner_email='[email protected]',
+            content_key=content_key,
+            content_quantity=content_price_cents,
+            state=LearnerContentAssignmentStateChoices.ACCEPTED,
+        )
+        cancelled_assignment = LearnerContentAssignmentFactory.create(
+            assignment_configuration=self.assignment_configuration,
+            learner_email='[email protected]',
+            content_key=content_key,
+            content_quantity=200,
+            state=LearnerContentAssignmentStateChoices.CANCELLED,
+        )
+        errored_assignment = LearnerContentAssignmentFactory.create(
+            assignment_configuration=self.assignment_configuration,
+            learner_email='[email protected]',
+            content_key=content_key,
+            content_quantity=200,
+            state=LearnerContentAssignmentStateChoices.ERRORED,
+        )
+
+        allocation_results = allocate_assignments(
+            self.assignment_configuration,
+            learners_to_assign,
+            content_key,
+            content_price_cents,
+        )
+
+        # Refresh from db to get any updates reflected in the python objects.
+        for record in (allocated_assignment, accepted_assignment, cancelled_assignment, errored_assignment):
+            record.refresh_from_db()
+
+        # The errored and cancelled assignments should be the only things updated
+        self.assertEqual(
+            {record.uuid for record in allocation_results['updated']},
+            {cancelled_assignment.uuid, errored_assignment.uuid},
+        )
+        for record in (cancelled_assignment, errored_assignment):
+            self.assertEqual(len(record.history.all()), 2)
+
+        # The allocated and accepted assignments should be the only things with no change
+        self.assertEqual(
+            {record.uuid for record in allocation_results['no_change']},
+            {allocated_assignment.uuid, accepted_assignment.uuid},
+        )
+        for record in (allocated_assignment, accepted_assignment):
+            self.assertEqual(len(record.history.all()), 1)
+
+        # The existing assignments should be 'allocated' now, except for the already accepted one
+        self.assertEqual(cancelled_assignment.state, LearnerContentAssignmentStateChoices.ALLOCATED)
+        self.assertEqual(errored_assignment.state, LearnerContentAssignmentStateChoices.ALLOCATED)
+        self.assertEqual(allocated_assignment.state, LearnerContentAssignmentStateChoices.ALLOCATED)
+        self.assertEqual(accepted_assignment.state, LearnerContentAssignmentStateChoices.ACCEPTED)
+
+        # We should have created only one new, allocated assignment for [email protected]
+        self.assertEqual(len(allocation_results['created']), 1)
+        created_assignment = allocation_results['created'][0]
+        self.assertEqual(created_assignment.assignment_configuration, self.assignment_configuration)
+        self.assertEqual(created_assignment.learner_email, '[email protected]')
+        self.assertEqual(created_assignment.content_key, content_key)
+        self.assertEqual(created_assignment.content_quantity, content_price_cents)
+        self.assertEqual(created_assignment.state, LearnerContentAssignmentStateChoices.ALLOCATED)

enterprise_access/apps/subsidy_access_policy/models.py

@@ -905,3 +905,11 @@ def can_allocate(self, number_of_learners, content_key, content_price_cents):
             return (False, REASON_POLICY_SPEND_LIMIT_REACHED)
 
         return (True, None)
+
+    def allocate(self, learner_emails, content_key, content_price_cents):
+        """
+        Creates allocated ``LearnerContentAssignment`` records.
+        """
+        # this will eventually lean on assignments_api.allocate_assignments()
+        # to do the heavy lifting.
+        raise NotImplementedError

enterprise_access/settings/devstack.py

@@ -93,6 +93,8 @@
 # shell_plus
 SHELL_PLUS_IMPORTS = [
     'from enterprise_access.apps.subsidy_request.utils import localized_utcnow',
+    'from enterprise_access.apps.content_assignments import api as assignments_api',
+    'from pprint import pprint',
 ]