From 1ab24d65fec8a3cef3188409643329eb04670268 Mon Sep 17 00:00:00 2001
From: Ali Akbar <52413434+Ali-D-Akbar@users.noreply.github.com>
Date: Tue, 20 Aug 2024 17:58:59 +0500
Subject: [PATCH] fix: use correct course_key in case of a key_for_rerun in RCM
 (#4411)

---
 .../course_metadata/data_loaders/__init__.py  | 12 +++++----
 .../apps/course_metadata/data_loaders/api.py  |  1 +
 .../data_loaders/tests/test_api.py            | 26 +++++++++++++++++++
 3 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/course_discovery/apps/course_metadata/data_loaders/__init__.py b/course_discovery/apps/course_metadata/data_loaders/__init__.py
index 94ea321744..646aa4bd00 100644
--- a/course_discovery/apps/course_metadata/data_loaders/__init__.py
+++ b/course_discovery/apps/course_metadata/data_loaders/__init__.py
@@ -3,7 +3,7 @@
 from dateutil.parser import parse
 from edx_rest_framework_extensions.auth.jwt.decoder import configured_jwt_decode_handler
 
-from course_discovery.apps.course_metadata.models import Image, Video
+from course_discovery.apps.course_metadata.models import Course, Image, Video
 
 
 class AbstractDataLoader(metaclass=abc.ABCMeta):
@@ -80,16 +80,18 @@ def parse_date(cls, date_string):
     @classmethod
     def get_course_key_from_course_run_key(cls, course_run_key):
         """
-        Given a serialized course run key, return the corresponding
-        serialized course key.
+        Given a serialized course run key, return the corresponding serialized course key.
+        If the course key is a 'key_for_rerun', return the original course's key.
 
         Args:
             course_run_key (CourseKey): Course run key.
 
         Returns:
-            str
+            str: The corresponding course key.
         """
-        return f'{course_run_key.org}+{course_run_key.course}'
+        course_key = f'{course_run_key.org}+{course_run_key.course}'
+        # Check if this course key is a rerun and return the original course key if it is
+        return Course.objects.filter(key_for_reruns=course_key).values_list('key', flat=True).first() or course_key
 
     @classmethod
     def _get_or_create_media(cls, media_type, url):
diff --git a/course_discovery/apps/course_metadata/data_loaders/api.py b/course_discovery/apps/course_metadata/data_loaders/api.py
index a63aca75e4..1c16ad920f 100644
--- a/course_discovery/apps/course_metadata/data_loaders/api.py
+++ b/course_discovery/apps/course_metadata/data_loaders/api.py
@@ -202,6 +202,7 @@ def create_course_run(self, course, body):
     def get_or_create_course(self, body):
         course_run_key = CourseKey.from_string(body['id'])
         course_key = self.get_course_key_from_course_run_key(course_run_key)
+
         defaults = self.format_course_data(body)
         # We need to add the key to the defaults because django ignores kwargs with __
         # separators when constructing the create request
diff --git a/course_discovery/apps/course_metadata/data_loaders/tests/test_api.py b/course_discovery/apps/course_metadata/data_loaders/tests/test_api.py
index ff616fa9be..57b1a0575e 100644
--- a/course_discovery/apps/course_metadata/data_loaders/tests/test_api.py
+++ b/course_discovery/apps/course_metadata/data_loaders/tests/test_api.py
@@ -534,6 +534,32 @@ def test_ingest_studio_made_run_with_existing_draft_course(self):
         assert draft_run != official_run
         assert official_run.draft_version == draft_run
 
+    @responses.activate
+    def test_ingest_with_key_for_rerun(self):
+        """
+        Verify that we correctly stitch up a course run made in Studio but having key_for_rerun.
+        The test makes sure that the course_run isn't incorrectly assigned to course having key = key_for_rerun.
+        """
+        datum = mock_data.COURSES_API_BODY_ORIGINAL
+        self.mock_api([datum])
+        course_type = CourseType.objects.get(slug=CourseType.AUDIT)
+        course_key = f"{datum['org']}+{datum['number']}"
+        course_key_with_key_for_rerun = 'test+key_for_rerun'
+        Course.objects.create(partner=self.partner, key=course_key, title='Title', type=course_type)
+        Course.objects.create(partner=self.partner, key=course_key_with_key_for_rerun, title='Key for Reruns',
+                              type=course_type, key_for_reruns=course_key)
+
+        self.loader.ingest()
+
+        test_course = Course.everything.get(partner=self.partner, key=course_key)
+        test_course_with_key_for_rerun = Course.everything.get(partner=self.partner, key=course_key_with_key_for_rerun)
+        test_course_run = CourseRun.everything.filter(course=test_course).first()
+        test_course_run_with_key_for_rerun = CourseRun.everything.filter(course=test_course_with_key_for_rerun).first()
+
+        assert test_course_with_key_for_rerun.key_for_reruns == test_course.key
+        assert test_course_run is None
+        assert test_course_run_with_key_for_rerun.course is not test_course
+
     def test_assigns_types(self):
         """
         Verify we set the special empty course and course run types when creating courses and runs.