From aaac71cf9411a72870c1abc7c2a3bb8a68516a79 Mon Sep 17 00:00:00 2001
From: Muhammad Afaq Shuaib <afaq.shoaib09@gmail.com>
Date: Thu, 25 Jul 2024 20:41:46 +0500
Subject: [PATCH] feat: Skipping price/couse updates in Course Entitlements for
 restricted runs (#4389)

---
 .../data_loaders/csv_loader.py                |  9 +++-
 .../data_loaders/tests/test_csv_loader.py     | 49 +++++++++++++++++++
 2 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/course_discovery/apps/course_metadata/data_loaders/csv_loader.py b/course_discovery/apps/course_metadata/data_loaders/csv_loader.py
index 8466f7a6cb..4ef9ee38e8 100644
--- a/course_discovery/apps/course_metadata/data_loaders/csv_loader.py
+++ b/course_discovery/apps/course_metadata/data_loaders/csv_loader.py
@@ -12,7 +12,7 @@
 
 from course_discovery.apps.core.utils import serialize_datetime
 from course_discovery.apps.course_metadata.choices import (
-    CourseRunStatus, ExternalCourseMarketingType, ExternalProductStatus
+    CourseRunRestrictionType, CourseRunStatus, ExternalCourseMarketingType, ExternalProductStatus
 )
 from course_discovery.apps.course_metadata.data_loaders import AbstractDataLoader
 from course_discovery.apps.course_metadata.data_loaders.constants import (
@@ -525,6 +525,11 @@ def _update_course_api_request_data(self, data, course, is_draft):
         Create and return the request data for making a patch call to update the course.
         """
         collaborator_uuids = self.process_collaborators(data.get('collaborators', ''), course.key)
+        price = (
+            self.get_pricing_representation(data['verified_price'], course.type)
+            if data.get('restriction_type', 'None') != CourseRunRestrictionType.CustomB2BEnterprise.value else {}
+        )
+
         subjects = self.get_subject_slugs(
             data.get('primary_subject'),
             data.get('secondary_subject'),
@@ -539,7 +544,7 @@ def _update_course_api_request_data(self, data, course, is_draft):
             'type': str(course.type.uuid),
             'subjects': subjects,
             'collaborators': collaborator_uuids,
-            'prices': self.get_pricing_representation(data['verified_price'], course.type),
+            'prices': price,
 
             'title': data['title'],
             'syllabus_raw': data.get('syllabus', ''),
diff --git a/course_discovery/apps/course_metadata/data_loaders/tests/test_csv_loader.py b/course_discovery/apps/course_metadata/data_loaders/tests/test_csv_loader.py
index abc066ce47..9ebe0268d0 100644
--- a/course_discovery/apps/course_metadata/data_loaders/tests/test_csv_loader.py
+++ b/course_discovery/apps/course_metadata/data_loaders/tests/test_csv_loader.py
@@ -821,6 +821,55 @@ def test_ingest_flow_for_minimal_course_data(self, jwt_decode_patch):  # pylint:
                     assert course.subjects.first().slug == "computer-science"
                     assert course_run.staff.exists() is False
 
+    @data(True, False)
+    @responses.activate
+    def test_entitlement_price_update_for_custom_presentation(self, reverse_order, jwt_decode_patch):  # pylint: disable=unused-argument
+        """
+        Verify that the loader does not update price for custom-b2b-enterprise in Course's Entitlement.
+        """
+        self._setup_prerequisites(self.partner)
+        self.mock_ecommerce_publication(self.partner)
+        self.mock_studio_calls(self.partner)
+        self.mock_image_response()
+        csv_key_order = list(self.CSV_DATA_KEYS_ORDER)
+        csv_key_order.append('restriction_type')
+
+        csv_data = [
+            {
+                **mock_data.VALID_COURSE_AND_COURSE_RUN_CSV_DICT,
+                "restriction_type": None,
+                "short_description": "ABC",
+            },
+            {
+                **mock_data.VALID_COURSE_AND_COURSE_RUN_CSV_DICT,
+                "restriction_type": "custom-b2b-enterprise",
+                "verified_price": "250",
+                "variant_id": "11111111-1111-1111-1111-111111111111",
+                "short_description": "ABC",
+            },
+        ]
+
+        if reverse_order:
+            csv_data = list(reversed(csv_data))
+
+        with NamedTemporaryFile() as csv:
+            csv = self._write_csv(csv, csv_data, csv_key_order)
+
+            with LogCapture(LOGGER_PATH) as log_capture:
+                with mock.patch.object(
+                        CSVDataLoader,
+                        '_call_course_api',
+                        self.mock_call_course_api
+                ):
+                    loader = CSVDataLoader(self.partner, csv_path=csv.name, product_source=self.source.slug)
+                    loader.ingest()
+
+                    self._assert_default_logs(log_capture)
+                    course = Course.objects.get(title='CSV Course')
+                    assert course.entitlements.count() == 1
+                    assert course.entitlements.first().price == 150
+                    assert course.short_description == '<p>ABC</p>'
+
     @responses.activate
     def test_ingest_product_metadata_flow_for_non_exec_ed(self, jwt_decode_patch):  # pylint: disable=unused-argument
         """