From 6ca47933bd28c1a32df0fb8cd0b6e10d6043cf5a Mon Sep 17 00:00:00 2001 From: Alex Dusenbery Date: Fri, 15 Sep 2023 15:18:06 -0400 Subject: [PATCH] feat: add allocate_assignments method to content_assignments app. ENT-7531 --- .../apps/content_assignments/api.py | 141 +++++++++++++++++- .../apps/content_assignments/constants.py | 2 + .../apps/content_assignments/models.py | 48 ++++++ .../content_assignments/tests/test_api.py | 88 ++++++++++- .../apps/subsidy_access_policy/models.py | 8 + enterprise_access/settings/devstack.py | 2 + 6 files changed, 283 insertions(+), 6 deletions(-) diff --git a/enterprise_access/apps/content_assignments/api.py b/enterprise_access/apps/content_assignments/api.py index 42b5c478..f489c692 100644 --- a/enterprise_access/apps/content_assignments/api.py +++ b/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) diff --git a/enterprise_access/apps/content_assignments/constants.py b/enterprise_access/apps/content_assignments/constants.py index 828b9fc7..a04834c3 100644 --- a/enterprise_access/apps/content_assignments/constants.py +++ b/enterprise_access/apps/content_assignments/constants.py @@ -17,3 +17,5 @@ class LearnerContentAssignmentStateChoices: (CANCELLED, 'Cancelled'), (ERRORED, 'Errored'), ) + + REALLOCATE_STATES = (CANCELLED, ERRORED) diff --git a/enterprise_access/apps/content_assignments/models.py b/enterprise_access/apps/content_assignments/models.py index f64b4746..579cc5f6 100644 --- a/enterprise_access/apps/content_assignments/models.py +++ b/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, + ) diff --git a/enterprise_access/apps/content_assignments/tests/test_api.py b/enterprise_access/apps/content_assignments/tests/test_api.py index de3dea6e..23229fba 100644 --- a/enterprise_access/apps/content_assignments/tests/test_api.py +++ b/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='alice@foo.com', + content_key=content_key, + content_quantity=content_price_cents, + state=LearnerContentAssignmentStateChoices.ALLOCATED, + ) + accepted_assignment = LearnerContentAssignmentFactory.create( + assignment_configuration=self.assignment_configuration, + learner_email='bob@foo.com', + content_key=content_key, + content_quantity=content_price_cents, + state=LearnerContentAssignmentStateChoices.ACCEPTED, + ) + cancelled_assignment = LearnerContentAssignmentFactory.create( + assignment_configuration=self.assignment_configuration, + learner_email='carol@foo.com', + content_key=content_key, + content_quantity=200, + state=LearnerContentAssignmentStateChoices.CANCELLED, + ) + errored_assignment = LearnerContentAssignmentFactory.create( + assignment_configuration=self.assignment_configuration, + learner_email='david@foo.com', + 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 eugene@foo.com + 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, 'eugene@foo.com') + self.assertEqual(created_assignment.content_key, content_key) + self.assertEqual(created_assignment.content_quantity, content_price_cents) + self.assertEqual(created_assignment.state, LearnerContentAssignmentStateChoices.ALLOCATED) diff --git a/enterprise_access/apps/subsidy_access_policy/models.py b/enterprise_access/apps/subsidy_access_policy/models.py index 43d6e8f2..4c80184f 100644 --- a/enterprise_access/apps/subsidy_access_policy/models.py +++ b/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 diff --git a/enterprise_access/settings/devstack.py b/enterprise_access/settings/devstack.py index de881eea..f26b66ca 100644 --- a/enterprise_access/settings/devstack.py +++ b/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', ]