feat: handle draft seats and entitlements in ecommerce data loader

course_discovery/apps/course_metadata/data_loaders/api.py

@@ -508,13 +508,19 @@ def _process_enrollment_codes(self, response):
             self.enrollment_skus.append(self.update_enrollment_code(body))
 
     def _delete_entitlements(self):
-        entitlements_to_delete = CourseEntitlement.objects.filter(
+        # The last two excludes ensure that we do not delete entitlements
+        # that have no official variants yet.
+        entitlements_to_delete = CourseEntitlement.everything.filter(
             partner=self.partner
-        ).exclude(sku__in=self.entitlement_skus)
+        ).exclude(sku__in=self.entitlement_skus).exclude(sku__isnull=True).exclude(sku__exact='')
 
         for entitlement in entitlements_to_delete:
-            msg = 'Deleting entitlement for course {course_title} with sku {sku} for partner {partner}'.format(
-                course_title=entitlement.course.title, sku=entitlement.sku, partner=entitlement.partner
+            # pylint: disable=line-too-long
+            msg = 'Deleting {draft_status} entitlement for course {course_title} with sku {sku} for partner {partner}'.format(
+                course_title=entitlement.course.title,
+                sku=entitlement.sku,
+                partner=entitlement.partner,
+                draft_status='draft' if entitlement.draft else 'non-draft'
             )
             logger.info(msg)
         entitlements_to_delete.delete()
@@ -556,6 +562,10 @@ def update_seats(self, body):
             )
         seats_to_remove.delete()
 
+        if course_run.draft_version:
+            draft_seats_to_remove = course_run.draft_version.seats.exclude(type__slug__in=certificate_types)
+            draft_seats_to_remove.delete()
+
     def update_seat(self, course_run, product_body):
         stock_record = product_body['stockrecords'][0]
         currency_code = stock_record['price_currency']
@@ -601,13 +611,24 @@ def update_seat(self, course_run, product_body):
             'credit_hours': credit_hours,
         }
 
-        _, created = course_run.seats.update_or_create(
+        seat, created = course_run.seats.update_or_create(
             type=seat_type,
             credit_provider=credit_provider,
             currency=currency,
             defaults=defaults
         )
 
+        if course_run.draft_version:
+            draft_seat, _ = course_run.draft_version.seats.update_or_create(
+                draft=True,
+                type=seat_type,
+                credit_provider=credit_provider,
+                currency=currency,
+                defaults=defaults
+            )
+            seat.draft_version = draft_seat
+            seat.save()
+
         if created:
             logger.info('Created seat for course with key [%s] and sku [%s].', course_run.key, sku)
 
@@ -740,7 +761,13 @@ def update_entitlement(self, body):
             title=title, sku=sku, partner=self.partner
         )
         logger.info(msg)
-        course.entitlements.update_or_create(mode=mode, defaults=defaults)
+        entitlement, _ = course.entitlements.update_or_create(mode=mode, defaults=defaults)
+        if course.draft_version:
+            draft_entitlement, _ = course.draft_version.entitlements.update_or_create(
+                draft=True, mode=mode, defaults=defaults
+            )
+            entitlement.draft_version = draft_entitlement
+            entitlement.save()
         return sku
 
     def update_enrollment_code(self, body):
@@ -788,7 +815,14 @@ def update_enrollment_code(self, body):
         )
         logger.info(msg)
 
-        course_run.seats.update_or_create(type=seat_type, defaults=defaults)
+        seat, _ = course_run.seats.update_or_create(type=seat_type, defaults=defaults)
+        if course_run.draft_version:
+            draft_seat, _ = course_run.draft_version.seats.update_or_create(
+                draft=True, type=seat_type, defaults=defaults
+            )
+            seat.draft_version = draft_seat
+            seat.save()
+
         return sku
 
     def get_certificate_type(self, product):

course_discovery/apps/course_metadata/data_loaders/tests/test_api.py

@@ -31,7 +31,7 @@
     SourceFactory
 )
 from course_discovery.apps.course_metadata.toggles import BYPASS_LMS_DATA_LOADER__END_DATE_UPDATED_CHECK
-from course_discovery.apps.course_metadata.utils import subtract_deadline_delta
+from course_discovery.apps.course_metadata.utils import ensure_draft_world, subtract_deadline_delta
 
 LOGGER_PATH = 'course_discovery.apps.course_metadata.data_loaders.api.logger'
 
@@ -640,6 +640,7 @@ def mock_products_api(self, alt_course=None, alt_currency=None, alt_mode=None, h
                           valid_stockrecord=True, product_class=None):
         """ Return a new Course Entitlement and Enrollment Code to be added by ingest """
         course = CourseFactory(type=CourseType.objects.get(slug=CourseType.VERIFIED_AUDIT), partner=self.partner)
+        ensure_draft_world(course)
 
         # If product_class is given, make sure it's either entitlement or enrollment_code
         if product_class:
@@ -787,6 +788,7 @@ def assert_seats_loaded(self, body, mock_products):
         products = [p for p in body['products'] if p['structure'] == 'child']
         # Verify that the old seat is removed
         assert course_run.seats.count() == len(products)
+        assert course_run.draft_version.seats.count() == len(products)
 
         # Validate each seat
         for product in products:
@@ -815,16 +817,25 @@ def assert_seats_loaded(self, body, mock_products):
             bulk_sku = self.get_product_bulk_sku(certificate_type, course_run, mock_products)
             seat = course_run.seats.get(type__slug=certificate_type, credit_provider=credit_provider,
                                         currency=price_currency)
+            seat_draft = course_run.draft_version.seats.get(type__slug=certificate_type,
+                                                            credit_provider=credit_provider,
+                                                            currency=price_currency)
 
             assert seat.course_run == course_run
             assert seat.type.slug == certificate_type
+            assert seat_draft.type.slug == certificate_type
             assert seat.price == price
+            assert seat_draft.price == price
             assert seat.currency.code == price_currency
+            assert seat_draft.currency.code == price_currency
             assert seat.credit_provider == credit_provider
+            assert seat_draft.credit_provider == credit_provider
             assert seat.credit_hours == credit_hours
             assert seat.upgrade_deadline == upgrade_deadline
             assert seat.sku == sku
+            assert seat_draft.sku == sku
             assert seat.bulk_sku == bulk_sku
+            assert seat_draft.bulk_sku == bulk_sku
 
     def assert_entitlements_loaded(self, body):
         """ Assert a Course Entitlement was loaded into the database for each entry in the specified data body. """
@@ -842,11 +853,14 @@ def assert_entitlements_loaded(self, body):
             mode = SeatType.objects.get(slug=mode_name)
 
             entitlement = course.entitlements.get(mode=mode)
-
+            entitlement_draft = course.draft_version.entitlements.get(mode=mode)
             assert entitlement.course == course
             assert entitlement.price == price
+            assert entitlement_draft.price == price
             assert entitlement.currency.code == price_currency
+            assert entitlement_draft.currency.code == price_currency
             assert entitlement.sku == sku
+            assert entitlement_draft.sku == sku
 
     def assert_enrollment_codes_loaded(self, body):
         """ Assert a Course Enrollment Code was loaded into the database for each entry in the specified data body. """
@@ -859,15 +873,19 @@ def assert_enrollment_codes_loaded(self, body):
 
             mode_name = attributes['seat_type']
             seat = course_run.seats.get(type__slug=mode_name)
-
+            draft_seat = course_run.draft_version.seats.get(type__slug=mode_name)
             assert seat.course_run == course_run
             assert seat.bulk_sku == bulk_sku
+            assert draft_seat.bulk_sku == bulk_sku
+            assert draft_seat == seat.draft_version
 
     @responses.activate
     def test_ingest(self):
         """ Verify the method ingests data from the E-Commerce API. """
         TieredCache.dangerous_clear_all_tiers()
         courses_api_data = self.mock_courses_api()
+        for course_run in CourseRun.objects.all():
+            ensure_draft_world(course_run)
         loaded_course_run_data = courses_api_data[:-1]
         loaded_seat_data = courses_api_data[:-2]
         assert CourseRun.objects.count() == len(loaded_course_run_data)
@@ -899,15 +917,24 @@ def test_ingest_deletes(self, mock_logger):
         self.mock_courses_api()
         products_api_data = self.mock_products_api()
         entitlement = CourseEntitlementFactory(partner=self.partner)
-
+        entitlement_draft_should_not_delete = CourseEntitlementFactory(draft=True, sku='', partner=self.partner)
+        entitlement_draft_should_delete = CourseEntitlementFactory(draft=True, partner=self.partner)
         self.loader.ingest()
         # Ensure that only entitlements retrieved from the Ecommerce API remain in Discovery,
         # and that the sku and partner of the deleted entitlement are logged
         self.assert_entitlements_loaded(products_api_data)
-        msg = 'Deleting entitlement for course {course_title} with sku {sku} for partner {partner}'.format(
+        msg = 'Deleting non-draft entitlement for course {course_title} with sku {sku} for partner {partner}'.format(
             course_title=entitlement.course.title, sku=entitlement.sku, partner=entitlement.partner
         )
+        msg_draft = 'Deleting draft entitlement for course {course_title} with sku {sku} for partner {partner}'.format(
+            course_title=entitlement_draft_should_delete.course.title,
+            sku=entitlement_draft_should_delete.sku,
+            partner=entitlement_draft_should_delete.partner
+        )
         mock_logger.info.assert_any_call(msg)
+        mock_logger.info.assert_any_call(msg_draft)
+        assert not CourseEntitlement.everything.filter(pk=entitlement_draft_should_delete.pk).exists()
+        assert CourseEntitlement.everything.filter(pk=entitlement_draft_should_not_delete.pk).exists()
 
     @responses.activate
     @mock.patch(LOGGER_PATH)